diff --git a/pkg/goldfish/storage.go b/pkg/goldfish/storage.go index 84ffc93c00c7e..5caacc67092d8 100644 --- a/pkg/goldfish/storage.go +++ b/pkg/goldfish/storage.go @@ -7,7 +7,7 @@ import ( // Storage defines the interface for storing and retrieving query samples and comparison results type Storage interface { // Write operations (used by querytee) - StoreQuerySample(ctx context.Context, sample *QuerySample) error + StoreQuerySample(ctx context.Context, sample *QuerySample, comparison *ComparisonResult) error StoreComparisonResult(ctx context.Context, result *ComparisonResult) error // Read operations (used by UI) diff --git a/pkg/goldfish/storage_mysql.go b/pkg/goldfish/storage_mysql.go index b25d20270955d..32e8ac9685ee1 100644 --- a/pkg/goldfish/storage_mysql.go +++ b/pkg/goldfish/storage_mysql.go @@ -73,25 +73,8 @@ func NewMySQLStorage(config StorageConfig, logger log.Logger) (*MySQLStorage, er }, nil } -// computeComparisonStatus determines the comparison status based on response codes and hashes -func computeComparisonStatus(sample *QuerySample) ComparisonStatus { - // If either cell has a non-2xx status code, it's an error - if sample.CellAStatusCode < 200 || sample.CellAStatusCode >= 300 || - sample.CellBStatusCode < 200 || sample.CellBStatusCode >= 300 { - return ComparisonStatusError - } - - // If response hashes match, queries produced identical results - if sample.CellAResponseHash == sample.CellBResponseHash { - return ComparisonStatusMatch - } - - // Otherwise, responses differ - return ComparisonStatusMismatch -} - // StoreQuerySample stores a sampled query with performance statistics -func (s *MySQLStorage) StoreQuerySample(ctx context.Context, sample *QuerySample) error { +func (s *MySQLStorage) StoreQuerySample(ctx context.Context, sample *QuerySample, comparison *ComparisonResult) error { query := ` INSERT INTO sampled_queries ( correlation_id, tenant_id, user, is_logs_drilldown, query, query_type, @@ -147,9 +130,6 @@ func (s *MySQLStorage) StoreQuerySample(ctx context.Context, sample *QuerySample } } - // Compute the comparison status before storing - comparisonStatus := computeComparisonStatus(sample) - _, err := s.db.ExecContext(ctx, query, sample.CorrelationID, sample.TenantID, @@ -197,7 +177,7 @@ func (s *MySQLStorage) StoreQuerySample(ctx context.Context, sample *QuerySample sample.CellAUsedNewEngine, sample.CellBUsedNewEngine, sample.SampledAt, - comparisonStatus, + comparison.ComparisonStatus, ) return err diff --git a/pkg/goldfish/storage_noop.go b/pkg/goldfish/storage_noop.go index 399ad277ee3b7..a0fb2d71cd6cd 100644 --- a/pkg/goldfish/storage_noop.go +++ b/pkg/goldfish/storage_noop.go @@ -14,7 +14,7 @@ func NewNoopStorage() *NoopStorage { } // StoreQuerySample is a no-op -func (n *NoopStorage) StoreQuerySample(_ context.Context, _ *QuerySample) error { +func (n *NoopStorage) StoreQuerySample(_ context.Context, _ *QuerySample, _ *ComparisonResult) error { return nil } diff --git a/pkg/ui/goldfish_test.go b/pkg/ui/goldfish_test.go index 412632d243c66..969b8fe817fc3 100644 --- a/pkg/ui/goldfish_test.go +++ b/pkg/ui/goldfish_test.go @@ -29,7 +29,7 @@ type mockStorage struct { getQueryFunc func(ctx context.Context, correlationID string) (*goldfish.QuerySample, error) } -func (m *mockStorage) StoreQuerySample(_ context.Context, _ *goldfish.QuerySample) error { +func (m *mockStorage) StoreQuerySample(_ context.Context, _ *goldfish.QuerySample, _ *goldfish.ComparisonResult) error { return nil } diff --git a/tools/querytee/goldfish/manager.go b/tools/querytee/goldfish/manager.go index b128ee200f0b3..36f87811b10ac 100644 --- a/tools/querytee/goldfish/manager.go +++ b/tools/querytee/goldfish/manager.go @@ -129,7 +129,7 @@ func (m *Manager) SendToGoldfish(httpReq *http.Request, cellAResp, cellBResp *Ba cellAData, err := CaptureResponse(&http.Response{ StatusCode: cellAResp.Status, Body: io.NopCloser(bytes.NewReader(cellAResp.Body)), - }, cellAResp.Duration, cellAResp.TraceID, cellAResp.SpanID) + }, cellAResp.Duration, cellAResp.TraceID, cellAResp.SpanID, m.logger) if err != nil { level.Error(m.logger).Log("msg", "failed to capture cell A response", "err", err) return @@ -139,7 +139,7 @@ func (m *Manager) SendToGoldfish(httpReq *http.Request, cellAResp, cellBResp *Ba cellBData, err := CaptureResponse(&http.Response{ StatusCode: cellBResp.Status, Body: io.NopCloser(bytes.NewReader(cellBResp.Body)), - }, cellBResp.Duration, cellBResp.TraceID, cellBResp.SpanID) + }, cellBResp.Duration, cellBResp.TraceID, cellBResp.SpanID, m.logger) if err != nil { level.Error(m.logger).Log("msg", "failed to capture cell B response", "err", err) return @@ -232,7 +232,7 @@ func (m *Manager) processQueryPair(req *http.Request, cellAResp, cellBResp *Resp sampleStored := false if m.storage != nil { - if err := m.storage.StoreQuerySample(ctx, sample); err != nil { + if err := m.storage.StoreQuerySample(ctx, sample, &result); err != nil { level.Error(m.logger).Log("msg", "failed to store query sample", "correlation_id", correlationID, "err", err) m.metrics.storageOperations.WithLabelValues("store_sample", "error").Inc() } else { @@ -426,7 +426,7 @@ type ResponseData struct { } // CaptureResponse captures response data for comparison including trace ID and span ID -func CaptureResponse(resp *http.Response, duration time.Duration, traceID, spanID string) (*ResponseData, error) { +func CaptureResponse(resp *http.Response, duration time.Duration, traceID, spanID string, logger log.Logger) (*ResponseData, error) { body, err := io.ReadAll(resp.Body) if err != nil { return nil, err @@ -446,7 +446,7 @@ func CaptureResponse(resp *http.Response, duration time.Duration, traceID, spanI stats, hash, size, usedNewEngine, err = extractor.ExtractResponseData(body, duration.Milliseconds()) if err != nil { // Log error but don't fail the capture - level.Warn(log.NewNopLogger()).Log("msg", "failed to extract response statistics", "err", err) + level.Warn(logger).Log("msg", "failed to extract response statistics", "err", err) } } else { size = int64(len(body)) diff --git a/tools/querytee/goldfish/manager_test.go b/tools/querytee/goldfish/manager_test.go index e57fd08d083f1..a3bb818354bab 100644 --- a/tools/querytee/goldfish/manager_test.go +++ b/tools/querytee/goldfish/manager_test.go @@ -2,6 +2,7 @@ package goldfish import ( "context" + "errors" "io" "net/http" "strings" @@ -11,6 +12,7 @@ import ( "github.com/go-kit/log" "github.com/grafana/loki/v3/pkg/goldfish" "github.com/grafana/loki/v3/pkg/storage/bucket" + "github.com/grafana/loki/v3/tools/querytee/comparator" "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -23,7 +25,8 @@ type mockStorage struct { closed bool } -func (m *mockStorage) StoreQuerySample(_ context.Context, sample *goldfish.QuerySample) error { +func (m *mockStorage) StoreQuerySample(_ context.Context, sample *goldfish.QuerySample, comparison *goldfish.ComparisonResult) error { + sample.ComparisonStatus = comparison.ComparisonStatus m.samples = append(m.samples, *sample) return nil } @@ -213,7 +216,7 @@ func Test_CaptureResponse_withTraceID(t *testing.T) { } // Call CaptureResponse with traceID and empty spanID - data, err := CaptureResponse(resp, time.Duration(100)*time.Millisecond, tt.traceID, "") + data, err := CaptureResponse(resp, time.Duration(100)*time.Millisecond, tt.traceID, "", log.NewNopLogger()) require.NoError(t, err) assert.Equal(t, tt.expected, data.TraceID) @@ -619,3 +622,85 @@ func (m *mockResultStore) Close(context.Context) error { m.closed = true return nil } + +// mockResponseComparator implements ResponsesComparator for testing +type mockResponseComparator struct { + match bool +} + +func (m *mockResponseComparator) Compare(_, _ []byte, _ time.Time) (*comparator.ComparisonSummary, error) { + if m.match { + return &comparator.ComparisonSummary{}, nil + } + return nil, errors.New("comparison failed") +} + +func TestManager_StoreQuerySample_UsesComparatorResult(t *testing.T) { + tests := []struct { + name string + cellAHash string + cellBHash string + comparatorMatch bool + expectedStatus goldfish.ComparisonStatus + }{ + { + name: "hash mismatch with tolerance match", + cellAHash: "hash1", + cellBHash: "hash2", + comparatorMatch: true, + expectedStatus: goldfish.ComparisonStatusMatch, + }, + { + name: "hash mismatch without tolerance match", + cellAHash: "hash1", + cellBHash: "hash2", + comparatorMatch: false, + expectedStatus: goldfish.ComparisonStatusMismatch, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + storage := &mockStorage{} + + mockComparator := &mockResponseComparator{ + match: tt.comparatorMatch, + } + + config := Config{ + Enabled: true, + ResultsStorage: ResultsStorageConfig{ + Mode: ResultsPersistenceModeAll, + }, + PerformanceTolerance: 0.1, + } + + manager, err := NewManager(config, mockComparator, storage, nil, log.NewNopLogger(), prometheus.NewRegistry()) + require.NoError(t, err) + + // Create responses with different hashes (both return 200) + cellAResp := &ResponseData{ + StatusCode: 200, + Hash: tt.cellAHash, + Body: []byte(`{"status":"success","data":{"result":[]}}`), + Stats: goldfish.QueryStats{ExecTimeMs: 100}, + } + cellBResp := &ResponseData{ + StatusCode: 200, + Hash: tt.cellBHash, + Body: []byte(`{"status":"success","data":{"result":[]}}`), + Stats: goldfish.QueryStats{ExecTimeMs: 100}, + } + + req, _ := http.NewRequest("GET", "/loki/api/v1/query_range?query=test", nil) + + // Process the query pair + manager.processQueryPair(req, cellAResp, cellBResp) + + // Verify the stored sample has the correct comparison status from the comparator + require.Len(t, storage.samples, 1) + assert.Equal(t, tt.expectedStatus, storage.samples[0].ComparisonStatus, + "comparison status should match what the comparator returned") + }) + } +} diff --git a/tools/querytee/proxy_endpoint_test.go b/tools/querytee/proxy_endpoint_test.go index d7869d4a01ad5..6a6cbe6ad4d2f 100644 --- a/tools/querytee/proxy_endpoint_test.go +++ b/tools/querytee/proxy_endpoint_test.go @@ -694,7 +694,7 @@ type mockGoldfishStorage struct { results []goldfish.ComparisonResult } -func (m *mockGoldfishStorage) StoreQuerySample(_ context.Context, sample *goldfish.QuerySample) error { +func (m *mockGoldfishStorage) StoreQuerySample(_ context.Context, sample *goldfish.QuerySample, _ *goldfish.ComparisonResult) error { m.samples = append(m.samples, *sample) return nil }