QueryCaching: Use CachingServiceClient for query caching (#112128)

* Integrate mt querier with query caching

* typo

* let the caller set cache status response header

* fix TestQueryAPI

* make gen-go

* handle CachingServiceClient being nil and make gen-go

* include namespace in cache key

* set signed in user namespace in query_test.go

* fix test

* remove commented out code

* undo services/query/query.go changes

* make gen-go

* remove namespace requirement

* fix tests

* fix test

* remove namespace from SignedInUser in tests

* make gen-go
This commit is contained in:
Bruno
2025-10-28 11:41:46 -03:00
committed by GitHub
parent 3131a69f04
commit 437dcc875c
12 changed files with 345 additions and 305 deletions
@@ -1,43 +0,0 @@
package clientmiddleware
import (
"github.com/grafana/grafana/pkg/infra/metrics"
contextmodel "github.com/grafana/grafana/pkg/services/contexthandler/model"
"github.com/prometheus/client_golang/prometheus"
)
const (
QueryPubdash = "pubdash"
QueryDashboard = "dashboard"
)
var QueryCachingRequestHistogram = prometheus.NewHistogramVec(prometheus.HistogramOpts{
Namespace: metrics.ExporterName,
Subsystem: "caching",
Name: "query_caching_request_duration_seconds",
Help: "histogram of grafana query endpoint requests in seconds",
Buckets: []float64{.005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5, 10, 25, 50, 100},
}, []string{"datasource_type", "cache", "query_type"})
var ShouldCacheQueryHistogram = prometheus.NewHistogramVec(prometheus.HistogramOpts{
Namespace: metrics.ExporterName,
Subsystem: "caching",
Name: "should_cache_query_request_duration_seconds",
Help: "histogram of grafana query endpoint requests in seconds",
Buckets: []float64{.005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5, 10, 25, 50, 100},
}, []string{"datasource_type", "cache", "shouldCache", "query_type"})
var ResourceCachingRequestHistogram = prometheus.NewHistogramVec(prometheus.HistogramOpts{
Namespace: metrics.ExporterName,
Subsystem: "caching",
Name: "resource_caching_request_duration_seconds",
Help: "histogram of grafana resource endpoint requests in seconds",
Buckets: []float64{.005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5, 10, 25, 50, 100},
}, []string{"plugin_id", "cache"})
func getQueryType(req *contextmodel.ReqContext) string {
if req.IsPublicDashboardView() {
return QueryPubdash
}
return QueryDashboard
}
@@ -2,222 +2,56 @@ package clientmiddleware
import (
"context"
"fmt"
"strconv"
"time"
"github.com/grafana/grafana-aws-sdk/pkg/awsds"
"github.com/grafana/grafana-plugin-sdk-go/backend"
"github.com/prometheus/client_golang/prometheus"
"golang.org/x/sync/singleflight"
"github.com/grafana/grafana/pkg/infra/log"
"github.com/grafana/grafana/pkg/services/caching"
"github.com/grafana/grafana/pkg/services/contexthandler"
"github.com/grafana/grafana/pkg/services/featuremgmt"
)
// needed to mock the function for testing
var shouldCacheQuery = awsds.ShouldCacheQuery
// NewCachingMiddleware creates a new backend.HandlerMiddleware that will
// attempt to read and write query results to the cache
func NewCachingMiddleware(cachingService caching.CachingService) backend.HandlerMiddleware {
return NewCachingMiddlewareWithFeatureManager(cachingService, nil)
}
// NewCachingMiddlewareWithFeatureManager creates a new backend.HandlerMiddleware that will
// attempt to read and write query results to the cache with a feature manager
func NewCachingMiddlewareWithFeatureManager(cachingService caching.CachingService, features featuremgmt.FeatureToggles) backend.HandlerMiddleware {
log := log.New("caching_middleware")
if err := prometheus.Register(QueryCachingRequestHistogram); err != nil {
log.Error("Error registering prometheus collector 'QueryRequestHistogram'", "error", err)
}
if err := prometheus.Register(ResourceCachingRequestHistogram); err != nil {
log.Error("Error registering prometheus collector 'ResourceRequestHistogram'", "error", err)
}
func NewCachingMiddleware(cachingServiceClient *caching.CachingServiceClient) backend.HandlerMiddleware {
cachingMiddlewareHandler := func(next backend.Handler) backend.Handler {
cachingMiddleware := &CachingMiddleware{
BaseHandler: backend.NewBaseHandler(next),
caching: cachingService,
log: log,
features: features,
return &CachingMiddleware{
BaseHandler: backend.NewBaseHandler(next),
cachingServiceClient: cachingServiceClient,
}
if features != nil && features.IsEnabled(context.Background(), featuremgmt.FlagQueryCacheRequestDeduplication) {
return newRequestDeduplicationMiddleware(log, cachingMiddleware)
}
return cachingMiddleware
}
return backend.HandlerMiddlewareFunc(cachingMiddlewareHandler)
}
// An adapter to use CachingServiceClient as a middleware. If possible prefer to use `CachingServiceClient` directly.
type CachingMiddleware struct {
backend.BaseHandler
caching caching.CachingService
log log.Logger
features featuremgmt.FeatureToggles
cachingServiceClient *caching.CachingServiceClient
}
// QueryData receives a data request and attempts to access results already stored in the cache for that request.
// If data is found, it will return it immediately. Otherwise, it will perform the queries as usual, then write the response to the cache.
// If the cache service is implemented, we capture the request duration as a metric. The service is expected to write any response headers.
func (m *CachingMiddleware) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
if req == nil {
return m.BaseHandler.QueryData(ctx, req)
}
reqCtx := contexthandler.FromContext(ctx)
if reqCtx == nil {
return m.BaseHandler.QueryData(ctx, req)
}
// time how long this request takes
start := time.Now()
// First look in the query cache if enabled
hit, cr := m.caching.HandleQueryRequest(ctx, req)
// record request duration if caching was used
ch := reqCtx.Resp.Header().Get(caching.XCacheHeader)
if ch != "" {
defer func() {
QueryCachingRequestHistogram.With(prometheus.Labels{
"datasource_type": req.PluginContext.DataSourceInstanceSettings.Type,
"cache": ch,
"query_type": getQueryType(reqCtx),
}).Observe(time.Since(start).Seconds())
}()
}
// Cache hit; return the response
if hit {
return cr.Response, nil
}
// Cache miss; do the actual queries
resp, err := m.BaseHandler.QueryData(ctx, req)
// Update the query cache with the result for this metrics request
if err == nil && cr.UpdateCacheFn != nil {
// If AWS async caching is not enabled, use the old code path
if m.features == nil || !m.features.IsEnabled(ctx, featuremgmt.FlagAwsAsyncQueryCaching) {
cr.UpdateCacheFn(ctx, resp)
} else {
// time how long shouldCacheQuery takes
startShouldCacheQuery := time.Now()
shouldCache := shouldCacheQuery(resp)
ShouldCacheQueryHistogram.With(prometheus.Labels{
"datasource_type": req.PluginContext.DataSourceInstanceSettings.Type,
"cache": ch,
"shouldCache": strconv.FormatBool(shouldCache),
"query_type": getQueryType(reqCtx),
}).Observe(time.Since(startShouldCacheQuery).Seconds())
// If AWS async caching is enabled and resp is for a running async query, don't cache it
if shouldCache {
cr.UpdateCacheFn(ctx, resp)
}
}
}
return resp, err
return m.cachingServiceClient.WithQueryDataCaching(ctx, req, func() (*backend.QueryDataResponse, error) {
return m.BaseHandler.QueryData(ctx, req)
})
}
// CallResource receives a resource request and attempts to access results already stored in the cache for that request.
// If data is found, it will return it immediately. Otherwise, it will perform the request as usual. The caller of CallResource is expected to explicitly update the cache with any responses.
// If the cache service is implemented, we capture the request duration as a metric. The service is expected to write any response headers.
func (m *CachingMiddleware) CallResource(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
if req == nil {
return m.BaseHandler.CallResource(ctx, req, sender)
}
reqCtx := contexthandler.FromContext(ctx)
if reqCtx == nil {
return m.BaseHandler.CallResource(ctx, req, sender)
}
// time how long this request takes
start := time.Now()
// First look in the resource cache if enabled
hit, cr := m.caching.HandleResourceRequest(ctx, req)
// record request duration if caching was used
if ch := reqCtx.Resp.Header().Get(caching.XCacheHeader); ch != "" {
defer func() {
ResourceCachingRequestHistogram.With(prometheus.Labels{
"plugin_id": req.PluginContext.PluginID,
"cache": ch,
}).Observe(time.Since(start).Seconds())
}()
}
// Cache hit; send the response and return
if hit {
return sender.Send(cr.Response)
}
// Cache miss; do the actual request
// If there is no update cache func, just pass in the original sender
if cr.UpdateCacheFn == nil {
return m.cachingServiceClient.WithCallResourceCaching(ctx, req, sender, func(sender backend.CallResourceResponseSender) error {
return m.BaseHandler.CallResource(ctx, req, sender)
}
// Otherwise, intercept the responses in a wrapped sender so we can cache them first
cacheSender := backend.CallResourceResponseSenderFunc(func(res *backend.CallResourceResponse) error {
cr.UpdateCacheFn(ctx, res)
return sender.Send(res)
})
return m.BaseHandler.CallResource(ctx, req, cacheSender)
}
// Given N requests happening at the same time and issuing the same query, only one request will execute
// and the other ones will wait for the response received by the request being executed.
type requestDeduplicationMiddleware struct {
backend.BaseHandler
log *log.ConcreteLogger
singleflight *singleflight.Group
}
func newRequestDeduplicationMiddleware(log *log.ConcreteLogger, next backend.Handler) *requestDeduplicationMiddleware {
return &requestDeduplicationMiddleware{log: log, BaseHandler: backend.NewBaseHandler(next), singleflight: &singleflight.Group{}}
}
func (m *requestDeduplicationMiddleware) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
if req.PluginContext.DataSourceInstanceSettings == nil || req.PluginContext.DataSourceInstanceSettings.UID == "" {
return m.BaseHandler.QueryData(ctx, req)
}
key, err := caching.GetKey(req.PluginContext.DataSourceInstanceSettings.UID, req)
if err != nil {
m.log.Error("error building cache key for request deduplication, skipping request deduplication", "error", err)
return m.BaseHandler.QueryData(ctx, req)
}
v, err, _ := m.singleflight.Do(key, func() (interface{}, error) {
return m.BaseHandler.QueryData(ctx, req)
})
if err != nil {
return nil, fmt.Errorf("request deduplication middleware: calling BaseHandler.QueryData: %w", err)
}
return v.(*backend.QueryDataResponse), nil
}
func (m *requestDeduplicationMiddleware) CallResource(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
if req.PluginContext.DataSourceInstanceSettings == nil || req.PluginContext.DataSourceInstanceSettings.UID == "" {
return m.BaseHandler.CallResource(ctx, req, sender)
}
key, err := caching.GetKey(req.PluginContext.DataSourceInstanceSettings.UID, req)
if err != nil {
m.log.Error("error building cache key for request deduplication, skipping request deduplication", "error", err)
return m.BaseHandler.CallResource(ctx, req, sender)
}
_, err, _ = m.singleflight.Do(key, func() (interface{}, error) {
return nil, m.BaseHandler.CallResource(ctx, req, sender)
})
if err != nil {
return fmt.Errorf("request deduplication middleware: calling BaseHandler.CallResource: %w", err)
}
return nil
}
@@ -4,10 +4,7 @@ import (
"context"
"encoding/json"
"net/http"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/grafana/grafana-plugin-sdk-go/backend"
"github.com/grafana/grafana-plugin-sdk-go/backend/handlertest"
@@ -25,9 +22,10 @@ func TestCachingMiddleware(t *testing.T) {
require.NoError(t, err)
cs := caching.NewFakeOSSCachingService()
cachingServiceClient := caching.ProvideCachingServiceClient(cs, nil)
cdt := handlertest.NewHandlerMiddlewareTest(t,
WithReqContext(req, &user.SignedInUser{}),
handlertest.WithMiddlewares(NewCachingMiddleware(cs)),
handlertest.WithMiddlewares(NewCachingMiddleware(cachingServiceClient)),
)
jsonDataMap := map[string]any{}
@@ -78,9 +76,9 @@ func TestCachingMiddleware(t *testing.T) {
})
t.Run("If cache returns a miss, queries are issued and the update cache function is called", func(t *testing.T) {
origShouldCacheQuery := shouldCacheQuery
origShouldCacheQuery := caching.ShouldCacheQuery
var shouldCacheQueryCalled bool
shouldCacheQuery = func(resp *backend.QueryDataResponse) bool {
caching.ShouldCacheQuery = func(resp *backend.QueryDataResponse) bool {
shouldCacheQueryCalled = true
return true
}
@@ -88,7 +86,7 @@ func TestCachingMiddleware(t *testing.T) {
t.Cleanup(func() {
updateCacheCalled = false
shouldCacheQueryCalled = false
shouldCacheQuery = origShouldCacheQuery
caching.ShouldCacheQuery = origShouldCacheQuery
cs.Reset()
})
@@ -108,15 +106,16 @@ func TestCachingMiddleware(t *testing.T) {
})
t.Run("with async queries", func(t *testing.T) {
cachingServiceClient := caching.ProvideCachingServiceClient(cs, featuremgmt.WithFeatures(featuremgmt.FlagAwsAsyncQueryCaching))
asyncCdt := handlertest.NewHandlerMiddlewareTest(t,
WithReqContext(req, &user.SignedInUser{}),
handlertest.WithMiddlewares(
NewCachingMiddlewareWithFeatureManager(cs, featuremgmt.WithFeatures(featuremgmt.FlagAwsAsyncQueryCaching))),
NewCachingMiddleware(cachingServiceClient)),
)
t.Run("If shoudCacheQuery returns true update cache function is called", func(t *testing.T) {
origShouldCacheQuery := shouldCacheQuery
origShouldCacheQuery := caching.ShouldCacheQuery
var shouldCacheQueryCalled bool
shouldCacheQuery = func(resp *backend.QueryDataResponse) bool {
caching.ShouldCacheQuery = func(resp *backend.QueryDataResponse) bool {
shouldCacheQueryCalled = true
return true
}
@@ -124,7 +123,7 @@ func TestCachingMiddleware(t *testing.T) {
t.Cleanup(func() {
updateCacheCalled = false
shouldCacheQueryCalled = false
shouldCacheQuery = origShouldCacheQuery
caching.ShouldCacheQuery = origShouldCacheQuery
cs.Reset()
})
@@ -144,9 +143,9 @@ func TestCachingMiddleware(t *testing.T) {
})
t.Run("If shoudCacheQuery returns false update cache function is not called", func(t *testing.T) {
origShouldCacheQuery := shouldCacheQuery
origShouldCacheQuery := caching.ShouldCacheQuery
var shouldCacheQueryCalled bool
shouldCacheQuery = func(resp *backend.QueryDataResponse) bool {
caching.ShouldCacheQuery = func(resp *backend.QueryDataResponse) bool {
shouldCacheQueryCalled = true
return false
}
@@ -154,7 +153,7 @@ func TestCachingMiddleware(t *testing.T) {
t.Cleanup(func() {
updateCacheCalled = false
shouldCacheQueryCalled = false
shouldCacheQuery = origShouldCacheQuery
caching.ShouldCacheQuery = origShouldCacheQuery
cs.Reset()
})
@@ -199,9 +198,10 @@ func TestCachingMiddleware(t *testing.T) {
}
cs := caching.NewFakeOSSCachingService()
cachingServiceClient := caching.ProvideCachingServiceClient(cs, nil)
cdt := handlertest.NewHandlerMiddlewareTest(t,
WithReqContext(req, &user.SignedInUser{}),
handlertest.WithMiddlewares(NewCachingMiddleware(cs)),
handlertest.WithMiddlewares(NewCachingMiddleware(cachingServiceClient)),
handlertest.WithResourceResponses([]*backend.CallResourceResponse{simulatedPluginResponse}),
)
@@ -275,9 +275,10 @@ func TestCachingMiddleware(t *testing.T) {
require.NoError(t, err)
cs := caching.NewFakeOSSCachingService()
cachingServiceClient := caching.ProvideCachingServiceClient(cs, nil)
cdt := handlertest.NewHandlerMiddlewareTest(t,
// Skip the request context in this case
handlertest.WithMiddlewares(NewCachingMiddleware(cs)),
handlertest.WithMiddlewares(NewCachingMiddleware(cachingServiceClient)),
)
reqCtx := contexthandler.FromContext(req.Context())
require.Nil(t, reqCtx)
@@ -325,86 +326,3 @@ func TestCachingMiddleware(t *testing.T) {
})
})
}
func TestRequestDeduplicationMiddleware(t *testing.T) {
t.Parallel()
t.Run("deduplicates requests issuing the same query", func(t *testing.T) {
t.Parallel()
handler := newMockMiddlewareHandler()
middleware := newRequestDeduplicationMiddleware(nil, handler)
req := backend.QueryDataRequest{
PluginContext: backend.PluginContext{
DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{
UID: "uid",
},
},
}
wg := &sync.WaitGroup{}
wg.Add(2)
for range 2 {
go func() {
defer wg.Done()
resp, err := middleware.QueryData(t.Context(), &req)
require.NoError(t, err)
require.Equal(t, &backend.QueryDataResponse{}, resp)
}()
}
wg.Wait()
require.EqualValues(t, 1, handler.QueryDataCalls)
})
t.Run("requests where DataSourceInstanceSettings is nil bypass request deduplication", func(t *testing.T) {
t.Parallel()
handler := newMockMiddlewareHandler()
middleware := newRequestDeduplicationMiddleware(nil, handler)
{
req := backend.QueryDataRequest{
PluginContext: backend.PluginContext{
DataSourceInstanceSettings: nil,
},
}
resp, err := middleware.QueryData(t.Context(), &req)
require.NoError(t, err)
require.Empty(t, resp)
}
{
req := backend.CallResourceRequest{
PluginContext: backend.PluginContext{
DataSourceInstanceSettings: nil,
},
}
require.NoError(t, middleware.CallResource(t.Context(), &req, nil))
}
})
}
type mockMiddlewareHandler struct {
backend.BaseHandler
QueryDataCalls int32
}
func newMockMiddlewareHandler() *mockMiddlewareHandler {
return &mockMiddlewareHandler{}
}
func (m *mockMiddlewareHandler) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
atomic.AddInt32(&m.QueryDataCalls, 1)
time.Sleep(10 * time.Millisecond)
return &backend.QueryDataResponse{}, nil
}
func (m *mockMiddlewareHandler) CallResource(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
return nil
}