diff --git a/pkg/storage/unified/resource/search_client.go b/pkg/storage/unified/resource/search_client.go index 551472cea50..79b6abf395b 100644 --- a/pkg/storage/unified/resource/search_client.go +++ b/pkg/storage/unified/resource/search_client.go @@ -12,6 +12,7 @@ import ( "github.com/grafana/grafana/pkg/infra/log" "github.com/grafana/grafana/pkg/services/featuremgmt" + "github.com/grafana/grafana/pkg/storage/legacysql/dualwrite" "github.com/grafana/grafana/pkg/storage/unified/resourcepb" ) @@ -38,6 +39,7 @@ var ( type DualWriter interface { IsEnabled(schema.GroupResource) bool ReadFromUnified(context.Context, schema.GroupResource) (bool, error) + Status(ctx context.Context, gr schema.GroupResource) (dualwrite.StorageStatus, error) } func NewSearchClient(dual DualWriter, gr schema.GroupResource, unifiedClient resourcepb.ResourceIndexClient, @@ -99,6 +101,28 @@ func calculateMatchPercentage(legacyUIDs, unifiedUIDs map[string]struct{}) float return float64(matches) / float64(len(legacyUIDs)) * 100.0 } +// If dual reader feature flag is enabled, and legacy is the main storage, +// and we are writing to unified (which means we are effectively dual writing), +// then make a background call to unified +func shouldMakeBackgroundCall(ctx context.Context, features featuremgmt.FeatureToggles, dual DualWriter, gr schema.GroupResource) (bool, error) { + unifiedIsMainStorage, err := dual.ReadFromUnified(ctx, gr) + if err != nil { + return false, err + } + + status, err := dual.Status(ctx, gr) + if err != nil { + return false, err + } + + res := features != nil && + features.IsEnabledGlobally(featuremgmt.FlagUnifiedStorageSearchDualReaderEnabled) && + !unifiedIsMainStorage && + status.WriteUnified + + return res, nil +} + func (s *searchWrapper) GetStats(ctx context.Context, in *resourcepb.ResourceStatsRequest, opts ...grpc.CallOption) (*resourcepb.ResourceStatsResponse, error) { client := s.legacyClient @@ -110,9 +134,12 @@ func (s *searchWrapper) GetStats(ctx context.Context, in *resourcepb.ResourceSta client = s.unifiedClient } - // If dual reader feature flag is enabled, and legacy is the main storage, - // make a background call to unified - if s.features != nil && s.features.IsEnabledGlobally(featuremgmt.FlagUnifiedStorageSearchDualReaderEnabled) && !unified { + makeBackgroundCall, err := shouldMakeBackgroundCall(ctx, s.features, s.dual, s.groupResource) + if err != nil { + return nil, err + } + + if makeBackgroundCall { // Create background context with timeout but ignore parent cancelation ctxBg := context.WithoutCancel(ctx) @@ -143,9 +170,12 @@ func (s *searchWrapper) Search(ctx context.Context, in *resourcepb.ResourceSearc client = s.unifiedClient } - // If dual reader feature flag is enabled, and legacy is the main storage, - // make a background call to unified and compare results - if s.features != nil && s.features.IsEnabledGlobally(featuremgmt.FlagUnifiedStorageSearchDualReaderEnabled) && !unified { + makeBackgroundCall, err := shouldMakeBackgroundCall(ctx, s.features, s.dual, s.groupResource) + if err != nil { + return nil, err + } + + if makeBackgroundCall { // Get the legacy result first legacyResponse, legacyErr := s.legacyClient.Search(ctx, in, opts...) if legacyErr != nil { diff --git a/pkg/storage/unified/resource/search_client_test.go b/pkg/storage/unified/resource/search_client_test.go index 927f9e9e60e..4b0a73895ec 100644 --- a/pkg/storage/unified/resource/search_client_test.go +++ b/pkg/storage/unified/resource/search_client_test.go @@ -13,6 +13,7 @@ import ( "github.com/grafana/grafana/pkg/infra/log" "github.com/grafana/grafana/pkg/services/featuremgmt" + "github.com/grafana/grafana/pkg/storage/legacysql/dualwrite" "github.com/grafana/grafana/pkg/storage/unified/resourcepb" "github.com/grafana/grafana/pkg/util/testutil" ) @@ -32,6 +33,11 @@ func (m *MockDualWriter) ReadFromUnified(ctx context.Context, gr schema.GroupRes return args.Bool(0), args.Error(1) } +func (m *MockDualWriter) Status(ctx context.Context, gr schema.GroupResource) (dualwrite.StorageStatus, error) { + args := m.Called(ctx, gr) + return args.Get(0).(dualwrite.StorageStatus), args.Error(1) +} + // Mock ResourceIndexClient with enhanced timeout testing capabilities type MockResourceIndexClient struct { mock.Mock @@ -145,15 +151,18 @@ func TestSearchClient_NewSearchClient(t *testing.T) { } func TestSearchWrapper_Search(t *testing.T) { - gr, unifiedClient, legacyClient, features := setupTestSearchClient(t) + // gr, unifiedClient, legacyClient, features := setupTestSearchClient(t) req := &resourcepb.ResourceSearchRequest{Query: "test"} expectedResponse := &resourcepb.ResourceSearchResponse{TotalHits: 0} t.Run("uses unified client when reading from unified", func(t *testing.T) { + gr, unifiedClient, legacyClient, features := setupTestSearchClient(t) + ctx := testutil.NewDefaultTestContext(t) dual := &MockDualWriter{} dual.On("ReadFromUnified", mock.Anything, gr).Return(true, nil) + dual.On("Status", mock.Anything, gr).Return(dualwrite.StorageStatus{ReadUnified: true, WriteUnified: true}, nil) unifiedClient.On("Search", mock.Anything, req, mock.Anything).Return(expectedResponse, nil) wrapper := setupTestSearchWrapper(t, dual, unifiedClient, legacyClient, features, gr) @@ -169,10 +178,13 @@ func TestSearchWrapper_Search(t *testing.T) { }) t.Run("uses legacy client when not reading from unified", func(t *testing.T) { + gr, unifiedClient, legacyClient, features := setupTestSearchClient(t) + ctx := testutil.NewDefaultTestContext(t) dual := &MockDualWriter{} dual.On("ReadFromUnified", mock.Anything, gr).Return(false, nil) + dual.On("Status", mock.Anything, gr).Return(dualwrite.StorageStatus{ReadUnified: false, WriteUnified: true}, nil) legacyClient.On("Search", mock.Anything, req, mock.Anything).Return(expectedResponse, nil) wrapper := setupTestSearchWrapper(t, dual, unifiedClient, legacyClient, features, gr) @@ -187,12 +199,44 @@ func TestSearchWrapper_Search(t *testing.T) { unifiedClient.AssertNotCalled(t, "Search") }) - t.Run("makes background call to unified when feature flag enabled and using legacy", func(t *testing.T) { + t.Run("do not make a background call to unified when feature flag enabled and using legacy with mode 0", func(t *testing.T) { + gr, unifiedClient, legacyClient, _ := setupTestSearchClient(t) + ctx := testutil.NewDefaultTestContext(t) dual := &MockDualWriter{} featuresWithFlag := featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorageSearchDualReaderEnabled) dual.On("ReadFromUnified", mock.Anything, gr).Return(false, nil) + dual.On("Status", mock.Anything, gr).Return(dualwrite.StorageStatus{ReadUnified: false, WriteUnified: false}, nil) + legacyClient.On("Search", mock.Anything, req, mock.Anything).Return(expectedResponse, nil) + + wrapper := setupTestSearchWrapper(t, dual, unifiedClient, legacyClient, featuresWithFlag, gr) + + resp, err := wrapper.Search(ctx, req) + + require.NoError(t, err) + assert.Equal(t, expectedResponse, resp) + + dual.AssertExpectations(t) + legacyClient.AssertExpectations(t) + unifiedClient.AssertExpectations(t) + + // Expect call to legacy client + legacyClient.AssertCalled(t, "Search", mock.Anything, req, mock.Anything) + + // Do not expect background call to unified client + unifiedClient.AssertNotCalled(t, "Search", mock.Anything, req, mock.Anything) + }) + + t.Run("makes background call to unified when feature flag enabled and using legacy", func(t *testing.T) { + gr, unifiedClient, legacyClient, _ := setupTestSearchClient(t) + + ctx := testutil.NewDefaultTestContext(t) + dual := &MockDualWriter{} + featuresWithFlag := featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorageSearchDualReaderEnabled) + + dual.On("ReadFromUnified", mock.Anything, gr).Return(false, nil) + dual.On("Status", mock.Anything, gr).Return(dualwrite.StorageStatus{ReadUnified: false, WriteUnified: true}, nil) legacyClient.On("Search", mock.Anything, req, mock.Anything).Return(expectedResponse, nil) // Expect background call to unified client @@ -220,11 +264,14 @@ func TestSearchWrapper_Search(t *testing.T) { }) t.Run("handles background call error gracefully", func(t *testing.T) { + gr, unifiedClient, legacyClient, _ := setupTestSearchClient(t) + ctx := testutil.NewDefaultTestContext(t) dual := &MockDualWriter{} featuresWithFlag := featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorageSearchDualReaderEnabled) dual.On("ReadFromUnified", mock.Anything, gr).Return(false, nil) + dual.On("Status", mock.Anything, gr).Return(dualwrite.StorageStatus{ReadUnified: false, WriteUnified: true}, nil) legacyClient.On("Search", mock.Anything, req, mock.Anything).Return(expectedResponse, nil) // Background call returns error - should be handled gracefully @@ -252,11 +299,14 @@ func TestSearchWrapper_Search(t *testing.T) { }) t.Run("background request times out after 500ms", func(t *testing.T) { + gr, unifiedClient, legacyClient, _ := setupTestSearchClient(t) + ctx := testutil.NewDefaultTestContext(t) dual := &MockDualWriter{} featuresWithFlag := featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorageSearchDualReaderEnabled) dual.On("ReadFromUnified", mock.Anything, gr).Return(false, nil) + dual.On("Status", mock.Anything, gr).Return(dualwrite.StorageStatus{ReadUnified: false, WriteUnified: true}, nil) legacyClient.On("Search", mock.Anything, req, mock.Anything).Return(expectedResponse, nil) // Configure unified client to take longer than the 500ms timeout @@ -289,11 +339,14 @@ func TestSearchWrapper_Search(t *testing.T) { }) t.Run("background request completes successfully when within timeout", func(t *testing.T) { + gr, unifiedClient, legacyClient, _ := setupTestSearchClient(t) + ctx := testutil.NewDefaultTestContext(t) dual := &MockDualWriter{} featuresWithFlag := featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorageSearchDualReaderEnabled) dual.On("ReadFromUnified", mock.Anything, gr).Return(false, nil) + dual.On("Status", mock.Anything, gr).Return(dualwrite.StorageStatus{ReadUnified: false, WriteUnified: true}, nil) legacyClient.On("Search", mock.Anything, req, mock.Anything).Return(expectedResponse, nil) // Configure unified client to respond within the 500ms timeout @@ -326,15 +379,18 @@ func TestSearchWrapper_Search(t *testing.T) { } func TestSearchWrapper_GetStats(t *testing.T) { - gr, unifiedClient, legacyClient, features := setupTestSearchClient(t) + // gr, unifiedClient, legacyClient, features := setupTestSearchClient(t) req := &resourcepb.ResourceStatsRequest{Namespace: "test"} expectedResponse := &resourcepb.ResourceStatsResponse{Stats: []*resourcepb.ResourceStatsResponse_Stats{{Count: 100}}} t.Run("uses unified client when reading from unified", func(t *testing.T) { + gr, unifiedClient, legacyClient, features := setupTestSearchClient(t) + ctx := testutil.NewDefaultTestContext(t) dual := &MockDualWriter{} dual.On("ReadFromUnified", mock.Anything, gr).Return(true, nil) + dual.On("Status", mock.Anything, gr).Return(dualwrite.StorageStatus{ReadUnified: true, WriteUnified: true}, nil) unifiedClient.On("GetStats", mock.Anything, req, mock.Anything).Return(expectedResponse, nil) wrapper := setupTestSearchWrapper(t, dual, unifiedClient, legacyClient, features, gr) @@ -349,12 +405,44 @@ func TestSearchWrapper_GetStats(t *testing.T) { legacyClient.AssertNotCalled(t, "GetStats") }) - t.Run("makes background call to unified when feature flag enabled and using legacy", func(t *testing.T) { + t.Run("Do not make background call to unified when feature flag enabled and using legacy with mode 0", func(t *testing.T) { + gr, unifiedClient, legacyClient, _ := setupTestSearchClient(t) + ctx := testutil.NewDefaultTestContext(t) dual := &MockDualWriter{} featuresWithFlag := featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorageSearchDualReaderEnabled) dual.On("ReadFromUnified", mock.Anything, gr).Return(false, nil) + dual.On("Status", mock.Anything, gr).Return(dualwrite.StorageStatus{ReadUnified: false, WriteUnified: false}, nil) + legacyClient.On("GetStats", mock.Anything, req, mock.Anything).Return(expectedResponse, nil) + + wrapper := setupTestSearchWrapper(t, dual, unifiedClient, legacyClient, featuresWithFlag, gr) + + resp, err := wrapper.GetStats(ctx, req) + + require.NoError(t, err) + assert.Equal(t, expectedResponse, resp) + + dual.AssertExpectations(t) + legacyClient.AssertExpectations(t) + unifiedClient.AssertExpectations(t) + + // Expect call to legacy client + legacyClient.AssertCalled(t, "GetStats", mock.Anything, req, mock.Anything) + + // Do not expect background call to unified client + unifiedClient.AssertNotCalled(t, "GetStats", mock.Anything, req, mock.Anything) + }) + + t.Run("makes background call to unified when feature flag enabled and using legacy", func(t *testing.T) { + gr, unifiedClient, legacyClient, _ := setupTestSearchClient(t) + + ctx := testutil.NewDefaultTestContext(t) + dual := &MockDualWriter{} + featuresWithFlag := featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorageSearchDualReaderEnabled) + + dual.On("ReadFromUnified", mock.Anything, gr).Return(false, nil) + dual.On("Status", mock.Anything, gr).Return(dualwrite.StorageStatus{ReadUnified: false, WriteUnified: true}, nil) legacyClient.On("GetStats", mock.Anything, req, mock.Anything).Return(expectedResponse, nil) // Expect background call to unified client @@ -382,11 +470,14 @@ func TestSearchWrapper_GetStats(t *testing.T) { }) t.Run("background GetStats request times out after 500ms", func(t *testing.T) { + gr, unifiedClient, legacyClient, _ := setupTestSearchClient(t) + ctx := testutil.NewDefaultTestContext(t) dual := &MockDualWriter{} featuresWithFlag := featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorageSearchDualReaderEnabled) dual.On("ReadFromUnified", mock.Anything, gr).Return(false, nil) + dual.On("Status", mock.Anything, gr).Return(dualwrite.StorageStatus{ReadUnified: false, WriteUnified: true}, nil) legacyClient.On("GetStats", mock.Anything, req, mock.Anything).Return(expectedResponse, nil) // Configure unified client to take longer than the 500ms timeout