Query caching: add request deduplication middleware (#110892)

* secrets: update test to accept []byte(nil) and []byte{} (#110630)

Co-authored-by: Matheus Macabu <macabu.matheus@gmail.com>

* Query caching: add request deduplication middleware

* log error if unable to build cache key

* remove TODO

* always use req.PluginContext.DataSourceInstanceSettings.UID

* make update-workspace

---------

Co-authored-by: Matheus Macabu <macabu.matheus@gmail.com>

pkg/services/pluginsintegration/clientmiddleware/caching_middleware.go

@@ -2,12 +2,14 @@ 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"
@@ -34,14 +36,20 @@ func NewCachingMiddlewareWithFeatureManager(cachingService caching.CachingServic
 	if err := prometheus.Register(ResourceCachingRequestHistogram); err != nil {
 		log.Error("Error registering prometheus collector 'ResourceRequestHistogram'", "error", err)
 	}
-	return backend.HandlerMiddlewareFunc(func(next backend.Handler) backend.Handler {
-		return &CachingMiddleware{
+	cachingMiddlewareHandler := func(next backend.Handler) backend.Handler {
+		cachingMiddleware := &CachingMiddleware{
 			BaseHandler: backend.NewBaseHandler(next),
 			caching:     cachingService,
 			log:         log,
 			features:    features,
 		}
-	})
+		if features != nil && features.IsEnabled(context.Background(), featuremgmt.FlagQueryCacheRequestDeduplication) {
+			return newRequestDeduplicationMiddleware(log, cachingMiddleware)
+		}
+		return cachingMiddleware
+	}
+
+	return backend.HandlerMiddlewareFunc(cachingMiddlewareHandler)
 }
 
 type CachingMiddleware struct {
@@ -164,3 +172,52 @@ func (m *CachingMiddleware) CallResource(ctx context.Context, req *backend.CallR
 
 	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.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("singleflight: calling next.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.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("singleflight: calling next.CallResource: %w", err)
+	}
+	return nil
+}

pkg/services/pluginsintegration/clientmiddleware/caching_middleware_test.go

@@ -4,7 +4,10 @@ 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"
@@ -322,3 +325,51 @@ func TestCachingMiddleware(t *testing.T) {
 		})
 	})
 }
+
+func TestRequestDeduplicationMiddleware(t *testing.T) {
+	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)
+	})
+}
+
+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
+}