From 27c7f3871ff056041597a77920128af0c88bf3b4 Mon Sep 17 00:00:00 2001 From: rim99 Date: Sun, 24 Nov 2024 10:56:28 +0800 Subject: [PATCH 01/12] refactor: add time window for trace get method in span store interface Signed-off-by: rim99 --- cmd/anonymizer/app/query/query.go | 8 ++- cmd/anonymizer/app/query/query_test.go | 8 +-- .../storageexporter/exporter_test.go | 3 +- .../internal/integration/span_reader.go | 6 +- cmd/query/app/apiv3/gateway_test.go | 15 ++--- cmd/query/app/apiv3/grpc_handler.go | 7 ++- cmd/query/app/apiv3/grpc_handler_test.go | 57 +++++++++++++++++-- cmd/query/app/apiv3/http_gateway.go | 10 +++- cmd/query/app/apiv3/http_gateway_test.go | 2 +- cmd/query/app/grpc_handler.go | 18 +++++- cmd/query/app/grpc_handler_test.go | 35 +++++++----- cmd/query/app/handler_archive_test.go | 16 +++--- cmd/query/app/http_handler.go | 16 +++++- cmd/query/app/http_handler_test.go | 30 +++++----- cmd/query/app/querysvc/query_service.go | 10 ++-- cmd/query/app/querysvc/query_service_test.go | 37 +++++++----- cmd/query/app/server_test.go | 3 +- .../badger/spanstore/read_write_test.go | 10 ++-- plugin/storage/badger/spanstore/reader.go | 4 +- .../badger/spanstore/rw_internal_test.go | 4 +- plugin/storage/blackhole/blackhole.go | 2 +- plugin/storage/blackhole/blackhole_test.go | 2 +- .../storage/cassandra/savetracetest/main.go | 2 +- plugin/storage/cassandra/spanstore/reader.go | 6 +- .../cassandra/spanstore/reader_test.go | 4 +- plugin/storage/es/spanstore/reader.go | 5 +- plugin/storage/es/spanstore/reader_test.go | 26 +++++---- plugin/storage/grpc/shared/archive.go | 6 +- plugin/storage/grpc/shared/archive_test.go | 8 ++- plugin/storage/grpc/shared/grpc_client.go | 6 +- .../storage/grpc/shared/grpc_client_test.go | 45 +++++++++++---- plugin/storage/grpc/shared/grpc_handler.go | 12 +++- .../storage/grpc/shared/grpc_handler_test.go | 14 ++--- plugin/storage/integration/integration.go | 8 +-- plugin/storage/integration/kafka_test.go | 4 +- plugin/storage/memory/memory.go | 4 +- plugin/storage/memory/memory_test.go | 26 +++++---- storage/spanstore/interface.go | 9 ++- storage/spanstore/mocks/Reader.go | 18 +++--- .../spanstoremetrics/decorator_test.go | 8 +-- .../spanstoremetrics/read_metrics.go | 4 +- storage_v2/factoryadapter/writer_test.go | 4 +- 42 files changed, 344 insertions(+), 178 deletions(-) diff --git a/cmd/anonymizer/app/query/query.go b/cmd/anonymizer/app/query/query.go index 755d8317d19..6cca46408ba 100644 --- a/cmd/anonymizer/app/query/query.go +++ b/cmd/anonymizer/app/query/query.go @@ -54,10 +54,12 @@ func (q *Query) QueryTrace(traceID string) ([]model.Span, error) { if err != nil { return nil, fmt.Errorf("failed to convert the provided trace id: %w", err) } - - stream, err := q.client.GetTrace(context.Background(), &api_v2.GetTraceRequest{ + // TODO: add start time & end time + request := api_v2.GetTraceRequest{ TraceID: mTraceID, - }) + } + + stream, err := q.client.GetTrace(context.Background(), &request) if err != nil { return nil, unwrapNotFoundErr(err) } diff --git a/cmd/anonymizer/app/query/query_test.go b/cmd/anonymizer/app/query/query_test.go index 86b7fd55ed8..957af00319e 100644 --- a/cmd/anonymizer/app/query/query_test.go +++ b/cmd/anonymizer/app/query/query_test.go @@ -24,8 +24,8 @@ import ( ) var ( - matchContext = mock.AnythingOfType("*context.valueCtx") - matchTraceID = mock.AnythingOfType("model.TraceID") + matchContext = mock.AnythingOfType("*context.valueCtx") + matchGetTraceParameters = mock.AnythingOfType("spanstore.GetTraceParameters") mockInvalidTraceID = "xyz" mockTraceID = model.NewTraceID(0, 123456) @@ -106,7 +106,7 @@ func TestQueryTrace(t *testing.T) { defer q.Close() t.Run("No error", func(t *testing.T) { - s.spanReader.On("GetTrace", matchContext, matchTraceID).Return( + s.spanReader.On("GetTrace", matchContext, matchGetTraceParameters).Return( mockTraceGRPC, nil).Once() spans, err := q.QueryTrace(mockTraceID.String()) @@ -120,7 +120,7 @@ func TestQueryTrace(t *testing.T) { }) t.Run("Trace not found", func(t *testing.T) { - s.spanReader.On("GetTrace", matchContext, matchTraceID).Return( + s.spanReader.On("GetTrace", matchContext, matchGetTraceParameters).Return( nil, spanstore.ErrTraceNotFound).Once() spans, err := q.QueryTrace(mockTraceID.String()) diff --git a/cmd/jaeger/internal/exporters/storageexporter/exporter_test.go b/cmd/jaeger/internal/exporters/storageexporter/exporter_test.go index 5ebe56c90aa..96fe9936b2a 100644 --- a/cmd/jaeger/internal/exporters/storageexporter/exporter_test.go +++ b/cmd/jaeger/internal/exporters/storageexporter/exporter_test.go @@ -27,6 +27,7 @@ import ( "github.com/jaegertracing/jaeger/plugin/storage/memory" "github.com/jaegertracing/jaeger/storage" factoryMocks "github.com/jaegertracing/jaeger/storage/mocks" + "github.com/jaegertracing/jaeger/storage/spanstore" ) type mockStorageExt struct { @@ -152,7 +153,7 @@ func TestExporter(t *testing.T) { spanReader, err := storageFactory.CreateSpanReader() require.NoError(t, err) requiredTraceID := model.NewTraceID(0, 1) // 00000000000000000000000000000001 - requiredTrace, err := spanReader.GetTrace(ctx, requiredTraceID) + requiredTrace, err := spanReader.GetTrace(ctx, spanstore.GetTraceParameters{TraceID: requiredTraceID}) require.NoError(t, err) assert.Equal(t, spanID.String(), requiredTrace.Spans[0].SpanID.String()) diff --git a/cmd/jaeger/internal/integration/span_reader.go b/cmd/jaeger/internal/integration/span_reader.go index abd98ecd6aa..3aaf0c09ece 100644 --- a/cmd/jaeger/internal/integration/span_reader.go +++ b/cmd/jaeger/internal/integration/span_reader.go @@ -66,9 +66,11 @@ func unwrapNotFoundErr(err error) error { return err } -func (r *spanReader) GetTrace(ctx context.Context, traceID model.TraceID) (*model.Trace, error) { +func (r *spanReader) GetTrace(ctx context.Context, query spanstore.GetTraceParameters) (*model.Trace, error) { stream, err := r.client.GetTrace(ctx, &api_v2.GetTraceRequest{ - TraceID: traceID, + TraceID: query.TraceID, + StartTime: query.StartTime, + EndTime: query.EndTime, }) if err != nil { return nil, unwrapNotFoundErr(err) diff --git a/cmd/query/app/apiv3/gateway_test.go b/cmd/query/app/apiv3/gateway_test.go index 6c942fcfb0f..22a08dc2641 100644 --- a/cmd/query/app/apiv3/gateway_test.go +++ b/cmd/query/app/apiv3/gateway_test.go @@ -80,8 +80,9 @@ func parseResponse(t *testing.T, body []byte, obj gogoproto.Message) { require.NoError(t, gogojsonpb.Unmarshal(bytes.NewBuffer(body), obj)) } -func makeTestTrace() (*model.Trace, model.TraceID) { +func makeTestTrace() (*model.Trace, spanstore.GetTraceParameters) { traceID := model.NewTraceID(150, 160) + query := spanstore.GetTraceParameters{TraceID: traceID} return &model.Trace{ Spans: []*model.Span{ { @@ -94,7 +95,7 @@ func makeTestTrace() (*model.Trace, model.TraceID) { }, }, }, - }, traceID + }, query } func runGatewayTests( @@ -140,18 +141,18 @@ func (gw *testGateway) runGatewayGetOperations(t *testing.T) { } func (gw *testGateway) runGatewayGetTrace(t *testing.T) { - trace, traceID := makeTestTrace() - gw.reader.On("GetTrace", matchContext, traceID).Return(trace, nil).Once() - gw.getTracesAndVerify(t, "/api/v3/traces/"+traceID.String(), traceID) + trace, query := makeTestTrace() + gw.reader.On("GetTrace", matchContext, query).Return(trace, nil).Once() + gw.getTracesAndVerify(t, "/api/v3/traces/"+query.TraceID.String(), query.TraceID) } func (gw *testGateway) runGatewayFindTraces(t *testing.T) { - trace, traceID := makeTestTrace() + trace, query := makeTestTrace() q, qp := mockFindQueries() gw.reader. On("FindTraces", matchContext, qp). Return([]*model.Trace{trace}, nil).Once() - gw.getTracesAndVerify(t, "/api/v3/traces?"+q.Encode(), traceID) + gw.getTracesAndVerify(t, "/api/v3/traces?"+q.Encode(), query.TraceID) } func (gw *testGateway) getTracesAndVerify(t *testing.T, url string, expectedTraceID model.TraceID) { diff --git a/cmd/query/app/apiv3/grpc_handler.go b/cmd/query/app/apiv3/grpc_handler.go index f27f5560c89..dfcf0480d35 100644 --- a/cmd/query/app/apiv3/grpc_handler.go +++ b/cmd/query/app/apiv3/grpc_handler.go @@ -32,7 +32,12 @@ func (h *Handler) GetTrace(request *api_v3.GetTraceRequest, stream api_v3.QueryS return fmt.Errorf("malform trace ID: %w", err) } - trace, err := h.QueryService.GetTrace(stream.Context(), traceID) + query := spanstore.GetTraceParameters{ + TraceID: traceID, + StartTime: request.GetStartTime(), + EndTime: request.GetEndTime(), + } + trace, err := h.QueryService.GetTrace(stream.Context(), query) if err != nil { return fmt.Errorf("cannot retrieve trace: %w", err) } diff --git a/cmd/query/app/apiv3/grpc_handler_test.go b/cmd/query/app/apiv3/grpc_handler_test.go index 7aaccb7c5d8..84f342eac68 100644 --- a/cmd/query/app/apiv3/grpc_handler_test.go +++ b/cmd/query/app/apiv3/grpc_handler_test.go @@ -26,8 +26,8 @@ import ( ) var ( - matchContext = mock.AnythingOfType("*context.valueCtx") - matchTraceID = mock.AnythingOfType("model.TraceID") + matchContext = mock.AnythingOfType("*context.valueCtx") + matchGetTraceParameters = mock.AnythingOfType("spanstore.GetTraceParameters") ) func newGrpcServer(t *testing.T, handler *Handler) (*grpc.Server, net.Addr) { @@ -79,7 +79,50 @@ func newTestServerClient(t *testing.T) *testServerClient { func TestGetTrace(t *testing.T) { tsc := newTestServerClient(t) - tsc.reader.On("GetTrace", matchContext, matchTraceID).Return( + traceIdLiteral := "156" + traceId, _ := model.TraceIDFromString(traceIdLiteral) + + expectedTraceGetParameters := spanstore.GetTraceParameters{ + TraceID: traceId, + StartTime: time.Time{}, + EndTime: time.Time{}, + } + tsc.reader.On("GetTrace", matchContext, expectedTraceGetParameters).Return( + &model.Trace{ + Spans: []*model.Span{ + { + OperationName: "foobar", + }, + }, + }, nil).Once() + + getTraceStream, err := tsc.client.GetTrace(context.Background(), + &api_v3.GetTraceRequest{ + TraceId: traceIdLiteral, + }, + ) + require.NoError(t, err) + recv, err := getTraceStream.Recv() + require.NoError(t, err) + td := recv.ToTraces() + require.EqualValues(t, 1, td.SpanCount()) + assert.Equal(t, "foobar", + td.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0).Name()) +} + +func TestGetTraceWithTimeWindow(t *testing.T) { + tsc := newTestServerClient(t) + traceIdLiteral := "156" + traceId, _ := model.TraceIDFromString(traceIdLiteral) + + start_ts := time.Unix(1, 2).UTC() + end_ts := time.Unix(3, 4).UTC() + expectedTraceGetParameters := spanstore.GetTraceParameters{ + TraceID: traceId, + StartTime: start_ts, + EndTime: end_ts, + } + tsc.reader.On("GetTrace", matchContext, expectedTraceGetParameters).Return( &model.Trace{ Spans: []*model.Span{ { @@ -90,7 +133,9 @@ func TestGetTrace(t *testing.T) { getTraceStream, err := tsc.client.GetTrace(context.Background(), &api_v3.GetTraceRequest{ - TraceId: "156", + TraceId: traceIdLiteral, + StartTime: start_ts, + EndTime: end_ts, }, ) require.NoError(t, err) @@ -104,7 +149,7 @@ func TestGetTrace(t *testing.T) { func TestGetTraceStorageError(t *testing.T) { tsc := newTestServerClient(t) - tsc.reader.On("GetTrace", matchContext, matchTraceID).Return( + tsc.reader.On("GetTrace", matchContext, matchGetTraceParameters).Return( nil, errors.New("storage_error")).Once() getTraceStream, err := tsc.client.GetTrace(context.Background(), &api_v3.GetTraceRequest{ @@ -118,7 +163,7 @@ func TestGetTraceStorageError(t *testing.T) { func TestGetTraceTraceIDError(t *testing.T) { tsc := newTestServerClient(t) - tsc.reader.On("GetTrace", matchContext, matchTraceID).Return( + tsc.reader.On("GetTrace", matchContext, matchGetTraceParameters).Return( &model.Trace{ Spans: []*model.Span{}, }, nil).Once() diff --git a/cmd/query/app/apiv3/http_gateway.go b/cmd/query/app/apiv3/http_gateway.go index 1abbcece1f8..534101e095d 100644 --- a/cmd/query/app/apiv3/http_gateway.go +++ b/cmd/query/app/apiv3/http_gateway.go @@ -27,7 +27,9 @@ import ( ) const ( - paramTraceID = "trace_id" // get trace by ID + paramTraceID = "trace_id" // get trace by ID + paramStartTime = "start_time" + paramEndTime = "end_time" paramServiceName = "query.service_name" // find traces paramOperationName = "query.operation_name" paramTimeMin = "query.start_time_min" @@ -136,7 +138,11 @@ func (h *HTTPGateway) getTrace(w http.ResponseWriter, r *http.Request) { if h.tryParamError(w, err, paramTraceID) { return } - trc, err := h.QueryService.GetTrace(r.Context(), traceID) + // TODO: add start time & end time + request := spanstore.GetTraceParameters{ + TraceID: traceID, + } + trc, err := h.QueryService.GetTrace(r.Context(), request) if h.tryHandleError(w, err, http.StatusInternalServerError) { return } diff --git a/cmd/query/app/apiv3/http_gateway_test.go b/cmd/query/app/apiv3/http_gateway_test.go index 6feaba989c6..3dda85d67d9 100644 --- a/cmd/query/app/apiv3/http_gateway_test.go +++ b/cmd/query/app/apiv3/http_gateway_test.go @@ -118,7 +118,7 @@ func TestHTTPGatewayGetTraceErrors(t *testing.T) { // error from span reader const simErr = "simulated error" gw.reader. - On("GetTrace", matchContext, matchTraceID). + On("GetTrace", matchContext, matchGetTraceParameters). Return(nil, errors.New(simErr)).Once() r, err = http.NewRequest(http.MethodGet, "/api/v3/traces/123", nil) diff --git a/cmd/query/app/grpc_handler.go b/cmd/query/app/grpc_handler.go index 6f70750e98f..c22780b6c5c 100644 --- a/cmd/query/app/grpc_handler.go +++ b/cmd/query/app/grpc_handler.go @@ -89,7 +89,12 @@ func (g *GRPCHandler) GetTrace(r *api_v2.GetTraceRequest, stream api_v2.QuerySer if r.TraceID == (model.TraceID{}) { return errUninitializedTraceID } - trace, err := g.queryService.GetTrace(stream.Context(), r.TraceID) + query := spanstore.GetTraceParameters{ + TraceID: r.TraceID, + StartTime: r.StartTime, + EndTime: r.EndTime, + } + trace, err := g.queryService.GetTrace(stream.Context(), query) if errors.Is(err, spanstore.ErrTraceNotFound) { g.logger.Warn(msgTraceNotFound, zap.Stringer("id", r.TraceID), zap.Error(err)) return status.Errorf(codes.NotFound, "%s: %v", msgTraceNotFound, err) @@ -109,7 +114,16 @@ func (g *GRPCHandler) ArchiveTrace(ctx context.Context, r *api_v2.ArchiveTraceRe if r.TraceID == (model.TraceID{}) { return nil, errUninitializedTraceID } - err := g.queryService.ArchiveTrace(ctx, r.TraceID) + query := spanstore.GetTraceParameters{ + TraceID: r.TraceID, + } + if r.StartTime != nil { + query.StartTime = *r.StartTime + } + if r.EndTime != nil { + query.EndTime = *r.EndTime + } + err := g.queryService.ArchiveTrace(ctx, query) if errors.Is(err, spanstore.ErrTraceNotFound) { g.logger.Warn(msgTraceNotFound, zap.Stringer("id", r.TraceID), zap.Error(err)) return nil, status.Errorf(codes.NotFound, "%s: %v", msgTraceNotFound, err) diff --git a/cmd/query/app/grpc_handler_test.go b/cmd/query/app/grpc_handler_test.go index 016ee19b77f..7ca990bc2f8 100644 --- a/cmd/query/app/grpc_handler_test.go +++ b/cmd/query/app/grpc_handler_test.go @@ -197,11 +197,20 @@ func withServerAndClient(t *testing.T, actualTest func(server *grpcServer, clien func TestGetTraceSuccessGRPC(t *testing.T) { withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + startTime := time.Date(2020, time.January, 1, 13, 0, 0, 0, time.UTC) + endTime := time.Date(2020, time.January, 1, 14, 0, 0, 0, time.UTC) + expectedTraceGetParameters := spanstore.GetTraceParameters{ + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, + } + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedTraceGetParameters). Return(mockTrace, nil).Once() res, err := client.GetTrace(context.Background(), &api_v2.GetTraceRequest{ - TraceID: mockTraceID, + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, }) spanResChunk, _ := res.Recv() @@ -220,7 +229,7 @@ func assertGRPCError(t *testing.T, err error, code codes.Code, msg string) { func TestGetTraceEmptyTraceIDFailure_GRPC(t *testing.T) { withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() res, err := client.GetTrace(context.Background(), &api_v2.GetTraceRequest{ @@ -237,7 +246,7 @@ func TestGetTraceEmptyTraceIDFailure_GRPC(t *testing.T) { func TestGetTraceDBFailureGRPC(t *testing.T) { withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, errStorageGRPC).Once() res, err := client.GetTrace(context.Background(), &api_v2.GetTraceRequest{ @@ -253,10 +262,10 @@ func TestGetTraceDBFailureGRPC(t *testing.T) { func TestGetTraceNotFoundGRPC(t *testing.T) { withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() - server.archiveSpanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + server.archiveSpanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() res, err := client.GetTrace(context.Background(), &api_v2.GetTraceRequest{ @@ -278,7 +287,7 @@ func TestGetTraceNilRequestOnHandlerGRPC(t *testing.T) { func TestArchiveTraceSuccessGRPC(t *testing.T) { withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). Return(nil).Times(2) @@ -293,9 +302,9 @@ func TestArchiveTraceSuccessGRPC(t *testing.T) { func TestArchiveTraceNotFoundGRPC(t *testing.T) { withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() - server.archiveSpanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + server.archiveSpanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() _, err := client.ArchiveTrace(context.Background(), &api_v2.ArchiveTraceRequest{ @@ -324,7 +333,7 @@ func TestArchiveTraceNilRequestOnHandlerGRPC(t *testing.T) { func TestArchiveTraceFailureGRPC(t *testing.T) { withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). Return(errStorageGRPC).Times(2) @@ -955,7 +964,7 @@ func TestSearchTenancyGRPC(t *testing.T) { Enabled: true, }) withTenantedServerAndClient(t, tm, func(server *grpcServer, client *grpcClient) { - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() // First try without tenancy header @@ -1012,7 +1021,7 @@ func TestSearchTenancyGRPCExplicitList(t *testing.T) { Tenants: []string{"mercury", "venus", "mars"}, }) withTenantedServerAndClient(t, tm, func(server *grpcServer, client *grpcClient) { - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() for _, tc := range []struct { @@ -1124,7 +1133,7 @@ func TestTenancyContextFlowGRPC(t *testing.T) { return false } return true - }), mock.AnythingOfType("model.TraceID")).Return(trace, err).Once() + }), mock.AnythingOfType("spanstore.GetTraceParameters")).Return(trace, err).Once() } for tenant, expected := range allExpectedResults { diff --git a/cmd/query/app/handler_archive_test.go b/cmd/query/app/handler_archive_test.go index 78a0afba11f..87ae3e304b2 100644 --- a/cmd/query/app/handler_archive_test.go +++ b/cmd/query/app/handler_archive_test.go @@ -20,7 +20,7 @@ import ( func TestGetArchivedTrace_NotFound(t *testing.T) { mockReader := &spanstoremocks.Reader{} - mockReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + mockReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() for _, tc := range []struct { name string @@ -37,7 +37,7 @@ func TestGetArchivedTrace_NotFound(t *testing.T) { } { t.Run(tc.name, func(t *testing.T) { withTestServer(t, func(ts *testServer) { - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() var response structuredResponse err := getJSON(ts.server.URL+"/api/traces/"+mockTraceID.String(), &response) @@ -52,11 +52,11 @@ func TestGetArchivedTrace_NotFound(t *testing.T) { func TestGetArchivedTraceSuccess(t *testing.T) { traceID := model.NewTraceID(0, 123456) mockReader := &spanstoremocks.Reader{} - mockReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + mockReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() withTestServer(t, func(ts *testServer) { // make main reader return NotFound - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() var response structuredTraceResponse err := getJSON(ts.server.URL+"/api/traces/"+mockTraceID.String(), &response) @@ -79,13 +79,13 @@ func TestArchiveTrace_BadTraceID(t *testing.T) { // Test return of 404 when trace is not found in APIHandler.archive func TestArchiveTrace_TraceNotFound(t *testing.T) { mockReader := &spanstoremocks.Reader{} - mockReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + mockReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() mockWriter := &spanstoremocks.Writer{} // Not actually going to write the trace, so no need to define mockWriter action withTestServer(t, func(ts *testServer) { // make main reader return NotFound - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() var response structuredResponse err := postJSON(ts.server.URL+"/api/archive/"+mockTraceID.String(), []string{}, &response) @@ -106,7 +106,7 @@ func TestArchiveTrace_Success(t *testing.T) { mockWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). Return(nil).Times(2) withTestServer(t, func(ts *testServer) { - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() var response structuredResponse err := postJSON(ts.server.URL+"/api/archive/"+mockTraceID.String(), []string{}, &response) @@ -119,7 +119,7 @@ func TestArchiveTrace_WriteErrors(t *testing.T) { mockWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). Return(errors.New("cannot save")).Times(2) withTestServer(t, func(ts *testServer) { - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() var response structuredResponse err := postJSON(ts.server.URL+"/api/archive/"+mockTraceID.String(), []string{}, &response) diff --git a/cmd/query/app/http_handler.go b/cmd/query/app/http_handler.go index 3eec204720e..97258e3161b 100644 --- a/cmd/query/app/http_handler.go +++ b/cmd/query/app/http_handler.go @@ -266,7 +266,10 @@ func (aH *APIHandler) tracesByIDs(ctx context.Context, traceIDs []model.TraceID) var traceErrors []structuredError retMe := make([]*model.Trace, 0, len(traceIDs)) for _, traceID := range traceIDs { - if trc, err := aH.queryService.GetTrace(ctx, traceID); err != nil { + query := spanstore.GetTraceParameters{ + TraceID: traceID, + } + if trc, err := aH.queryService.GetTrace(ctx, query); err != nil { if !errors.Is(err, spanstore.ErrTraceNotFound) { return nil, nil, err } @@ -429,7 +432,10 @@ func (aH *APIHandler) getTrace(w http.ResponseWriter, r *http.Request) { if !ok { return } - trc, err := aH.queryService.GetTrace(r.Context(), traceID) + query := spanstore.GetTraceParameters{ + TraceID: traceID, + } + trc, err := aH.queryService.GetTrace(r.Context(), query) if errors.Is(err, spanstore.ErrTraceNotFound) { aH.handleError(w, err, http.StatusNotFound) return @@ -458,7 +464,11 @@ func (aH *APIHandler) archiveTrace(w http.ResponseWriter, r *http.Request) { } // QueryService.ArchiveTrace can now archive this traceID. - err := aH.queryService.ArchiveTrace(r.Context(), traceID) + // TODO: add start time & end time + query := spanstore.GetTraceParameters{ + TraceID: traceID, + } + err := aH.queryService.ArchiveTrace(r.Context(), query) if errors.Is(err, spanstore.ErrTraceNotFound) { aH.handleError(w, err, http.StatusNotFound) return diff --git a/cmd/query/app/http_handler_test.go b/cmd/query/app/http_handler_test.go index 5da9482af87..32c15d6aa91 100644 --- a/cmd/query/app/http_handler_test.go +++ b/cmd/query/app/http_handler_test.go @@ -153,7 +153,7 @@ func withTestServer(t *testing.T, doTest func(s *testServer), queryOptions query func TestGetTraceSuccess(t *testing.T) { ts := initializeTestServer(t) - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() var response structuredResponse @@ -180,7 +180,7 @@ func TestGetTraceDedupeSuccess(t *testing.T) { } ts := initializeTestServer(t) - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(dupedMockTrace, nil).Once() var response structuredResponse @@ -321,7 +321,7 @@ func TestGetTrace(t *testing.T) { ts := initializeTestServer(t, HandlerOptions.Tracer(jTracer.OTEL)) - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), model.NewTraceID(0, 0x123456abc)). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), spanstore.GetTraceParameters{TraceID: model.NewTraceID(0, 0x123456abc)}). Return(makeMockTrace(t), nil).Once() var response structuredResponse @@ -338,7 +338,7 @@ func TestGetTrace(t *testing.T) { func TestGetTraceDBFailure(t *testing.T) { ts := initializeTestServer(t) - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, errStorage).Once() var response structuredResponse @@ -348,7 +348,7 @@ func TestGetTraceDBFailure(t *testing.T) { func TestGetTraceNotFound(t *testing.T) { ts := initializeTestServer(t) - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() var response structuredResponse @@ -365,7 +365,7 @@ func TestGetTraceAdjustmentFailure(t *testing.T) { }), }, ) - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() var response structuredResponse @@ -396,7 +396,7 @@ func TestSearchSuccess(t *testing.T) { func TestSearchByTraceIDSuccess(t *testing.T) { ts := initializeTestServer(t) - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Twice() var response structuredResponse @@ -411,9 +411,9 @@ func TestSearchByTraceIDSuccessWithArchive(t *testing.T) { ts := initializeTestServerWithOptions(t, &tenancy.Manager{}, querysvc.QueryServiceOptions{ ArchiveSpanReader: archiveReadMock, }) - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Twice() - archiveReadMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + archiveReadMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Twice() var response structuredResponse @@ -425,7 +425,7 @@ func TestSearchByTraceIDSuccessWithArchive(t *testing.T) { func TestSearchByTraceIDNotFound(t *testing.T) { ts := initializeTestServer(t) - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() var response structuredResponse @@ -438,7 +438,7 @@ func TestSearchByTraceIDNotFound(t *testing.T) { func TestSearchByTraceIDFailure(t *testing.T) { ts := initializeTestServer(t) whatsamattayou := "whatsamattayou" - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, errors.New(whatsamattayou)).Once() var response structuredResponse @@ -912,7 +912,7 @@ func TestSearchTenancyHTTP(t *testing.T) { ts := initializeTestServerWithOptions(t, tenancy.NewManager(&tenancyOptions), querysvc.QueryServiceOptions{}) - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Twice() var response structuredResponse @@ -936,7 +936,7 @@ func TestSearchTenancyRejectionHTTP(t *testing.T) { Enabled: true, } ts := initializeTestServerWithOptions(t, tenancy.NewManager(&tenancyOptions), querysvc.QueryServiceOptions{}) - ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Twice() req, err := http.NewRequest(http.MethodGet, ts.server.URL+`/api/traces?traceID=1&traceID=2`, nil) @@ -968,14 +968,14 @@ func TestSearchTenancyFlowTenantHTTP(t *testing.T) { return false } return true - }), mock.AnythingOfType("model.TraceID")).Return(mockTrace, nil).Twice() + }), mock.AnythingOfType("spanstore.GetTraceParameters")).Return(mockTrace, nil).Twice() ts.spanReader.On("GetTrace", mock.MatchedBy(func(v any) bool { ctx, ok := v.(context.Context) if !ok || tenancy.GetTenant(ctx) != "megacorp" { return false } return true - }), mock.AnythingOfType("model.TraceID")).Return(nil, errStorage).Once() + }), mock.AnythingOfType("spanstore.GetTraceParameters")).Return(nil, errStorage).Once() var responseAcme structuredResponse err := getJSONCustomHeaders( diff --git a/cmd/query/app/querysvc/query_service.go b/cmd/query/app/querysvc/query_service.go index 135ecc60bbe..4f1939ce88c 100644 --- a/cmd/query/app/querysvc/query_service.go +++ b/cmd/query/app/querysvc/query_service.go @@ -60,13 +60,13 @@ func NewQueryService(spanReader spanstore.Reader, dependencyReader dependencysto } // GetTrace is the queryService implementation of spanstore.Reader.GetTrace -func (qs QueryService) GetTrace(ctx context.Context, traceID model.TraceID) (*model.Trace, error) { - trace, err := qs.spanReader.GetTrace(ctx, traceID) +func (qs QueryService) GetTrace(ctx context.Context, query spanstore.GetTraceParameters) (*model.Trace, error) { + trace, err := qs.spanReader.GetTrace(ctx, query) if errors.Is(err, spanstore.ErrTraceNotFound) { if qs.options.ArchiveSpanReader == nil { return nil, err } - trace, err = qs.options.ArchiveSpanReader.GetTrace(ctx, traceID) + trace, err = qs.options.ArchiveSpanReader.GetTrace(ctx, query) } return trace, err } @@ -90,11 +90,11 @@ func (qs QueryService) FindTraces(ctx context.Context, query *spanstore.TraceQue } // ArchiveTrace is the queryService utility to archive traces. -func (qs QueryService) ArchiveTrace(ctx context.Context, traceID model.TraceID) error { +func (qs QueryService) ArchiveTrace(ctx context.Context, query spanstore.GetTraceParameters) error { if qs.options.ArchiveSpanWriter == nil { return errNoArchiveSpanStorage } - trace, err := qs.GetTrace(ctx, traceID) + trace, err := qs.GetTrace(ctx, query) if err != nil { return err } diff --git a/cmd/query/app/querysvc/query_service_test.go b/cmd/query/app/querysvc/query_service_test.go index 582a4568509..27a8949ddf2 100644 --- a/cmd/query/app/querysvc/query_service_test.go +++ b/cmd/query/app/querysvc/query_service_test.go @@ -107,12 +107,13 @@ func initializeTestService(optionAppliers ...testOption) *testQueryService { // Test QueryService.GetTrace() func TestGetTraceSuccess(t *testing.T) { tqs := initializeTestService() - tqs.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + tqs.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() type contextKey string ctx := context.Background() - res, err := tqs.queryService.GetTrace(context.WithValue(ctx, contextKey("foo"), "bar"), mockTraceID) + query := spanstore.GetTraceParameters{TraceID: mockTraceID} + res, err := tqs.queryService.GetTrace(context.WithValue(ctx, contextKey("foo"), "bar"), query) require.NoError(t, err) assert.Equal(t, res, mockTrace) } @@ -120,26 +121,28 @@ func TestGetTraceSuccess(t *testing.T) { // Test QueryService.GetTrace() without ArchiveSpanReader func TestGetTraceNotFound(t *testing.T) { tqs := initializeTestService() - tqs.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + tqs.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() type contextKey string ctx := context.Background() - _, err := tqs.queryService.GetTrace(context.WithValue(ctx, contextKey("foo"), "bar"), mockTraceID) + query := spanstore.GetTraceParameters{TraceID: mockTraceID} + _, err := tqs.queryService.GetTrace(context.WithValue(ctx, contextKey("foo"), "bar"), query) assert.Equal(t, err, spanstore.ErrTraceNotFound) } // Test QueryService.GetTrace() with ArchiveSpanReader func TestGetTraceFromArchiveStorage(t *testing.T) { tqs := initializeTestService(withArchiveSpanReader()) - tqs.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + tqs.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() - tqs.archiveSpanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + tqs.archiveSpanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() type contextKey string ctx := context.Background() - res, err := tqs.queryService.GetTrace(context.WithValue(ctx, contextKey("foo"), "bar"), mockTraceID) + query := spanstore.GetTraceParameters{TraceID: mockTraceID} + res, err := tqs.queryService.GetTrace(context.WithValue(ctx, contextKey("foo"), "bar"), query) require.NoError(t, err) assert.Equal(t, res, mockTrace) } @@ -202,35 +205,38 @@ func TestArchiveTraceNoOptions(t *testing.T) { type contextKey string ctx := context.Background() - err := tqs.queryService.ArchiveTrace(context.WithValue(ctx, contextKey("foo"), "bar"), mockTraceID) + query := spanstore.GetTraceParameters{TraceID: mockTraceID} + err := tqs.queryService.ArchiveTrace(context.WithValue(ctx, contextKey("foo"), "bar"), query) assert.Equal(t, errNoArchiveSpanStorage, err) } // Test QueryService.ArchiveTrace() with ArchiveSpanWriter but invalid traceID. func TestArchiveTraceWithInvalidTraceID(t *testing.T) { tqs := initializeTestService(withArchiveSpanReader(), withArchiveSpanWriter()) - tqs.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + tqs.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() - tqs.archiveSpanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + tqs.archiveSpanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(nil, spanstore.ErrTraceNotFound).Once() type contextKey string ctx := context.Background() - err := tqs.queryService.ArchiveTrace(context.WithValue(ctx, contextKey("foo"), "bar"), mockTraceID) + query := spanstore.GetTraceParameters{TraceID: mockTraceID} + err := tqs.queryService.ArchiveTrace(context.WithValue(ctx, contextKey("foo"), "bar"), query) assert.Equal(t, spanstore.ErrTraceNotFound, err) } // Test QueryService.ArchiveTrace(), save error with ArchiveSpanWriter. func TestArchiveTraceWithArchiveWriterError(t *testing.T) { tqs := initializeTestService(withArchiveSpanWriter()) - tqs.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + tqs.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() tqs.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). Return(errors.New("cannot save")).Times(2) type contextKey string ctx := context.Background() - joinErr := tqs.queryService.ArchiveTrace(context.WithValue(ctx, contextKey("foo"), "bar"), mockTraceID) + query := spanstore.GetTraceParameters{TraceID: mockTraceID} + joinErr := tqs.queryService.ArchiveTrace(context.WithValue(ctx, contextKey("foo"), "bar"), query) // There are two spans in the mockTrace, ArchiveTrace should return a wrapped error. require.EqualError(t, joinErr, "cannot save\ncannot save") } @@ -238,14 +244,15 @@ func TestArchiveTraceWithArchiveWriterError(t *testing.T) { // Test QueryService.ArchiveTrace() with correctly configured ArchiveSpanWriter. func TestArchiveTraceSuccess(t *testing.T) { tqs := initializeTestService(withArchiveSpanWriter()) - tqs.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")). + tqs.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). Return(mockTrace, nil).Once() tqs.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). Return(nil).Times(2) type contextKey string ctx := context.Background() - err := tqs.queryService.ArchiveTrace(context.WithValue(ctx, contextKey("foo"), "bar"), mockTraceID) + query := spanstore.GetTraceParameters{TraceID: mockTraceID} + err := tqs.queryService.ArchiveTrace(context.WithValue(ctx, contextKey("foo"), "bar"), query) require.NoError(t, err) } diff --git a/cmd/query/app/server_test.go b/cmd/query/app/server_test.go index 7a7d1bed674..097f1317189 100644 --- a/cmd/query/app/server_test.go +++ b/cmd/query/app/server_test.go @@ -44,6 +44,7 @@ import ( "github.com/jaegertracing/jaeger/ports" "github.com/jaegertracing/jaeger/proto-gen/api_v2" depsmocks "github.com/jaegertracing/jaeger/storage/dependencystore/mocks" + "github.com/jaegertracing/jaeger/storage/spanstore" spanstoremocks "github.com/jaegertracing/jaeger/storage/spanstore/mocks" ) @@ -901,7 +902,7 @@ func TestServerHTTP_TracesRequest(t *testing.T) { tenancyMgr := tenancy.NewManager(&serverOptions.Tenancy) querySvc := makeQuerySvc() - querySvc.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), model.NewTraceID(0, 0x123456abc)). + querySvc.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), spanstore.GetTraceParameters{TraceID: model.NewTraceID(0, 0x123456abc)}). Return(makeMockTrace(t), nil).Once() telset := initTelSet(zaptest.NewLogger(t), &tracer, healthcheck.New()) diff --git a/plugin/storage/badger/spanstore/read_write_test.go b/plugin/storage/badger/spanstore/read_write_test.go index 9c4d8f22be7..b324a9ffccc 100644 --- a/plugin/storage/badger/spanstore/read_write_test.go +++ b/plugin/storage/badger/spanstore/read_write_test.go @@ -68,10 +68,10 @@ func TestWriteReadBack(t *testing.T) { } for i := 0; i < traces; i++ { - tr, err := sr.GetTrace(context.Background(), model.TraceID{ + tr, err := sr.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{ Low: uint64(i), High: 1, - }) + }}) require.NoError(t, err) assert.Len(t, tr.Spans, spans) @@ -291,7 +291,7 @@ func TestFindNothing(t *testing.T) { require.NoError(t, err) assert.Empty(t, trs) - tr, err := sr.GetTrace(context.Background(), model.TraceID{High: 0, Low: 0}) + tr, err := sr.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{Low: 0, High: 0}}) assert.Equal(t, spanstore.ErrTraceNotFound, err) assert.Nil(t, tr) }) @@ -418,10 +418,10 @@ func TestPersist(t *testing.T) { }) p(t, dir, func(t *testing.T, _ spanstore.Writer, sr spanstore.Reader) { - trace, err := sr.GetTrace(context.Background(), model.TraceID{ + trace, err := sr.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{ Low: uint64(1), High: 1, - }) + }}) require.NoError(t, err) assert.Equal(t, "operation-p", trace.Spans[0].OperationName) diff --git a/plugin/storage/badger/spanstore/reader.go b/plugin/storage/badger/spanstore/reader.go index 84c5212d775..810f774a2c2 100644 --- a/plugin/storage/badger/spanstore/reader.go +++ b/plugin/storage/badger/spanstore/reader.go @@ -146,8 +146,8 @@ func (r *TraceReader) getTraces(traceIDs []model.TraceID) ([]*model.Trace, error } // GetTrace takes a traceID and returns a Trace associated with that traceID -func (r *TraceReader) GetTrace(_ context.Context, traceID model.TraceID) (*model.Trace, error) { - traces, err := r.getTraces([]model.TraceID{traceID}) +func (r *TraceReader) GetTrace(_ context.Context, query spanstore.GetTraceParameters) (*model.Trace, error) { + traces, err := r.getTraces([]model.TraceID{query.TraceID}) if err != nil { return nil, err } diff --git a/plugin/storage/badger/spanstore/rw_internal_test.go b/plugin/storage/badger/spanstore/rw_internal_test.go index 60828a8b80a..01a78a820cd 100644 --- a/plugin/storage/badger/spanstore/rw_internal_test.go +++ b/plugin/storage/badger/spanstore/rw_internal_test.go @@ -31,7 +31,7 @@ func TestEncodingTypes(t *testing.T) { err := sw.WriteSpan(context.Background(), &testSpan) require.NoError(t, err) - tr, err := rw.GetTrace(context.Background(), model.TraceID{Low: 0, High: 1}) + tr, err := rw.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{Low: 0, High: 1}}) require.NoError(t, err) assert.Len(t, tr.Spans, 1) }) @@ -74,7 +74,7 @@ func TestEncodingTypes(t *testing.T) { return nil }) - _, err = rw.GetTrace(context.Background(), model.TraceID{Low: 0, High: 1}) + _, err = rw.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{Low: 0, High: 1}}) require.EqualError(t, err, "unknown encoding type: 0x04") }) } diff --git a/plugin/storage/blackhole/blackhole.go b/plugin/storage/blackhole/blackhole.go index 5ffa51bb3af..e75d2b1be17 100644 --- a/plugin/storage/blackhole/blackhole.go +++ b/plugin/storage/blackhole/blackhole.go @@ -36,7 +36,7 @@ func (*Store) WriteSpan(context.Context, *model.Span) error { } // GetTrace gets nothing. -func (*Store) GetTrace(context.Context, model.TraceID) (*model.Trace, error) { +func (*Store) GetTrace(context.Context, spanstore.GetTraceParameters) (*model.Trace, error) { return nil, spanstore.ErrTraceNotFound } diff --git a/plugin/storage/blackhole/blackhole_test.go b/plugin/storage/blackhole/blackhole_test.go index 236ce82406d..3c71fa9896d 100644 --- a/plugin/storage/blackhole/blackhole_test.go +++ b/plugin/storage/blackhole/blackhole_test.go @@ -37,7 +37,7 @@ func TestStoreWriteSpan(t *testing.T) { func TestStoreGetTrace(t *testing.T) { withBlackhole(func(store *Store) { - trace, err := store.GetTrace(context.Background(), model.NewTraceID(1, 2)) + trace, err := store.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.NewTraceID(1, 2)}) require.Error(t, err) assert.Nil(t, trace) }) diff --git a/plugin/storage/cassandra/savetracetest/main.go b/plugin/storage/cassandra/savetracetest/main.go index d2190f5511e..8d6f0d5e721 100644 --- a/plugin/storage/cassandra/savetracetest/main.go +++ b/plugin/storage/cassandra/savetracetest/main.go @@ -59,7 +59,7 @@ func main() { logger.Info("Saved span", zap.String("spanID", getSomeSpan().SpanID.String())) } s := getSomeSpan() - trace, err := spanReader.GetTrace(ctx, s.TraceID) + trace, err := spanReader.GetTrace(ctx, spanstore.GetTraceParameters{TraceID: s.TraceID}) if err != nil { logger.Fatal("Failed to read", zap.Error(err)) } else { diff --git a/plugin/storage/cassandra/spanstore/reader.go b/plugin/storage/cassandra/spanstore/reader.go index 88ba7b8f335..e79d8224ae0 100644 --- a/plugin/storage/cassandra/spanstore/reader.go +++ b/plugin/storage/cassandra/spanstore/reader.go @@ -204,8 +204,8 @@ func (s *SpanReader) readTraceInSpan(_ context.Context, traceID dbmodel.TraceID) } // GetTrace takes a traceID and returns a Trace associated with that traceID -func (s *SpanReader) GetTrace(ctx context.Context, traceID model.TraceID) (*model.Trace, error) { - return s.readTrace(ctx, dbmodel.TraceIDFromDomain(traceID)) +func (s *SpanReader) GetTrace(ctx context.Context, query spanstore.GetTraceParameters) (*model.Trace, error) { + return s.readTrace(ctx, dbmodel.TraceIDFromDomain(query.TraceID)) } func validateQuery(p *spanstore.TraceQueryParameters) error { @@ -238,7 +238,7 @@ func (s *SpanReader) FindTraces(ctx context.Context, traceQuery *spanstore.Trace } var retMe []*model.Trace for _, traceID := range uniqueTraceIDs { - jTrace, err := s.GetTrace(ctx, traceID) + jTrace, err := s.GetTrace(ctx, spanstore.GetTraceParameters{TraceID: traceID}) if err != nil { s.logger.Error("Failure to read trace", zap.String("trace_id", traceID.String()), zap.Error(err)) continue diff --git a/plugin/storage/cassandra/spanstore/reader_test.go b/plugin/storage/cassandra/spanstore/reader_test.go index 7dbef606ccb..a55ced3e7dd 100644 --- a/plugin/storage/cassandra/spanstore/reader_test.go +++ b/plugin/storage/cassandra/spanstore/reader_test.go @@ -171,7 +171,7 @@ func TestSpanReaderGetTrace(t *testing.T) { r.session.On("Query", mock.AnythingOfType("string"), matchEverything()).Return(query) - trace, err := r.reader.GetTrace(context.Background(), model.TraceID{}) + trace, err := r.reader.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}) if testCase.expectedErr == "" { require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.NoError(t, err) @@ -197,7 +197,7 @@ func TestSpanReaderGetTrace_TraceNotFound(t *testing.T) { r.session.On("Query", mock.AnythingOfType("string"), matchEverything()).Return(query) - trace, err := r.reader.GetTrace(context.Background(), model.TraceID{}) + trace, err := r.reader.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") assert.Nil(t, trace) require.EqualError(t, err, "trace not found") diff --git a/plugin/storage/es/spanstore/reader.go b/plugin/storage/es/spanstore/reader.go index 47bd6dbcec2..1b3cc39911f 100644 --- a/plugin/storage/es/spanstore/reader.go +++ b/plugin/storage/es/spanstore/reader.go @@ -233,11 +233,12 @@ func timeRangeIndices(indexName, indexDateLayout string, startTime time.Time, en } // GetTrace takes a traceID and returns a Trace associated with that traceID -func (s *SpanReader) GetTrace(ctx context.Context, traceID model.TraceID) (*model.Trace, error) { +func (s *SpanReader) GetTrace(ctx context.Context, query spanstore.GetTraceParameters) (*model.Trace, error) { ctx, span := s.tracer.Start(ctx, "GetTrace") defer span.End() currentTime := time.Now() - traces, err := s.multiRead(ctx, []model.TraceID{traceID}, currentTime.Add(-s.maxSpanAge), currentTime) + // TODO: use start time & end time in "query" struct + traces, err := s.multiRead(ctx, []model.TraceID{query.TraceID}, currentTime.Add(-s.maxSpanAge), currentTime) if err != nil { return nil, es.DetailedError(err) } diff --git a/plugin/storage/es/spanstore/reader_test.go b/plugin/storage/es/spanstore/reader_test.go index 3b292008dcd..a91baafc2bd 100644 --- a/plugin/storage/es/spanstore/reader_test.go +++ b/plugin/storage/es/spanstore/reader_test.go @@ -330,8 +330,8 @@ func TestSpanReader_GetTrace(t *testing.T) { {Hits: searchHits}, }, }, nil) - - trace, err := r.reader.GetTrace(context.Background(), model.NewTraceID(0, 1)) + query := spanstore.GetTraceParameters{TraceID: model.NewTraceID(0, 1)} + trace, err := r.reader.GetTrace(context.Background(), query) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.NoError(t, err) require.NotNil(t, trace) @@ -446,7 +446,8 @@ func TestSpanReader_SearchAfter(t *testing.T) { }, }, nil).Times(2) - trace, err := r.reader.GetTrace(context.Background(), model.NewTraceID(0, 1)) + query := spanstore.GetTraceParameters{TraceID: model.NewTraceID(0, 1)} + trace, err := r.reader.GetTrace(context.Background(), query) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.NoError(t, err) require.NotNil(t, trace) @@ -466,7 +467,8 @@ func TestSpanReader_GetTraceQueryError(t *testing.T) { Return(&elastic.MultiSearchResult{ Responses: []*elastic.SearchResult{}, }, nil) - trace, err := r.reader.GetTrace(context.Background(), model.NewTraceID(0, 1)) + query := spanstore.GetTraceParameters{TraceID: model.NewTraceID(0, 1)} + trace, err := r.reader.GetTrace(context.Background(), query) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.EqualError(t, err, "trace not found") require.Nil(t, trace) @@ -486,7 +488,8 @@ func TestSpanReader_GetTraceNilHits(t *testing.T) { }, }, nil) - trace, err := r.reader.GetTrace(context.Background(), model.NewTraceID(0, 1)) + query := spanstore.GetTraceParameters{TraceID: model.NewTraceID(0, 1)} + trace, err := r.reader.GetTrace(context.Background(), query) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.EqualError(t, err, "trace not found") require.Nil(t, trace) @@ -510,7 +513,8 @@ func TestSpanReader_GetTraceInvalidSpanError(t *testing.T) { }, }, nil) - trace, err := r.reader.GetTrace(context.Background(), model.NewTraceID(0, 1)) + query := spanstore.GetTraceParameters{TraceID: model.NewTraceID(0, 1)} + trace, err := r.reader.GetTrace(context.Background(), query) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.Error(t, err, "invalid span") require.Nil(t, trace) @@ -535,7 +539,8 @@ func TestSpanReader_GetTraceSpanConversionError(t *testing.T) { }, }, nil) - trace, err := r.reader.GetTrace(context.Background(), model.NewTraceID(0, 1)) + query := spanstore.GetTraceParameters{TraceID: model.NewTraceID(0, 1)} + trace, err := r.reader.GetTrace(context.Background(), query) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.Error(t, err, "span conversion error, because lacks elements") require.Nil(t, trace) @@ -1272,8 +1277,8 @@ func TestSpanReader_ArchiveTraces(t *testing.T) { Return(&elastic.MultiSearchResult{ Responses: []*elastic.SearchResult{}, }, nil) - - trace, err := r.reader.GetTrace(context.Background(), model.TraceID{}) + query := spanstore.GetTraceParameters{TraceID: model.TraceID{}} + trace, err := r.reader.GetTrace(context.Background(), query) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.Nil(t, trace) require.EqualError(t, err, "trace not found") @@ -1289,7 +1294,8 @@ func TestSpanReader_ArchiveTraces_ReadAlias(t *testing.T) { Responses: []*elastic.SearchResult{}, }, nil) - trace, err := r.reader.GetTrace(context.Background(), model.TraceID{}) + query := spanstore.GetTraceParameters{TraceID: model.TraceID{}} + trace, err := r.reader.GetTrace(context.Background(), query) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.Nil(t, trace) require.EqualError(t, err, "trace not found") diff --git a/plugin/storage/grpc/shared/archive.go b/plugin/storage/grpc/shared/archive.go index 2bd14685d27..8eeb3b609b2 100644 --- a/plugin/storage/grpc/shared/archive.go +++ b/plugin/storage/grpc/shared/archive.go @@ -32,9 +32,11 @@ type archiveWriter struct { } // GetTrace takes a traceID and returns a Trace associated with that traceID from Archive Storage -func (r *archiveReader) GetTrace(ctx context.Context, traceID model.TraceID) (*model.Trace, error) { +func (r *archiveReader) GetTrace(ctx context.Context, q spanstore.GetTraceParameters) (*model.Trace, error) { stream, err := r.client.GetArchiveTrace(ctx, &storage_v1.GetTraceRequest{ - TraceID: traceID, + TraceID: q.TraceID, + StartTime: q.StartTime, + EndTime: q.EndTime, }) if status.Code(err) == codes.NotFound { return nil, spanstore.ErrTraceNotFound diff --git a/plugin/storage/grpc/shared/archive_test.go b/plugin/storage/grpc/shared/archive_test.go index 75857f9e3cc..3f5128ff590 100644 --- a/plugin/storage/grpc/shared/archive_test.go +++ b/plugin/storage/grpc/shared/archive_test.go @@ -59,7 +59,9 @@ func TestArchiveReader_GetTrace(t *testing.T) { }).Return(traceClient, nil) reader := &archiveReader{client: archiveSpanReader} - trace, err := reader.GetTrace(context.Background(), mockTraceID) + trace, err := reader.GetTrace(context.Background(), spanstore.GetTraceParameters{ + TraceID: mockTraceID, + }) require.NoError(t, err) assert.Equal(t, expected, trace) } @@ -73,7 +75,9 @@ func TestArchiveReaderGetTrace_NoTrace(t *testing.T) { }).Return(nil, status.Errorf(codes.NotFound, "")) reader := &archiveReader{client: archiveSpanReader} - _, err := reader.GetTrace(context.Background(), mockTraceID) + _, err := reader.GetTrace(context.Background(), spanstore.GetTraceParameters{ + TraceID: mockTraceID, + }) assert.Equal(t, spanstore.ErrTraceNotFound, err) } diff --git a/plugin/storage/grpc/shared/grpc_client.go b/plugin/storage/grpc/shared/grpc_client.go index abe296b70e3..f6323be4568 100644 --- a/plugin/storage/grpc/shared/grpc_client.go +++ b/plugin/storage/grpc/shared/grpc_client.go @@ -81,9 +81,11 @@ func (c *GRPCClient) ArchiveSpanWriter() spanstore.Writer { } // GetTrace takes a traceID and returns a Trace associated with that traceID -func (c *GRPCClient) GetTrace(ctx context.Context, traceID model.TraceID) (*model.Trace, error) { +func (c *GRPCClient) GetTrace(ctx context.Context, query spanstore.GetTraceParameters) (*model.Trace, error) { stream, err := c.readerClient.GetTrace(ctx, &storage_v1.GetTraceRequest{ - TraceID: traceID, + TraceID: query.TraceID, + StartTime: query.StartTime, + EndTime: query.EndTime, }) if status.Code(err) == codes.NotFound { return nil, spanstore.ErrTraceNotFound diff --git a/plugin/storage/grpc/shared/grpc_client_test.go b/plugin/storage/grpc/shared/grpc_client_test.go index a8bb88fcf34..85814a3163a 100644 --- a/plugin/storage/grpc/shared/grpc_client_test.go +++ b/plugin/storage/grpc/shared/grpc_client_test.go @@ -157,13 +157,17 @@ func TestGRPCClientGetOperationsV2(t *testing.T) { func TestGRPCClientGetTrace(t *testing.T) { withGRPCClient(func(r *grpcClientTest) { + startTime := time.Date(2020, time.January, 1, 13, 0, 0, 0, time.UTC) + endTime := time.Date(2020, time.January, 1, 14, 0, 0, 0, time.UTC) traceClient := new(grpcMocks.SpanReaderPlugin_GetTraceClient) traceClient.On("Recv").Return(&storage_v1.SpansResponseChunk{ Spans: mockTraceSpans, }, nil).Once() traceClient.On("Recv").Return(nil, io.EOF) r.spanReader.On("GetTrace", mock.Anything, &storage_v1.GetTraceRequest{ - TraceID: mockTraceID, + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, }).Return(traceClient, nil) var expectedSpans []*model.Span @@ -171,7 +175,11 @@ func TestGRPCClientGetTrace(t *testing.T) { expectedSpans = append(expectedSpans, &mockTraceSpans[i]) } - s, err := r.client.GetTrace(context.Background(), mockTraceID) + s, err := r.client.GetTrace(context.Background(), spanstore.GetTraceParameters{ + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, + }) require.NoError(t, err) assert.Equal(t, &model.Trace{ Spans: expectedSpans, @@ -187,7 +195,7 @@ func TestGRPCClientGetTrace_StreamError(t *testing.T) { TraceID: mockTraceID, }).Return(traceClient, nil) - s, err := r.client.GetTrace(context.Background(), mockTraceID) + s, err := r.client.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: mockTraceID}) require.Error(t, err) assert.Nil(t, s) }) @@ -199,7 +207,7 @@ func TestGRPCClientGetTrace_NoTrace(t *testing.T) { TraceID: mockTraceID, }).Return(nil, status.Errorf(codes.NotFound, "")) - s, err := r.client.GetTrace(context.Background(), mockTraceID) + s, err := r.client.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: mockTraceID}) assert.Equal(t, spanstore.ErrTraceNotFound, err) assert.Nil(t, s) }) @@ -215,7 +223,7 @@ func TestGRPCClientGetTrace_StreamErrorTraceNotFound(t *testing.T) { TraceID: mockTraceID, }).Return(traceClient, nil) - s, err := r.client.GetTrace(context.Background(), mockTraceID) + s, err := r.client.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: mockTraceID}) assert.Equal(t, spanstore.ErrTraceNotFound, err) assert.Nil(t, s) }) @@ -364,13 +372,17 @@ func TestGrpcClientStreamWriterWriteSpan(t *testing.T) { func TestGrpcClientGetArchiveTrace(t *testing.T) { withGRPCClient(func(r *grpcClientTest) { + startTime := time.Date(2020, time.January, 1, 13, 0, 0, 0, time.UTC) + endTime := time.Date(2020, time.January, 1, 14, 0, 0, 0, time.UTC) traceClient := new(grpcMocks.ArchiveSpanReaderPlugin_GetArchiveTraceClient) traceClient.On("Recv").Return(&storage_v1.SpansResponseChunk{ Spans: mockTraceSpans, }, nil).Once() traceClient.On("Recv").Return(nil, io.EOF) r.archiveReader.On("GetArchiveTrace", mock.Anything, &storage_v1.GetTraceRequest{ - TraceID: mockTraceID, + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, }).Return(traceClient, nil) var expectedSpans []*model.Span @@ -378,7 +390,11 @@ func TestGrpcClientGetArchiveTrace(t *testing.T) { expectedSpans = append(expectedSpans, &mockTraceSpans[i]) } - s, err := r.client.ArchiveSpanReader().GetTrace(context.Background(), mockTraceID) + s, err := r.client.ArchiveSpanReader().GetTrace(context.Background(), spanstore.GetTraceParameters{ + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, + }) require.NoError(t, err) assert.Equal(t, &model.Trace{ Spans: expectedSpans, @@ -394,7 +410,9 @@ func TestGrpcClientGetArchiveTrace_StreamError(t *testing.T) { TraceID: mockTraceID, }).Return(traceClient, nil) - s, err := r.client.ArchiveSpanReader().GetTrace(context.Background(), mockTraceID) + s, err := r.client.ArchiveSpanReader().GetTrace(context.Background(), spanstore.GetTraceParameters{ + TraceID: mockTraceID, + }) require.Error(t, err) assert.Nil(t, s) }) @@ -406,7 +424,12 @@ func TestGrpcClientGetArchiveTrace_NoTrace(t *testing.T) { TraceID: mockTraceID, }).Return(nil, spanstore.ErrTraceNotFound) - s, err := r.client.ArchiveSpanReader().GetTrace(context.Background(), mockTraceID) + s, err := r.client.ArchiveSpanReader().GetTrace( + context.Background(), + spanstore.GetTraceParameters{ + TraceID: mockTraceID, + }, + ) require.Error(t, err) assert.Nil(t, s) }) @@ -420,7 +443,9 @@ func TestGrpcClientGetArchiveTrace_StreamErrorTraceNotFound(t *testing.T) { TraceID: mockTraceID, }).Return(traceClient, nil) - s, err := r.client.ArchiveSpanReader().GetTrace(context.Background(), mockTraceID) + s, err := r.client.ArchiveSpanReader().GetTrace(context.Background(), spanstore.GetTraceParameters{ + TraceID: mockTraceID, + }) assert.Equal(t, spanstore.ErrTraceNotFound, err) assert.Nil(t, s) }) diff --git a/plugin/storage/grpc/shared/grpc_handler.go b/plugin/storage/grpc/shared/grpc_handler.go index 564a625ebf0..28dc0adf8cb 100644 --- a/plugin/storage/grpc/shared/grpc_handler.go +++ b/plugin/storage/grpc/shared/grpc_handler.go @@ -151,7 +151,11 @@ func (s *GRPCHandler) Close(context.Context, *storage_v1.CloseWriterRequest) (*s // GetTrace takes a traceID and streams a Trace associated with that traceID func (s *GRPCHandler) GetTrace(r *storage_v1.GetTraceRequest, stream storage_v1.SpanReaderPlugin_GetTraceServer) error { - trace, err := s.impl.SpanReader().GetTrace(stream.Context(), r.TraceID) + trace, err := s.impl.SpanReader().GetTrace(stream.Context(), spanstore.GetTraceParameters{ + TraceID: r.TraceID, + StartTime: r.StartTime, + EndTime: r.EndTime, + }) if errors.Is(err, spanstore.ErrTraceNotFound) { return status.Error(codes.NotFound, spanstore.ErrTraceNotFound.Error()) } @@ -276,7 +280,11 @@ func (s *GRPCHandler) GetArchiveTrace(r *storage_v1.GetTraceRequest, stream stor if reader == nil { return status.Error(codes.Unimplemented, "not implemented") } - trace, err := reader.GetTrace(stream.Context(), r.TraceID) + trace, err := reader.GetTrace(stream.Context(), spanstore.GetTraceParameters{ + TraceID: r.TraceID, + StartTime: r.StartTime, + EndTime: r.EndTime, + }) if errors.Is(err, spanstore.ErrTraceNotFound) { return status.Error(codes.NotFound, spanstore.ErrTraceNotFound.Error()) } diff --git a/plugin/storage/grpc/shared/grpc_handler_test.go b/plugin/storage/grpc/shared/grpc_handler_test.go index 6fe228387c7..726ea30c6b4 100644 --- a/plugin/storage/grpc/shared/grpc_handler_test.go +++ b/plugin/storage/grpc/shared/grpc_handler_test.go @@ -136,7 +136,7 @@ func TestGRPCServerGetTrace(t *testing.T) { for i := range mockTraceSpans { traceSpans = append(traceSpans, &mockTraceSpans[i]) } - r.impl.spanReader.On("GetTrace", mock.Anything, mockTraceID). + r.impl.spanReader.On("GetTrace", mock.Anything, spanstore.GetTraceParameters{TraceID: mockTraceID}). Return(&model.Trace{Spans: traceSpans}, nil) err := r.server.GetTrace(&storage_v1.GetTraceRequest{ @@ -151,7 +151,7 @@ func TestGRPCServerGetTrace_NotFound(t *testing.T) { traceSteam := new(grpcMocks.SpanReaderPlugin_GetTraceServer) traceSteam.On("Context").Return(context.Background()) - r.impl.spanReader.On("GetTrace", mock.Anything, mockTraceID). + r.impl.spanReader.On("GetTrace", mock.Anything, spanstore.GetTraceParameters{TraceID: mockTraceID}). Return(nil, spanstore.ErrTraceNotFound) err := r.server.GetTrace(&storage_v1.GetTraceRequest{ @@ -284,7 +284,7 @@ func TestGRPCServerGetArchiveTrace(t *testing.T) { for i := range mockTraceSpans { traceSpans = append(traceSpans, &mockTraceSpans[i]) } - r.impl.archiveReader.On("GetTrace", mock.Anything, mockTraceID). + r.impl.archiveReader.On("GetTrace", mock.Anything, spanstore.GetTraceParameters{TraceID: mockTraceID}). Return(&model.Trace{Spans: traceSpans}, nil) err := r.server.GetArchiveTrace(&storage_v1.GetTraceRequest{ @@ -299,7 +299,7 @@ func TestGRPCServerGetArchiveTrace_NotFound(t *testing.T) { traceSteam := new(grpcMocks.SpanReaderPlugin_GetTraceServer) traceSteam.On("Context").Return(context.Background()) - r.impl.archiveReader.On("GetTrace", mock.Anything, mockTraceID). + r.impl.archiveReader.On("GetTrace", mock.Anything, spanstore.GetTraceParameters{TraceID: mockTraceID}). Return(nil, spanstore.ErrTraceNotFound) err := r.server.GetArchiveTrace(&storage_v1.GetTraceRequest{ @@ -314,7 +314,7 @@ func TestGRPCServerGetArchiveTrace_Error(t *testing.T) { traceSteam := new(grpcMocks.SpanReaderPlugin_GetTraceServer) traceSteam.On("Context").Return(context.Background()) - r.impl.archiveReader.On("GetTrace", mock.Anything, mockTraceID). + r.impl.archiveReader.On("GetTrace", mock.Anything, spanstore.GetTraceParameters{TraceID: mockTraceID}). Return(nil, errors.New("some error")) err := r.server.GetArchiveTrace(&storage_v1.GetTraceRequest{ @@ -329,7 +329,7 @@ func TestGRPCServerGetArchiveTrace_NoImpl(t *testing.T) { r.server.impl.ArchiveSpanReader = func() spanstore.Reader { return nil } traceSteam := new(grpcMocks.SpanReaderPlugin_GetTraceServer) - r.impl.archiveReader.On("GetTrace", mock.Anything, mockTraceID). + r.impl.archiveReader.On("GetTrace", mock.Anything, spanstore.GetTraceParameters{TraceID: mockTraceID}). Return(nil, errors.New("some error")) err := r.server.GetArchiveTrace(&storage_v1.GetTraceRequest{ @@ -350,7 +350,7 @@ func TestGRPCServerGetArchiveTrace_StreamError(t *testing.T) { for i := range mockTraceSpans { traceSpans = append(traceSpans, &mockTraceSpans[i]) } - r.impl.archiveReader.On("GetTrace", mock.Anything, mockTraceID). + r.impl.archiveReader.On("GetTrace", mock.Anything, spanstore.GetTraceParameters{TraceID: mockTraceID}). Return(&model.Trace{Spans: traceSpans}, nil) err := r.server.GetArchiveTrace(&storage_v1.GetTraceRequest{ diff --git a/plugin/storage/integration/integration.go b/plugin/storage/integration/integration.go index 243cc8c9a1c..1e1d8751a22 100644 --- a/plugin/storage/integration/integration.go +++ b/plugin/storage/integration/integration.go @@ -198,7 +198,7 @@ func (s *StorageIntegration) testArchiveTrace(t *testing.T) { var actual *model.Trace found := s.waitForCondition(t, func(_ *testing.T) bool { var err error - actual, err = s.ArchiveSpanReader.GetTrace(context.Background(), tID) + actual, err = s.ArchiveSpanReader.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: tID}) return err == nil && len(actual.Spans) == 1 }) require.True(t, found) @@ -216,7 +216,7 @@ func (s *StorageIntegration) testGetLargeSpan(t *testing.T) { var actual *model.Trace found := s.waitForCondition(t, func(_ *testing.T) bool { var err error - actual, err = s.SpanReader.GetTrace(context.Background(), expectedTraceID) + actual, err = s.SpanReader.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: expectedTraceID}) return err == nil && len(actual.Spans) >= len(expected.Spans) }) if !assert.True(t, found) { @@ -276,7 +276,7 @@ func (s *StorageIntegration) testGetTrace(t *testing.T) { var actual *model.Trace found := s.waitForCondition(t, func(t *testing.T) bool { var err error - actual, err = s.SpanReader.GetTrace(context.Background(), expectedTraceID) + actual, err = s.SpanReader.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: expectedTraceID}) if err != nil { t.Log(err) } @@ -288,7 +288,7 @@ func (s *StorageIntegration) testGetTrace(t *testing.T) { t.Run("NotFound error", func(t *testing.T) { fakeTraceID := model.TraceID{High: 0, Low: 1} - trace, err := s.SpanReader.GetTrace(context.Background(), fakeTraceID) + trace, err := s.SpanReader.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: fakeTraceID}) assert.Equal(t, spanstore.ErrTraceNotFound, err) assert.Nil(t, trace) }) diff --git a/plugin/storage/integration/kafka_test.go b/plugin/storage/integration/kafka_test.go index ef3528baf65..06664f1a3b5 100644 --- a/plugin/storage/integration/kafka_test.go +++ b/plugin/storage/integration/kafka_test.go @@ -95,8 +95,8 @@ type ingester struct { traceStore *memory.Store } -func (r *ingester) GetTrace(ctx context.Context, traceID model.TraceID) (*model.Trace, error) { - return r.traceStore.GetTrace(ctx, traceID) +func (r *ingester) GetTrace(ctx context.Context, query spanstore.GetTraceParameters) (*model.Trace, error) { + return r.traceStore.GetTrace(ctx, query) } func (*ingester) GetServices(context.Context) ([]string, error) { diff --git a/plugin/storage/memory/memory.go b/plugin/storage/memory/memory.go index 48276aabffb..a20e1a5afc6 100644 --- a/plugin/storage/memory/memory.go +++ b/plugin/storage/memory/memory.go @@ -182,11 +182,11 @@ func (st *Store) WriteSpan(ctx context.Context, span *model.Span) error { } // GetTrace gets a trace -func (st *Store) GetTrace(ctx context.Context, traceID model.TraceID) (*model.Trace, error) { +func (st *Store) GetTrace(ctx context.Context, query spanstore.GetTraceParameters) (*model.Trace, error) { m := st.getTenant(tenancy.GetTenant(ctx)) m.RLock() defer m.RUnlock() - trace, ok := m.traces[traceID] + trace, ok := m.traces[query.TraceID] if !ok { return nil, spanstore.ErrTraceNotFound } diff --git a/plugin/storage/memory/memory_test.go b/plugin/storage/memory/memory_test.go index 7005c273b33..84de6cbd917 100644 --- a/plugin/storage/memory/memory_test.go +++ b/plugin/storage/memory/memory_test.go @@ -187,7 +187,8 @@ func TestStoreWithLimit(t *testing.T) { func TestStoreGetTraceSuccess(t *testing.T) { withPopulatedMemoryStore(func(store *Store) { - trace, err := store.GetTrace(context.Background(), testingSpan.TraceID) + query := spanstore.GetTraceParameters{TraceID: testingSpan.TraceID} + trace, err := store.GetTrace(context.Background(), query) require.NoError(t, err) assert.Len(t, trace.Spans, 1) assert.Equal(t, testingSpan, trace.Spans[0]) @@ -196,7 +197,8 @@ func TestStoreGetTraceSuccess(t *testing.T) { func TestStoreGetAndMutateTrace(t *testing.T) { withPopulatedMemoryStore(func(store *Store) { - trace, err := store.GetTrace(context.Background(), testingSpan.TraceID) + query := spanstore.GetTraceParameters{TraceID: testingSpan.TraceID} + trace, err := store.GetTrace(context.Background(), query) require.NoError(t, err) assert.Len(t, trace.Spans, 1) assert.Equal(t, testingSpan, trace.Spans[0]) @@ -204,7 +206,7 @@ func TestStoreGetAndMutateTrace(t *testing.T) { trace.Spans[0].Warnings = append(trace.Spans[0].Warnings, "the end is near") - trace, err = store.GetTrace(context.Background(), testingSpan.TraceID) + trace, err = store.GetTrace(context.Background(), query) require.NoError(t, err) assert.Len(t, trace.Spans, 1) assert.Equal(t, testingSpan, trace.Spans[0]) @@ -217,14 +219,16 @@ func TestStoreGetTraceError(t *testing.T) { store.getTenant("").traces[testingSpan.TraceID] = &model.Trace{ Spans: []*model.Span{nonSerializableSpan}, } - _, err := store.GetTrace(context.Background(), testingSpan.TraceID) + query := spanstore.GetTraceParameters{TraceID: testingSpan.TraceID} + _, err := store.GetTrace(context.Background(), query) require.Error(t, err) }) } func TestStoreGetTraceFailure(t *testing.T) { withPopulatedMemoryStore(func(store *Store) { - trace, err := store.GetTrace(context.Background(), model.TraceID{}) + query := spanstore.GetTraceParameters{TraceID: model.TraceID{}} + trace, err := store.GetTrace(context.Background(), query) require.EqualError(t, err, spanstore.ErrTraceNotFound.Error()) assert.Nil(t, trace) }) @@ -448,12 +452,14 @@ func TestTenantStore(t *testing.T) { require.NoError(t, store.WriteSpan(ctxWonka, testingSpan2)) // Can we retrieve the spans with correct tenancy - trace1, err := store.GetTrace(ctxAcme, testingSpan.TraceID) + query := spanstore.GetTraceParameters{TraceID: testingSpan.TraceID} + trace1, err := store.GetTrace(ctxAcme, query) require.NoError(t, err) assert.Len(t, trace1.Spans, 1) assert.Equal(t, testingSpan, trace1.Spans[0]) - trace2, err := store.GetTrace(ctxWonka, testingSpan2.TraceID) + query2 := spanstore.GetTraceParameters{TraceID: testingSpan2.TraceID} + trace2, err := store.GetTrace(ctxWonka, query2) require.NoError(t, err) assert.Len(t, trace2.Spans, 1) assert.Equal(t, testingSpan2, trace2.Spans[0]) @@ -476,13 +482,13 @@ func TestTenantStore(t *testing.T) { assert.Equal(t, testingSpan2, traces2[0].Spans[0]) // Do the spans fail with incorrect tenancy? - _, err = store.GetTrace(ctxAcme, testingSpan2.TraceID) + _, err = store.GetTrace(ctxAcme, query2) require.Error(t, err) - _, err = store.GetTrace(ctxWonka, testingSpan.TraceID) + _, err = store.GetTrace(ctxWonka, query) require.Error(t, err) - _, err = store.GetTrace(context.Background(), testingSpan.TraceID) + _, err = store.GetTrace(context.Background(), query) require.Error(t, err) }) } diff --git a/storage/spanstore/interface.go b/storage/spanstore/interface.go index c4c29181502..5b91e392ef8 100644 --- a/storage/spanstore/interface.go +++ b/storage/spanstore/interface.go @@ -25,7 +25,7 @@ type Reader interface { // GetTrace retrieves the trace with a given id. // // If no spans are stored for this trace, it returns ErrTraceNotFound. - GetTrace(ctx context.Context, traceID model.TraceID) (*model.Trace, error) + GetTrace(ctx context.Context, query GetTraceParameters) (*model.Trace, error) // GetServices returns all service names known to the backend from spans // within its retention period. @@ -50,6 +50,13 @@ type Reader interface { FindTraceIDs(ctx context.Context, query *TraceQueryParameters) ([]model.TraceID, error) } +// GetTraceParameters contains parameters of a trace get. +type GetTraceParameters struct { + TraceID model.TraceID + StartTime time.Time + EndTime time.Time +} + // TraceQueryParameters contains parameters of a trace query. type TraceQueryParameters struct { ServiceName string diff --git a/storage/spanstore/mocks/Reader.go b/storage/spanstore/mocks/Reader.go index d48658937c3..29d745996f3 100644 --- a/storage/spanstore/mocks/Reader.go +++ b/storage/spanstore/mocks/Reader.go @@ -141,9 +141,9 @@ func (_m *Reader) GetServices(ctx context.Context) ([]string, error) { return r0, r1 } -// GetTrace provides a mock function with given fields: ctx, traceID -func (_m *Reader) GetTrace(ctx context.Context, traceID model.TraceID) (*model.Trace, error) { - ret := _m.Called(ctx, traceID) +// GetTrace provides a mock function with given fields: ctx, query +func (_m *Reader) GetTrace(ctx context.Context, query spanstore.GetTraceParameters) (*model.Trace, error) { + ret := _m.Called(ctx, query) if len(ret) == 0 { panic("no return value specified for GetTrace") @@ -151,19 +151,19 @@ func (_m *Reader) GetTrace(ctx context.Context, traceID model.TraceID) (*model.T var r0 *model.Trace var r1 error - if rf, ok := ret.Get(0).(func(context.Context, model.TraceID) (*model.Trace, error)); ok { - return rf(ctx, traceID) + if rf, ok := ret.Get(0).(func(context.Context, spanstore.GetTraceParameters) (*model.Trace, error)); ok { + return rf(ctx, query) } - if rf, ok := ret.Get(0).(func(context.Context, model.TraceID) *model.Trace); ok { - r0 = rf(ctx, traceID) + if rf, ok := ret.Get(0).(func(context.Context, spanstore.GetTraceParameters) *model.Trace); ok { + r0 = rf(ctx, query) } else { if ret.Get(0) != nil { r0 = ret.Get(0).(*model.Trace) } } - if rf, ok := ret.Get(1).(func(context.Context, model.TraceID) error); ok { - r1 = rf(ctx, traceID) + if rf, ok := ret.Get(1).(func(context.Context, spanstore.GetTraceParameters) error); ok { + r1 = rf(ctx, query) } else { r1 = ret.Error(1) } diff --git a/storage/spanstore/spanstoremetrics/decorator_test.go b/storage/spanstore/spanstoremetrics/decorator_test.go index 8f9c40863b5..bf366c9b06c 100644 --- a/storage/spanstore/spanstoremetrics/decorator_test.go +++ b/storage/spanstore/spanstoremetrics/decorator_test.go @@ -29,8 +29,8 @@ func TestSuccessfulUnderlyingCalls(t *testing.T) { mockReader.On("GetOperations", context.Background(), operationQuery). Return([]spanstore.Operation{}, nil) mrs.GetOperations(context.Background(), operationQuery) - mockReader.On("GetTrace", context.Background(), model.TraceID{}).Return(&model.Trace{}, nil) - mrs.GetTrace(context.Background(), model.TraceID{}) + mockReader.On("GetTrace", context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}).Return(&model.Trace{}, nil) + mrs.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}) mockReader.On("FindTraces", context.Background(), &spanstore.TraceQueryParameters{}). Return([]*model.Trace{}, nil) mrs.FindTraces(context.Background(), &spanstore.TraceQueryParameters{}) @@ -97,9 +97,9 @@ func TestFailingUnderlyingCalls(t *testing.T) { mockReader.On("GetOperations", context.Background(), operationQuery). Return(nil, errors.New("Failure")) mrs.GetOperations(context.Background(), operationQuery) - mockReader.On("GetTrace", context.Background(), model.TraceID{}). + mockReader.On("GetTrace", context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}). Return(nil, errors.New("Failure")) - mrs.GetTrace(context.Background(), model.TraceID{}) + mrs.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}) mockReader.On("FindTraces", context.Background(), &spanstore.TraceQueryParameters{}). Return(nil, errors.New("Failure")) mrs.FindTraces(context.Background(), &spanstore.TraceQueryParameters{}) diff --git a/storage/spanstore/spanstoremetrics/read_metrics.go b/storage/spanstore/spanstoremetrics/read_metrics.go index 0f4ca13d0c0..8a42e7d2081 100644 --- a/storage/spanstore/spanstoremetrics/read_metrics.go +++ b/storage/spanstore/spanstoremetrics/read_metrics.go @@ -78,9 +78,9 @@ func (m *ReadMetricsDecorator) FindTraceIDs(ctx context.Context, traceQuery *spa } // GetTrace implements spanstore.Reader#GetTrace -func (m *ReadMetricsDecorator) GetTrace(ctx context.Context, traceID model.TraceID) (*model.Trace, error) { +func (m *ReadMetricsDecorator) GetTrace(ctx context.Context, traceGet spanstore.GetTraceParameters) (*model.Trace, error) { start := time.Now() - retMe, err := m.spanReader.GetTrace(ctx, traceID) + retMe, err := m.spanReader.GetTrace(ctx, traceGet) m.getTraceMetrics.emit(err, time.Since(start), 1) return retMe, err } diff --git a/storage_v2/factoryadapter/writer_test.go b/storage_v2/factoryadapter/writer_test.go index df5966cd33f..68a09af6e4c 100644 --- a/storage_v2/factoryadapter/writer_test.go +++ b/storage_v2/factoryadapter/writer_test.go @@ -16,6 +16,7 @@ import ( "github.com/jaegertracing/jaeger/model" "github.com/jaegertracing/jaeger/plugin/storage/memory" + "github.com/jaegertracing/jaeger/storage/spanstore" spanstoreMocks "github.com/jaegertracing/jaeger/storage/spanstore/mocks" ) @@ -32,7 +33,8 @@ func TestWriteTraces(t *testing.T) { tdID := td.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0).TraceID() traceID, err := model.TraceIDFromBytes(tdID[:]) require.NoError(t, err) - trace, err := memstore.GetTrace(context.Background(), traceID) + query := spanstore.GetTraceParameters{TraceID: traceID} + trace, err := memstore.GetTrace(context.Background(), query) require.NoError(t, err) require.NotNil(t, trace) assert.Len(t, trace.Spans, 1) From 097f36a7af1111a992059bbd2d1320ed609668a5 Mon Sep 17 00:00:00 2001 From: rim99 Date: Sun, 24 Nov 2024 12:03:29 +0800 Subject: [PATCH 02/12] add unit test in grpc_handler_test.go Signed-off-by: rim99 --- cmd/query/app/grpc_handler_test.go | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/cmd/query/app/grpc_handler_test.go b/cmd/query/app/grpc_handler_test.go index 7ca990bc2f8..197047d0c7a 100644 --- a/cmd/query/app/grpc_handler_test.go +++ b/cmd/query/app/grpc_handler_test.go @@ -300,6 +300,30 @@ func TestArchiveTraceSuccessGRPC(t *testing.T) { }) } +func TestArchiveTraceWithTimeWindowSuccessGRPC(t *testing.T) { + startTime := time.Unix(1, 2).UTC() + endTime := time.Unix(3, 4).UTC() + withServerAndClient(t, func(server *grpcServer, client *grpcClient) { + expectedGetTraceParameters := spanstore.GetTraceParameters{ + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, + } + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedGetTraceParameters). + Return(mockTrace, nil).Once() + server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). + Return(nil).Times(2) + + _, err := client.ArchiveTrace(context.Background(), &api_v2.ArchiveTraceRequest{ + TraceID: mockTraceID, + StartTime: &startTime, + EndTime: &endTime, + }) + + require.NoError(t, err) + }) +} + func TestArchiveTraceNotFoundGRPC(t *testing.T) { withServerAndClient(t, func(server *grpcServer, client *grpcClient) { server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). From cfecbea0542c2b608070af4c10702b241cd97d53 Mon Sep 17 00:00:00 2001 From: Zhang Xin Date: Mon, 25 Nov 2024 07:31:45 +0800 Subject: [PATCH 03/12] Update storage/spanstore/interface.go Co-authored-by: Yuri Shkuro Signed-off-by: Zhang Xin --- storage/spanstore/interface.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/storage/spanstore/interface.go b/storage/spanstore/interface.go index 5b91e392ef8..3c4dd0dd4e7 100644 --- a/storage/spanstore/interface.go +++ b/storage/spanstore/interface.go @@ -53,8 +53,8 @@ type Reader interface { // GetTraceParameters contains parameters of a trace get. type GetTraceParameters struct { TraceID model.TraceID - StartTime time.Time - EndTime time.Time + StartTime time.Time // optional + EndTime time.Time // optional } // TraceQueryParameters contains parameters of a trace query. From 1ab1419d06dffc37cdd74895a72a254332f2a7d7 Mon Sep 17 00:00:00 2001 From: Zhang Xin Date: Mon, 25 Nov 2024 07:31:59 +0800 Subject: [PATCH 04/12] Update storage/spanstore/metrics/decorator_test.go Co-authored-by: Yuri Shkuro Signed-off-by: Zhang Xin --- storage/spanstore/spanstoremetrics/decorator_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/storage/spanstore/spanstoremetrics/decorator_test.go b/storage/spanstore/spanstoremetrics/decorator_test.go index bf366c9b06c..3f6f26a52f7 100644 --- a/storage/spanstore/spanstoremetrics/decorator_test.go +++ b/storage/spanstore/spanstoremetrics/decorator_test.go @@ -30,7 +30,7 @@ func TestSuccessfulUnderlyingCalls(t *testing.T) { Return([]spanstore.Operation{}, nil) mrs.GetOperations(context.Background(), operationQuery) mockReader.On("GetTrace", context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}).Return(&model.Trace{}, nil) - mrs.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}) + mrs.GetTrace(context.Background(), spanstore.GetTraceParameters{}) mockReader.On("FindTraces", context.Background(), &spanstore.TraceQueryParameters{}). Return([]*model.Trace{}, nil) mrs.FindTraces(context.Background(), &spanstore.TraceQueryParameters{}) From 512f6cffbf5da5d8c3ad6182c7ac47337356bca9 Mon Sep 17 00:00:00 2001 From: Zhang Xin Date: Mon, 25 Nov 2024 07:32:07 +0800 Subject: [PATCH 05/12] Update storage/spanstore/metrics/decorator.go Co-authored-by: Yuri Shkuro Signed-off-by: Zhang Xin --- storage/spanstore/spanstoremetrics/read_metrics.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/storage/spanstore/spanstoremetrics/read_metrics.go b/storage/spanstore/spanstoremetrics/read_metrics.go index 8a42e7d2081..b8d46197744 100644 --- a/storage/spanstore/spanstoremetrics/read_metrics.go +++ b/storage/spanstore/spanstoremetrics/read_metrics.go @@ -78,7 +78,7 @@ func (m *ReadMetricsDecorator) FindTraceIDs(ctx context.Context, traceQuery *spa } // GetTrace implements spanstore.Reader#GetTrace -func (m *ReadMetricsDecorator) GetTrace(ctx context.Context, traceGet spanstore.GetTraceParameters) (*model.Trace, error) { +func (m *ReadMetricsDecorator) GetTrace(ctx context.Context, query spanstore.GetTraceParameters) (*model.Trace, error) { start := time.Now() retMe, err := m.spanReader.GetTrace(ctx, traceGet) m.getTraceMetrics.emit(err, time.Since(start), 1) From 811df658b06d914988532af4c5f2719719de1391 Mon Sep 17 00:00:00 2001 From: rim99 Date: Sun, 24 Nov 2024 12:03:29 +0800 Subject: [PATCH 06/12] add unit test in grpc_handler_test.go Signed-off-by: rim99 --- cmd/query/app/grpc_handler_test.go | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/cmd/query/app/grpc_handler_test.go b/cmd/query/app/grpc_handler_test.go index 197047d0c7a..52a98987316 100644 --- a/cmd/query/app/grpc_handler_test.go +++ b/cmd/query/app/grpc_handler_test.go @@ -196,6 +196,25 @@ func withServerAndClient(t *testing.T, actualTest func(server *grpcServer, clien } func TestGetTraceSuccessGRPC(t *testing.T) { + withServerAndClient(t, func(server *grpcServer, client *grpcClient) { + expectedTraceGetParameters := spanstore.GetTraceParameters{ + TraceID: mockTraceID, + } + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedTraceGetParameters). + Return(mockTrace, nil).Once() + + res, err := client.GetTrace(context.Background(), &api_v2.GetTraceRequest{ + TraceID: mockTraceID, + }) + + spanResChunk, _ := res.Recv() + + require.NoError(t, err) + assert.Equal(t, spanResChunk.Spans[0].TraceID, mockTraceID) + }) +} + +func TestGetTraceWithTimeWindowSuccessGRPC(t *testing.T) { withServerAndClient(t, func(server *grpcServer, client *grpcClient) { startTime := time.Date(2020, time.January, 1, 13, 0, 0, 0, time.UTC) endTime := time.Date(2020, time.January, 1, 14, 0, 0, 0, time.UTC) From 64588a506a85e14deb56c30cc1fe02cf693c19b5 Mon Sep 17 00:00:00 2001 From: rim99 Date: Mon, 25 Nov 2024 07:43:27 +0800 Subject: [PATCH 07/12] update unit tests Signed-off-by: rim99 --- cmd/query/app/grpc_handler_test.go | 129 ++++++++---------- cmd/query/app/http_handler_test.go | 2 + .../cassandra/spanstore/reader_test.go | 4 +- plugin/storage/es/spanstore/reader_test.go | 4 +- plugin/storage/memory/memory_test.go | 2 +- .../spanstoremetrics/decorator_test.go | 6 +- .../spanstoremetrics/read_metrics.go | 2 +- 7 files changed, 71 insertions(+), 78 deletions(-) diff --git a/cmd/query/app/grpc_handler_test.go b/cmd/query/app/grpc_handler_test.go index 52a98987316..abba3a3628c 100644 --- a/cmd/query/app/grpc_handler_test.go +++ b/cmd/query/app/grpc_handler_test.go @@ -196,47 +196,41 @@ func withServerAndClient(t *testing.T, actualTest func(server *grpcServer, clien } func TestGetTraceSuccessGRPC(t *testing.T) { - withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - expectedTraceGetParameters := spanstore.GetTraceParameters{ - TraceID: mockTraceID, - } - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedTraceGetParameters). - Return(mockTrace, nil).Once() - - res, err := client.GetTrace(context.Background(), &api_v2.GetTraceRequest{ - TraceID: mockTraceID, - }) + inputs := []struct { + expectedQuery spanstore.GetTraceParameters + request api_v2.GetTraceRequest + }{ + { + spanstore.GetTraceParameters{TraceID: mockTraceID}, + api_v2.GetTraceRequest{TraceID: mockTraceID}, + }, + { + spanstore.GetTraceParameters{ + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, + }, + api_v2.GetTraceRequest{ + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, + }, + }, + } - spanResChunk, _ := res.Recv() + for _, input := range inputs { + withServerAndClient(t, func(server *grpcServer, client *grpcClient) { + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), input.expectedQuery). + Return(mockTrace, nil).Once() - require.NoError(t, err) - assert.Equal(t, spanResChunk.Spans[0].TraceID, mockTraceID) - }) -} + res, err := client.GetTrace(context.Background(), &input.request) -func TestGetTraceWithTimeWindowSuccessGRPC(t *testing.T) { - withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - startTime := time.Date(2020, time.January, 1, 13, 0, 0, 0, time.UTC) - endTime := time.Date(2020, time.January, 1, 14, 0, 0, 0, time.UTC) - expectedTraceGetParameters := spanstore.GetTraceParameters{ - TraceID: mockTraceID, - StartTime: startTime, - EndTime: endTime, - } - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedTraceGetParameters). - Return(mockTrace, nil).Once() + spanResChunk, _ := res.Recv() - res, err := client.GetTrace(context.Background(), &api_v2.GetTraceRequest{ - TraceID: mockTraceID, - StartTime: startTime, - EndTime: endTime, + require.NoError(t, err) + assert.Equal(t, spanResChunk.Spans[0].TraceID, mockTraceID) }) - - spanResChunk, _ := res.Recv() - - require.NoError(t, err) - assert.Equal(t, spanResChunk.Spans[0].TraceID, mockTraceID) - }) + } } func assertGRPCError(t *testing.T, err error, code codes.Code, msg string) { @@ -305,42 +299,39 @@ func TestGetTraceNilRequestOnHandlerGRPC(t *testing.T) { } func TestArchiveTraceSuccessGRPC(t *testing.T) { - withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). - Return(mockTrace, nil).Once() - server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). - Return(nil).Times(2) - - _, err := client.ArchiveTrace(context.Background(), &api_v2.ArchiveTraceRequest{ - TraceID: mockTraceID, - }) - - require.NoError(t, err) - }) -} + inputs := []struct { + expectedQuery spanstore.GetTraceParameters + request api_v2.ArchiveTraceRequest + }{ + { + spanstore.GetTraceParameters{TraceID: mockTraceID}, + api_v2.ArchiveTraceRequest{TraceID: mockTraceID}, + }, + { + spanstore.GetTraceParameters{ + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, + }, + api_v2.ArchiveTraceRequest{ + TraceID: mockTraceID, + StartTime: &startTime, + EndTime: &endTime, + }, + }, + } + for _, input := range inputs { + withServerAndClient(t, func(server *grpcServer, client *grpcClient) { + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), input.expectedQuery). + Return(mockTrace, nil).Once() + server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). + Return(nil).Times(2) -func TestArchiveTraceWithTimeWindowSuccessGRPC(t *testing.T) { - startTime := time.Unix(1, 2).UTC() - endTime := time.Unix(3, 4).UTC() - withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - expectedGetTraceParameters := spanstore.GetTraceParameters{ - TraceID: mockTraceID, - StartTime: startTime, - EndTime: endTime, - } - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedGetTraceParameters). - Return(mockTrace, nil).Once() - server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). - Return(nil).Times(2) + _, err := client.ArchiveTrace(context.Background(), &input.request) - _, err := client.ArchiveTrace(context.Background(), &api_v2.ArchiveTraceRequest{ - TraceID: mockTraceID, - StartTime: &startTime, - EndTime: &endTime, + require.NoError(t, err) }) - - require.NoError(t, err) - }) + } } func TestArchiveTraceNotFoundGRPC(t *testing.T) { diff --git a/cmd/query/app/http_handler_test.go b/cmd/query/app/http_handler_test.go index 32c15d6aa91..3fc33f6bebf 100644 --- a/cmd/query/app/http_handler_test.go +++ b/cmd/query/app/http_handler_test.go @@ -65,6 +65,8 @@ var ( } mockTraceID = model.NewTraceID(0, 123456) + startTime = time.Date(2020, time.January, 1, 13, 0, 0, 0, time.UTC) + endTime = time.Date(2020, time.January, 1, 14, 0, 0, 0, time.UTC) mockTrace = &model.Trace{ Spans: []*model.Span{ { diff --git a/plugin/storage/cassandra/spanstore/reader_test.go b/plugin/storage/cassandra/spanstore/reader_test.go index a55ced3e7dd..afbcf48f0fb 100644 --- a/plugin/storage/cassandra/spanstore/reader_test.go +++ b/plugin/storage/cassandra/spanstore/reader_test.go @@ -171,7 +171,7 @@ func TestSpanReaderGetTrace(t *testing.T) { r.session.On("Query", mock.AnythingOfType("string"), matchEverything()).Return(query) - trace, err := r.reader.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}) + trace, err := r.reader.GetTrace(context.Background(), spanstore.GetTraceParameters{}) if testCase.expectedErr == "" { require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.NoError(t, err) @@ -197,7 +197,7 @@ func TestSpanReaderGetTrace_TraceNotFound(t *testing.T) { r.session.On("Query", mock.AnythingOfType("string"), matchEverything()).Return(query) - trace, err := r.reader.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}) + trace, err := r.reader.GetTrace(context.Background(), spanstore.GetTraceParameters{}) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") assert.Nil(t, trace) require.EqualError(t, err, "trace not found") diff --git a/plugin/storage/es/spanstore/reader_test.go b/plugin/storage/es/spanstore/reader_test.go index a91baafc2bd..3cc8292e640 100644 --- a/plugin/storage/es/spanstore/reader_test.go +++ b/plugin/storage/es/spanstore/reader_test.go @@ -1277,7 +1277,7 @@ func TestSpanReader_ArchiveTraces(t *testing.T) { Return(&elastic.MultiSearchResult{ Responses: []*elastic.SearchResult{}, }, nil) - query := spanstore.GetTraceParameters{TraceID: model.TraceID{}} + query := spanstore.GetTraceParameters{} trace, err := r.reader.GetTrace(context.Background(), query) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.Nil(t, trace) @@ -1294,7 +1294,7 @@ func TestSpanReader_ArchiveTraces_ReadAlias(t *testing.T) { Responses: []*elastic.SearchResult{}, }, nil) - query := spanstore.GetTraceParameters{TraceID: model.TraceID{}} + query := spanstore.GetTraceParameters{} trace, err := r.reader.GetTrace(context.Background(), query) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.Nil(t, trace) diff --git a/plugin/storage/memory/memory_test.go b/plugin/storage/memory/memory_test.go index 84de6cbd917..af130edd566 100644 --- a/plugin/storage/memory/memory_test.go +++ b/plugin/storage/memory/memory_test.go @@ -227,7 +227,7 @@ func TestStoreGetTraceError(t *testing.T) { func TestStoreGetTraceFailure(t *testing.T) { withPopulatedMemoryStore(func(store *Store) { - query := spanstore.GetTraceParameters{TraceID: model.TraceID{}} + query := spanstore.GetTraceParameters{} trace, err := store.GetTrace(context.Background(), query) require.EqualError(t, err, spanstore.ErrTraceNotFound.Error()) assert.Nil(t, trace) diff --git a/storage/spanstore/spanstoremetrics/decorator_test.go b/storage/spanstore/spanstoremetrics/decorator_test.go index 3f6f26a52f7..8581cbf9875 100644 --- a/storage/spanstore/spanstoremetrics/decorator_test.go +++ b/storage/spanstore/spanstoremetrics/decorator_test.go @@ -29,7 +29,7 @@ func TestSuccessfulUnderlyingCalls(t *testing.T) { mockReader.On("GetOperations", context.Background(), operationQuery). Return([]spanstore.Operation{}, nil) mrs.GetOperations(context.Background(), operationQuery) - mockReader.On("GetTrace", context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}).Return(&model.Trace{}, nil) + mockReader.On("GetTrace", context.Background(), spanstore.GetTraceParameters{}).Return(&model.Trace{}, nil) mrs.GetTrace(context.Background(), spanstore.GetTraceParameters{}) mockReader.On("FindTraces", context.Background(), &spanstore.TraceQueryParameters{}). Return([]*model.Trace{}, nil) @@ -97,9 +97,9 @@ func TestFailingUnderlyingCalls(t *testing.T) { mockReader.On("GetOperations", context.Background(), operationQuery). Return(nil, errors.New("Failure")) mrs.GetOperations(context.Background(), operationQuery) - mockReader.On("GetTrace", context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}). + mockReader.On("GetTrace", context.Background(), spanstore.GetTraceParameters{}). Return(nil, errors.New("Failure")) - mrs.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}) + mrs.GetTrace(context.Background(), spanstore.GetTraceParameters{}) mockReader.On("FindTraces", context.Background(), &spanstore.TraceQueryParameters{}). Return(nil, errors.New("Failure")) mrs.FindTraces(context.Background(), &spanstore.TraceQueryParameters{}) diff --git a/storage/spanstore/spanstoremetrics/read_metrics.go b/storage/spanstore/spanstoremetrics/read_metrics.go index b8d46197744..b87453cc053 100644 --- a/storage/spanstore/spanstoremetrics/read_metrics.go +++ b/storage/spanstore/spanstoremetrics/read_metrics.go @@ -80,7 +80,7 @@ func (m *ReadMetricsDecorator) FindTraceIDs(ctx context.Context, traceQuery *spa // GetTrace implements spanstore.Reader#GetTrace func (m *ReadMetricsDecorator) GetTrace(ctx context.Context, query spanstore.GetTraceParameters) (*model.Trace, error) { start := time.Now() - retMe, err := m.spanReader.GetTrace(ctx, traceGet) + retMe, err := m.spanReader.GetTrace(ctx, query) m.getTraceMetrics.emit(err, time.Since(start), 1) return retMe, err } From 743a8d339da6467fffece26631b5e117d62b7035 Mon Sep 17 00:00:00 2001 From: rim99 Date: Mon, 25 Nov 2024 07:46:20 +0800 Subject: [PATCH 08/12] add more TODOs Signed-off-by: rim99 --- cmd/query/app/http_handler.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cmd/query/app/http_handler.go b/cmd/query/app/http_handler.go index 97258e3161b..899ca208411 100644 --- a/cmd/query/app/http_handler.go +++ b/cmd/query/app/http_handler.go @@ -266,6 +266,7 @@ func (aH *APIHandler) tracesByIDs(ctx context.Context, traceIDs []model.TraceID) var traceErrors []structuredError retMe := make([]*model.Trace, 0, len(traceIDs)) for _, traceID := range traceIDs { + // TODO: add start time & end time query := spanstore.GetTraceParameters{ TraceID: traceID, } @@ -432,6 +433,7 @@ func (aH *APIHandler) getTrace(w http.ResponseWriter, r *http.Request) { if !ok { return } + // TODO: add start time & end time query := spanstore.GetTraceParameters{ TraceID: traceID, } From b0556391f7bc97fae98dda702c714afe225fce23 Mon Sep 17 00:00:00 2001 From: rim99 Date: Mon, 25 Nov 2024 07:59:47 +0800 Subject: [PATCH 09/12] update type of time fields of api_v2 ArchiveTraceRequest Signed-off-by: rim99 --- cmd/query/app/grpc_handler.go | 10 +++------- cmd/query/app/grpc_handler_test.go | 4 ++-- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/cmd/query/app/grpc_handler.go b/cmd/query/app/grpc_handler.go index c22780b6c5c..1e007d10e5b 100644 --- a/cmd/query/app/grpc_handler.go +++ b/cmd/query/app/grpc_handler.go @@ -115,13 +115,9 @@ func (g *GRPCHandler) ArchiveTrace(ctx context.Context, r *api_v2.ArchiveTraceRe return nil, errUninitializedTraceID } query := spanstore.GetTraceParameters{ - TraceID: r.TraceID, - } - if r.StartTime != nil { - query.StartTime = *r.StartTime - } - if r.EndTime != nil { - query.EndTime = *r.EndTime + TraceID: r.TraceID, + StartTime: r.StartTime, + EndTime: r.EndTime, } err := g.queryService.ArchiveTrace(ctx, query) if errors.Is(err, spanstore.ErrTraceNotFound) { diff --git a/cmd/query/app/grpc_handler_test.go b/cmd/query/app/grpc_handler_test.go index abba3a3628c..00a4088b82e 100644 --- a/cmd/query/app/grpc_handler_test.go +++ b/cmd/query/app/grpc_handler_test.go @@ -315,8 +315,8 @@ func TestArchiveTraceSuccessGRPC(t *testing.T) { }, api_v2.ArchiveTraceRequest{ TraceID: mockTraceID, - StartTime: &startTime, - EndTime: &endTime, + StartTime: startTime, + EndTime: endTime, }, }, } From ebb2b61f05ef265a624055b6fbc2f3f01e52825d Mon Sep 17 00:00:00 2001 From: rim99 Date: Tue, 26 Nov 2024 07:21:56 +0800 Subject: [PATCH 10/12] refactor test Signed-off-by: rim99 --- cmd/query/app/apiv3/grpc_handler_test.go | 103 +++++++++-------------- 1 file changed, 42 insertions(+), 61 deletions(-) diff --git a/cmd/query/app/apiv3/grpc_handler_test.go b/cmd/query/app/apiv3/grpc_handler_test.go index 84f342eac68..27351962b8f 100644 --- a/cmd/query/app/apiv3/grpc_handler_test.go +++ b/cmd/query/app/apiv3/grpc_handler_test.go @@ -78,73 +78,54 @@ func newTestServerClient(t *testing.T) *testServerClient { } func TestGetTrace(t *testing.T) { - tsc := newTestServerClient(t) traceIdLiteral := "156" traceId, _ := model.TraceIDFromString(traceIdLiteral) - - expectedTraceGetParameters := spanstore.GetTraceParameters{ - TraceID: traceId, - StartTime: time.Time{}, - EndTime: time.Time{}, - } - tsc.reader.On("GetTrace", matchContext, expectedTraceGetParameters).Return( - &model.Trace{ - Spans: []*model.Span{ - { - OperationName: "foobar", - }, + inputs := []struct { + expectedQuery spanstore.GetTraceParameters + request api_v3.GetTraceRequest + }{ + { + spanstore.GetTraceParameters{ + TraceID: traceId, + StartTime: time.Time{}, + EndTime: time.Time{}, }, - }, nil).Once() - - getTraceStream, err := tsc.client.GetTrace(context.Background(), - &api_v3.GetTraceRequest{ - TraceId: traceIdLiteral, + api_v3.GetTraceRequest{TraceId: traceIdLiteral}, }, - ) - require.NoError(t, err) - recv, err := getTraceStream.Recv() - require.NoError(t, err) - td := recv.ToTraces() - require.EqualValues(t, 1, td.SpanCount()) - assert.Equal(t, "foobar", - td.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0).Name()) -} - -func TestGetTraceWithTimeWindow(t *testing.T) { - tsc := newTestServerClient(t) - traceIdLiteral := "156" - traceId, _ := model.TraceIDFromString(traceIdLiteral) - - start_ts := time.Unix(1, 2).UTC() - end_ts := time.Unix(3, 4).UTC() - expectedTraceGetParameters := spanstore.GetTraceParameters{ - TraceID: traceId, - StartTime: start_ts, - EndTime: end_ts, - } - tsc.reader.On("GetTrace", matchContext, expectedTraceGetParameters).Return( - &model.Trace{ - Spans: []*model.Span{ - { - OperationName: "foobar", - }, + { + spanstore.GetTraceParameters{ + TraceID: traceId, + StartTime: time.Unix(1, 2).UTC(), + EndTime: time.Unix(3, 4).UTC(), + }, + api_v3.GetTraceRequest{ + TraceId: traceIdLiteral, + StartTime: time.Unix(1, 2).UTC(), + EndTime: time.Unix(3, 4).UTC(), }, - }, nil).Once() - - getTraceStream, err := tsc.client.GetTrace(context.Background(), - &api_v3.GetTraceRequest{ - TraceId: traceIdLiteral, - StartTime: start_ts, - EndTime: end_ts, }, - ) - require.NoError(t, err) - recv, err := getTraceStream.Recv() - require.NoError(t, err) - td := recv.ToTraces() - require.EqualValues(t, 1, td.SpanCount()) - assert.Equal(t, "foobar", - td.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0).Name()) + } + + for _, input := range inputs { + tsc := newTestServerClient(t) + tsc.reader.On("GetTrace", matchContext, input.expectedQuery).Return( + &model.Trace{ + Spans: []*model.Span{ + { + OperationName: "foobar", + }, + }, + }, nil).Once() + + getTraceStream, err := tsc.client.GetTrace(context.Background(), &input.request) + require.NoError(t, err) + recv, err := getTraceStream.Recv() + require.NoError(t, err) + td := recv.ToTraces() + require.EqualValues(t, 1, td.SpanCount()) + assert.Equal(t, "foobar", + td.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0).Name()) + } } func TestGetTraceStorageError(t *testing.T) { From 491c8264b490cbbd200046fb61dec14544c672a3 Mon Sep 17 00:00:00 2001 From: rim99 Date: Tue, 26 Nov 2024 20:11:08 +0800 Subject: [PATCH 11/12] refactor unit tests Signed-off-by: rim99 --- cmd/query/app/apiv3/grpc_handler_test.go | 50 +++++++++++++----------- 1 file changed, 27 insertions(+), 23 deletions(-) diff --git a/cmd/query/app/apiv3/grpc_handler_test.go b/cmd/query/app/apiv3/grpc_handler_test.go index 27351962b8f..d6ff4f212ea 100644 --- a/cmd/query/app/apiv3/grpc_handler_test.go +++ b/cmd/query/app/apiv3/grpc_handler_test.go @@ -78,53 +78,57 @@ func newTestServerClient(t *testing.T) *testServerClient { } func TestGetTrace(t *testing.T) { - traceIdLiteral := "156" - traceId, _ := model.TraceIDFromString(traceIdLiteral) - inputs := []struct { + traceId, _ := model.TraceIDFromString("156") + testCases := []struct { + name string expectedQuery spanstore.GetTraceParameters request api_v3.GetTraceRequest }{ { + "TestGetTrace", spanstore.GetTraceParameters{ TraceID: traceId, StartTime: time.Time{}, EndTime: time.Time{}, }, - api_v3.GetTraceRequest{TraceId: traceIdLiteral}, + api_v3.GetTraceRequest{TraceId: "156"}, }, { + "TestGetTraceWithTimeWindow", spanstore.GetTraceParameters{ TraceID: traceId, StartTime: time.Unix(1, 2).UTC(), EndTime: time.Unix(3, 4).UTC(), }, api_v3.GetTraceRequest{ - TraceId: traceIdLiteral, + TraceId: "156", StartTime: time.Unix(1, 2).UTC(), EndTime: time.Unix(3, 4).UTC(), }, }, } - for _, input := range inputs { - tsc := newTestServerClient(t) - tsc.reader.On("GetTrace", matchContext, input.expectedQuery).Return( - &model.Trace{ - Spans: []*model.Span{ - { - OperationName: "foobar", + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + tsc := newTestServerClient(t) + tsc.reader.On("GetTrace", matchContext, tc.expectedQuery).Return( + &model.Trace{ + Spans: []*model.Span{ + { + OperationName: "foobar", + }, }, - }, - }, nil).Once() - - getTraceStream, err := tsc.client.GetTrace(context.Background(), &input.request) - require.NoError(t, err) - recv, err := getTraceStream.Recv() - require.NoError(t, err) - td := recv.ToTraces() - require.EqualValues(t, 1, td.SpanCount()) - assert.Equal(t, "foobar", - td.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0).Name()) + }, nil).Once() + + getTraceStream, err := tsc.client.GetTrace(context.Background(), &tc.request) + require.NoError(t, err) + recv, err := getTraceStream.Recv() + require.NoError(t, err) + td := recv.ToTraces() + require.EqualValues(t, 1, td.SpanCount()) + assert.Equal(t, "foobar", + td.ResourceSpans().At(0).ScopeSpans().At(0).Spans().At(0).Name()) + }) } } From b9c4cebb32323d8421d89f7d2b5b0ed92fa4a86b Mon Sep 17 00:00:00 2001 From: Zhang Xin Date: Sun, 15 Dec 2024 10:10:11 +0800 Subject: [PATCH 12/12] fix tests Signed-off-by: Zhang Xin --- cmd/query/app/querysvc/query_service_test.go | 3 ++- storage_v2/factoryadapter/reader.go | 5 ++++- storage_v2/factoryadapter/reader_test.go | 3 ++- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/cmd/query/app/querysvc/query_service_test.go b/cmd/query/app/querysvc/query_service_test.go index d5c6f3547c4..10a3b21e5af 100644 --- a/cmd/query/app/querysvc/query_service_test.go +++ b/cmd/query/app/querysvc/query_service_test.go @@ -140,7 +140,8 @@ func TestGetTrace_V1ReaderNotFound(t *testing.T) { qs := QueryService{ traceReader: fr, } - _, err := qs.GetTrace(context.Background(), mockTraceID) + query := spanstore.GetTraceParameters{TraceID: mockTraceID} + _, err := qs.GetTrace(context.Background(), query) require.Error(t, err) } diff --git a/storage_v2/factoryadapter/reader.go b/storage_v2/factoryadapter/reader.go index c5239abd8ab..0074127a3ce 100644 --- a/storage_v2/factoryadapter/reader.go +++ b/storage_v2/factoryadapter/reader.go @@ -42,7 +42,10 @@ func NewTraceReader(spanReader spanstore.Reader) *TraceReader { func (tr *TraceReader) GetTrace(ctx context.Context, traceID pcommon.TraceID) iter.Seq2[ptrace.Traces, error] { return func(yield func(ptrace.Traces, error) bool) { - t, err := tr.spanReader.GetTrace(ctx, model.TraceIDFromOTEL(traceID)) + query := spanstore.GetTraceParameters{ + TraceID: model.TraceIDFromOTEL(traceID), + } + t, err := tr.spanReader.GetTrace(ctx, query) if err != nil { // if not found, return an empty iterator, otherwire yield the error. if !errors.Is(err, spanstore.ErrTraceNotFound) { diff --git a/storage_v2/factoryadapter/reader_test.go b/storage_v2/factoryadapter/reader_test.go index bd7c9c4a91a..dbbe81fb98c 100644 --- a/storage_v2/factoryadapter/reader_test.go +++ b/storage_v2/factoryadapter/reader_test.go @@ -58,7 +58,8 @@ func TestTraceReader_GetTraceDelegatesSuccessResponse(t *testing.T) { }, }, } - sr.On("GetTrace", mock.Anything, model.NewTraceID(2, 3)).Return(modelTrace, nil) + expectedQuery := spanstore.GetTraceParameters{TraceID: model.NewTraceID(2, 3)} + sr.On("GetTrace", mock.Anything, expectedQuery).Return(modelTrace, nil) traceReader := &TraceReader{ spanReader: sr, }