add optional time window when fetching trace by id

Signed-off-by: rim99 <[email protected]>

Makefile.Protobuf.mk

@@ -14,9 +14,10 @@
 # instead of the go_package's declared by the imported protof files.
 #
 
-DOCKER_PROTOBUF_VERSION=0.5.0
-DOCKER_PROTOBUF=jaegertracing/protobuf:$(DOCKER_PROTOBUF_VERSION)
-PROTOC := docker run --rm -u ${shell id -u} -v${PWD}:${PWD} -w${PWD} ${DOCKER_PROTOBUF} --proto_path=${PWD}
+CONTAINER=docker
+CONTAINER_PROTOBUF_VERSION=0.5.0
+CONTAINER_PROTOBUF=jaegertracing/protobuf:$(CONTAINER_PROTOBUF_VERSION)
+PROTOC := ${CONTAINER} run --rm -u ${shell id -u} -v${PWD}:${PWD} -w${PWD} ${CONTAINER_PROTOBUF} --proto_path=${PWD}
 
 PATCHED_OTEL_PROTO_DIR = proto-gen/.patched-otel-proto
 

cmd/anonymizer/app/flags.go

@@ -17,6 +17,8 @@ type Options struct {
 	HashCustomTags    bool
 	HashLogs          bool
 	HashProcess       bool
+	StartTime         int64
+	EndTime           int64
 }
 
 const (
@@ -28,6 +30,8 @@ const (
 	hashLogsFlag          = "hash-logs"
 	hashProcessFlag       = "hash-process"
 	maxSpansCount         = "max-spans-count"
+	startTime             = "start-time"
+	endTime               = "end-time"
 )
 
 // AddFlags adds flags for anonymizer main program
@@ -72,6 +76,16 @@ func (o *Options) AddFlags(command *cobra.Command) {
 		maxSpansCount,
 		-1,
 		"The maximum number of spans to anonymize")
+	command.Flags().Int64Var(
+		&o.StartTime,
+		startTime,
+		-1,
+		"The start time of time window for searching trace, timestampe in unix microseconds")
+	command.Flags().Int64Var(
+		&o.EndTime,
+		endTime,
+		-1,
+		"The end time of time window for searching trace, timestampe in unix microseconds")
 
 	// mark traceid flag as mandatory
 	command.MarkFlagRequired(traceIDFlag)

cmd/anonymizer/app/query/query.go

@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"io"
 	"strings"
+	"time"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
@@ -49,15 +50,19 @@ func unwrapNotFoundErr(err error) error {
 }
 
 // QueryTrace queries for a trace and returns all spans inside it
-func (q *Query) QueryTrace(traceID string) ([]model.Span, error) {
+func (q *Query) QueryTrace(traceID string, startTime time.Time, endTime time.Time) ([]model.Span, error) {
 	mTraceID, err := model.TraceIDFromString(traceID)
 	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{
-		TraceID: mTraceID,
-	})
+	request := api_v2.GetTraceRequest{
+		TraceID:   mTraceID,
+		StartTime: startTime,
+		EndTime:   endTime,
+	}
+
+	stream, err := q.client.GetTrace(context.Background(), &request)
 	if err != nil {
 		return nil, unwrapNotFoundErr(err)
 	}

cmd/anonymizer/app/query/query_test.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -24,8 +25,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,24 +107,31 @@ func TestQueryTrace(t *testing.T) {
 	defer q.Close()
 
 	t.Run("No error", func(t *testing.T) {
-		s.spanReader.On("GetTrace", matchContext, matchTraceID).Return(
+		startTime := time.Date(1970, time.January, 1, 0, 0, 0, 1000, time.UTC)
+		endTime := time.Date(1970, time.January, 1, 0, 0, 0, 2000, time.UTC)
+		expectedGetTraceParameters := spanstore.GetTraceParameters{
+			TraceID:   mockTraceID,
+			StartTime: startTime,
+			EndTime:   endTime,
+		}
+		s.spanReader.On("GetTrace", matchContext, expectedGetTraceParameters).Return(
 			mockTraceGRPC, nil).Once()
 
-		spans, err := q.QueryTrace(mockTraceID.String())
+		spans, err := q.QueryTrace(mockTraceID.String(), startTime, endTime)
 		require.NoError(t, err)
 		assert.Equal(t, len(spans), len(mockTraceGRPC.Spans))
 	})
 
 	t.Run("Invalid TraceID", func(t *testing.T) {
-		_, err := q.QueryTrace(mockInvalidTraceID)
+		_, err := q.QueryTrace(mockInvalidTraceID, time.Time{}, time.Time{})
 		assert.ErrorContains(t, err, "failed to convert the provided trace id")
 	})
 
 	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())
+		spans, err := q.QueryTrace(mockTraceID.String(), time.Time{}, time.Time{})
 		assert.Nil(t, spans)
 		assert.ErrorIs(t, err, spanstore.ErrTraceNotFound)
 	})

cmd/anonymizer/main.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"time"
 
 	"github.com/spf13/cobra"
 	"go.uber.org/zap"
@@ -53,7 +54,7 @@ func main() {
 				logger.Fatal("error while creating query object", zap.Error(err))
 			}
 
-			spans, err := query.QueryTrace(options.TraceID)
+			spans, err := query.QueryTrace(options.TraceID, time.UnixMicro(options.StartTime), time.UnixMicro(options.EndTime))
 			if err != nil {
 				logger.Fatal("error while querying for trace", zap.Error(err))
 			}

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())
 

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)

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) {

cmd/query/app/apiv3/grpc_handler.go

@@ -32,7 +32,13 @@ 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)
 	}

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)
+
+	expectedGetTraceParameters := spanstore.GetTraceParameters{
+		TraceID:   traceId,
+		StartTime: time.Time{},
+		EndTime:   time.Time{},
+	}
+	tsc.reader.On("GetTrace", matchContext, expectedGetTraceParameters).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()
+	expectedGetTraceParameters := spanstore.GetTraceParameters{
+		TraceID:   traceId,
+		StartTime: start_ts,
+		EndTime:   end_ts,
+	}
+	tsc.reader.On("GetTrace", matchContext, expectedGetTraceParameters).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()

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,27 @@ 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)
+	query := spanstore.GetTraceParameters{
+		TraceID: traceID,
+	}
+	http_query := r.URL.Query()
+	startTime := http_query.Get(paramStartTime)
+	if startTime != "" {
+		timeParsed, err := time.Parse(time.RFC3339Nano, startTime)
+		if h.tryParamError(w, err, paramStartTime) {
+			return
+		}
+		query.StartTime = timeParsed
+	}
+	endTime := http_query.Get(paramEndTime)
+	if endTime != "" {
+		timeParsed, err := time.Parse(time.RFC3339Nano, endTime)
+		if h.tryParamError(w, err, paramEndTime) {
+			return
+		}
+		query.EndTime = timeParsed
+	}
+	trc, err := h.QueryService.GetTrace(r.Context(), query)
 	if h.tryHandleError(w, err, http.StatusInternalServerError) {
 		return
 	}