Add optinal time window in TraceGetParameters

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

Author: rim99

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,27 @@ 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 int64, endTime int64) ([]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{
+	request := api_v2.GetTraceRequest{
 		TraceID: mTraceID,
-	})
+	}
+
+	if startTime != -1 {
+		t := time.UnixMicro(startTime)
+		request.StartTime = &t
+	}
+
+	if endTime != -1 {
+		t := time.UnixMicro(endTime)
+		request.EndTime = &t
+	}
+
+	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")
+	matchTraceGetParameters = mock.AnythingOfType("spanstore.TraceGetParameters")
 
 	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)
+		expectedTraceGetParameters := spanstore.TraceGetParameters{
+			TraceID:   mockTraceID,
+			StartTime: &startTime,
+			EndTime:   &endTime,
+		}
+		s.spanReader.On("GetTrace", matchContext, expectedTraceGetParameters).Return(
 			mockTraceGRPC, nil).Once()
 
-		spans, err := q.QueryTrace(mockTraceID.String())
+		spans, err := q.QueryTrace(mockTraceID.String(), 1, 2)
 		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, -1, -1)
 		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, matchTraceGetParameters).Return(
 			nil, spanstore.ErrTraceNotFound).Once()
 
-		spans, err := q.QueryTrace(mockTraceID.String())
+		spans, err := q.QueryTrace(mockTraceID.String(), -1, -1)
 		assert.Nil(t, spans)
 		assert.ErrorIs(t, err, spanstore.ErrTraceNotFound)
 	})

cmd/anonymizer/main.go

@@ -53,7 +53,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, options.StartTime, 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 {
@@ -154,7 +155,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.TraceGetParameters{TraceID: requiredTraceID})
 	require.NoError(t, err)
 	assert.Equal(t, spanID.String(), requiredTrace.Spans[0].SpanID.String())
 

cmd/jaeger/internal/extension/jaegerstorage/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.TraceGetParameters{TraceID: traceID}
+	trace, err := memstore.GetTrace(context.Background(), query)
 	require.NoError(t, err)
 	require.NotNil(t, trace)
 	assert.Len(t, trace.Spans, 1)

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.TraceGetParameters) (*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.TraceGetParameters) {
 	traceID := model.NewTraceID(150, 160)
+	query := spanstore.TraceGetParameters{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

@@ -33,7 +33,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.TraceGetParameters{
+		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)
 	}