Unistore: Only Shadow Search Traffic when running on modes > 0 (#110302)

* Unistore: Only Shadow Search Traffic when running on modes > 0

Signed-off-by: Maicon Costa <maiconscosta@gmail.com>

---------

Signed-off-by: Maicon Costa <maiconscosta@gmail.com>
This commit is contained in:
maicon
2025-09-01 16:14:53 -03:00
committed by GitHub
parent 2f769b0a62
commit fccf58add7
2 changed files with 131 additions and 10 deletions
+36 -6
View File
@@ -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 {
@@ -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