[Draft] Mix test

.gitignore

@@ -42,7 +42,6 @@ cmd/tracegen/tracegen-*
 crossdock/crossdock-*
 
 run-crossdock.log
-proto-gen/.patched-otel-proto/
 __pycache__
 .asset-manifest.json
 deploy/

Makefile.Protobuf.mk

@@ -14,15 +14,15 @@
 # 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
 
 PROTO_INCLUDES := \
 	-Iidl/proto/api_v2 \
-	-Iidl/proto/api_v3 \
 	-Imodel/proto/metrics \
 	-I/usr/include/github.com/gogo/protobuf
 
@@ -127,9 +127,18 @@ proto-zipkin:
 # Note that the .pb.go types must be generated into the same internal package $(API_V3_PATH)
 # where a manually defined traces.go file is located.
 API_V3_PATH=cmd/query/app/internal/api_v3
+API_V3_PATCHED_DIR=proto-gen/.patched/api_v3
+API_V3_PATCHED=$(API_V3_PATCHED_DIR)/query_service.proto
+.PHONY: patch-api-v3
+patch-api-v3:
+	mkdir -p $(API_V3_PATCHED_DIR)
+	cat idl/proto/api_v3/query_service.proto | \
+		$(SED) -f ./proto-gen/patch-api-v3.sed \
+		> $(API_V3_PATCHED)
+
 .PHONY: proto-api-v3
-proto-api-v3:
-	$(call proto_compile, $(API_V3_PATH), idl/proto/api_v3/query_service.proto, -Iidl/opentelemetry-proto)
+proto-api-v3: patch-api-v3
+	$(call proto_compile, $(API_V3_PATH), $(API_V3_PATCHED), -I$(API_V3_PATCHED_DIR) -Iidl/opentelemetry-proto)
 	@echo "🏗️  replace TracesData with internal custom type"
 	$(SED) -i 's/v1.TracesData/TracesData/g' $(API_V3_PATH)/query_service.pb.go
 	@echo "🏗️  remove OTEL import because we're not using any other OTLP types"

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,25 @@ 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 {
+		request.StartTime = time.UnixMicro(startTime)
+	}
+
+	if endTime != -1 {
+		request.EndTime = time.UnixMicro(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")
+	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 {
@@ -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.TraceGetParameters{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.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

@@ -8,7 +8,6 @@ import (
 	"errors"
 	"fmt"
 
-	"github.com/gogo/protobuf/types"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
@@ -33,7 +32,16 @@ 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()
+	query.StartTime = &startTime
+	endTime := request.GetEndTime()
+	query.EndTime = &endTime
+
+	trace, err := h.QueryService.GetTrace(stream.Context(), query)
 	if err != nil {
 		return fmt.Errorf("cannot retrieve trace: %w", err)
 	}
@@ -51,44 +59,28 @@ func (h *Handler) FindTraces(request *api_v3.FindTracesRequest, stream api_v3.Qu
 	if query == nil {
 		return status.Error(codes.InvalidArgument, "missing query")
 	}
-	if query.GetStartTimeMin() == nil ||
-		query.GetStartTimeMax() == nil {
+	if query.GetStartTimeMin().IsZero() ||
+		query.GetStartTimeMax().IsZero() {
 		return errors.New("start time min and max are required parameters")
 	}
 
 	queryParams := &spanstore.TraceQueryParameters{
 		ServiceName:   query.GetServiceName(),
 		OperationName: query.GetOperationName(),
 		Tags:          query.GetAttributes(),
-		NumTraces:     int(query.GetNumTraces()),
+		NumTraces:     int(query.GetSearchDepth()),
 	}
-	if query.GetStartTimeMin() != nil {
-		startTimeMin, err := types.TimestampFromProto(query.GetStartTimeMin())
-		if err != nil {
-			return err
-		}
-		queryParams.StartTimeMin = startTimeMin
+	if ts := query.GetStartTimeMin(); !ts.IsZero() {
+		queryParams.StartTimeMin = ts
 	}
-	if query.GetStartTimeMax() != nil {
-		startTimeMax, err := types.TimestampFromProto(query.GetStartTimeMax())
-		if err != nil {
-			return err
-		}
-		queryParams.StartTimeMax = startTimeMax
+	if ts := query.GetStartTimeMax(); !ts.IsZero() {
+		queryParams.StartTimeMax = ts
 	}
-	if query.GetDurationMin() != nil {
-		durationMin, err := types.DurationFromProto(query.GetDurationMin())
-		if err != nil {
-			return err
-		}
-		queryParams.DurationMin = durationMin
+	if d := query.GetDurationMin(); d != 0 {
+		queryParams.DurationMin = d
 	}
-	if query.GetDurationMax() != nil {
-		durationMax, err := types.DurationFromProto(query.GetDurationMax())
-		if err != nil {
-			return err
-		}
-		queryParams.DurationMax = durationMax
+	if d := query.GetDurationMax(); d != 0 {
+		queryParams.DurationMax = d
 	}
 
 	traces, err := h.QueryService.FindTraces(stream.Context(), queryParams)