Add FindTraceID to the spanstore interface (#1246)

Fixes #1241

cmd/query/app/handler_archive_test.go

@@ -29,13 +29,13 @@ import (
 
 func TestGetArchivedTrace_NotFound(t *testing.T) {
 	mockReader := &spanstoremocks.Reader{}
-	mockReader.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+	mockReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(nil, spanstore.ErrTraceNotFound).Once()
 	for _, tc := range []spanstore.Reader{nil, mockReader} {
 		archiveReader := tc // capture loop var
 		t.Run(fmt.Sprint(archiveReader), func(t *testing.T) {
 			withTestServer(t, func(ts *testServer) {
-				ts.spanReader.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+				ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 					Return(nil, spanstore.ErrTraceNotFound).Once()
 				var response structuredResponse
 				err := getJSON(ts.server.URL+"/api/traces/"+mockTraceID.String(), &response)
@@ -50,11 +50,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("model.TraceID")).
+	mockReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(mockTrace, nil).Once()
 	withTestServer(t, func(ts *testServer) {
 		// make main reader return NotFound
-		ts.spanReader.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+		ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 			Return(nil, spanstore.ErrTraceNotFound).Once()
 		var response structuredTraceResponse
 		err := getJSON(ts.server.URL+"/api/traces/"+mockTraceID.String(), &response)
@@ -78,7 +78,7 @@ func TestArchiveTrace_Success(t *testing.T) {
 	mockWriter.On("WriteSpan", mock.AnythingOfType("*model.Span")).
 		Return(nil).Times(2)
 	withTestServer(t, func(ts *testServer) {
-		ts.spanReader.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+		ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 			Return(mockTrace, nil).Once()
 		var response structuredResponse
 		err := postJSON(ts.server.URL+"/api/archive/"+mockTraceID.String(), []string{}, &response)
@@ -91,7 +91,7 @@ func TestArchiveTrace_WriteErrors(t *testing.T) {
 	mockWriter.On("WriteSpan", mock.AnythingOfType("*model.Span")).
 		Return(errors.New("cannot save")).Times(2)
 	withTestServer(t, func(ts *testServer) {
-		ts.spanReader.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+		ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 			Return(mockTrace, nil).Once()
 		var response structuredResponse
 		err := postJSON(ts.server.URL+"/api/archive/"+mockTraceID.String(), []string{}, &response)

cmd/query/app/handler_test.go

@@ -133,7 +133,7 @@ func withTestServer(t *testing.T, doTest func(s *testServer), options ...Handler
 func TestGetTraceSuccess(t *testing.T) {
 	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	readMock.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+	readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(mockTrace, nil).Once()
 
 	var response structuredResponse
@@ -252,7 +252,7 @@ func TestGetTrace(t *testing.T) {
 			server, readMock, _ := initializeTestServer(HandlerOptions.Tracer(jaegerTracer))
 			defer server.Close()
 
-			readMock.On("GetTrace", model.NewTraceID(0, 0x123456abc)).
+			readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), model.NewTraceID(0, 0x123456abc)).
 				Return(makeMockTrace(t), nil).Once()
 
 			var response structuredResponse
@@ -273,7 +273,7 @@ func TestGetTrace(t *testing.T) {
 func TestGetTraceDBFailure(t *testing.T) {
 	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	readMock.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+	readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(nil, errStorage).Once()
 
 	var response structuredResponse
@@ -284,7 +284,7 @@ func TestGetTraceDBFailure(t *testing.T) {
 func TestGetTraceNotFound(t *testing.T) {
 	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	readMock.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+	readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(nil, spanstore.ErrTraceNotFound).Once()
 
 	var response structuredResponse
@@ -301,7 +301,7 @@ func TestGetTraceAdjustmentFailure(t *testing.T) {
 		),
 	)
 	defer server.Close()
-	readMock.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+	readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(mockTrace, nil).Once()
 
 	var response structuredResponse
@@ -323,7 +323,7 @@ func TestGetTraceBadTraceID(t *testing.T) {
 func TestSearchSuccess(t *testing.T) {
 	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	readMock.On("FindTraces", mock.AnythingOfType("*spanstore.TraceQueryParameters")).
+	readMock.On("FindTraces", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*spanstore.TraceQueryParameters")).
 		Return([]*model.Trace{mockTrace}, nil).Once()
 
 	var response structuredResponse
@@ -335,7 +335,7 @@ func TestSearchSuccess(t *testing.T) {
 func TestSearchByTraceIDSuccess(t *testing.T) {
 	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	readMock.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+	readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(mockTrace, nil).Twice()
 
 	var response structuredResponse
@@ -349,9 +349,9 @@ func TestSearchByTraceIDSuccessWithArchive(t *testing.T) {
 	archiveReadMock := &spanstoremocks.Reader{}
 	server, readMock, _ := initializeTestServer(HandlerOptions.ArchiveSpanReader(archiveReadMock))
 	defer server.Close()
-	readMock.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+	readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(nil, spanstore.ErrTraceNotFound).Twice()
-	archiveReadMock.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+	archiveReadMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(mockTrace, nil).Twice()
 
 	var response structuredResponse
@@ -364,7 +364,7 @@ func TestSearchByTraceIDSuccessWithArchive(t *testing.T) {
 func TestSearchByTraceIDNotFound(t *testing.T) {
 	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	readMock.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+	readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(nil, spanstore.ErrTraceNotFound).Once()
 
 	var response structuredResponse
@@ -378,7 +378,7 @@ func TestSearchByTraceIDFailure(t *testing.T) {
 	server, readMock, _ := initializeTestServer()
 	defer server.Close()
 	whatsamattayou := "https://youtu.be/WrKFOCg13QQ"
-	readMock.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+	readMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("model.TraceID")).
 		Return(nil, fmt.Errorf(whatsamattayou)).Once()
 
 	var response structuredResponse
@@ -395,7 +395,7 @@ func TestSearchModelConversionFailure(t *testing.T) {
 		),
 	)
 	defer server.Close()
-	readMock.On("FindTraces", mock.AnythingOfType("*spanstore.TraceQueryParameters")).
+	readMock.On("FindTraces", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*spanstore.TraceQueryParameters")).
 		Return([]*model.Trace{mockTrace}, nil).Once()
 	var response structuredResponse
 	err := getJSON(server.URL+`/api/traces?service=service&start=0&end=0&operation=operation&limit=200&minDuration=20ms`, &response)
@@ -407,7 +407,7 @@ func TestSearchModelConversionFailure(t *testing.T) {
 func TestSearchDBFailure(t *testing.T) {
 	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	readMock.On("FindTraces", mock.AnythingOfType("*spanstore.TraceQueryParameters")).
+	readMock.On("FindTraces", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*spanstore.TraceQueryParameters")).
 		Return(nil, fmt.Errorf("whatsamattayou")).Once()
 
 	var response structuredResponse
@@ -449,7 +449,7 @@ func TestGetServicesSuccess(t *testing.T) {
 	server, readMock, _ := initializeTestServer()
 	defer server.Close()
 	expectedServices := []string{"trifle", "bling"}
-	readMock.On("GetServices").Return(expectedServices, nil).Once()
+	readMock.On("GetServices", mock.AnythingOfType("*context.valueCtx")).Return(expectedServices, nil).Once()
 
 	var response structuredResponse
 	err := getJSON(server.URL+"/api/services", &response)
@@ -462,20 +462,20 @@ func TestGetServicesSuccess(t *testing.T) {
 }
 
 func TestGetServicesStorageFailure(t *testing.T) {
-	server, mock, _ := initializeTestServer()
+	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	mock.On("GetServices").Return(nil, errStorage).Once()
+	readMock.On("GetServices", mock.AnythingOfType("*context.valueCtx")).Return(nil, errStorage).Once()
 
 	var response structuredResponse
 	err := getJSON(server.URL+"/api/services", &response)
 	assert.Error(t, err)
 }
 
 func TestGetOperationsSuccess(t *testing.T) {
-	server, mock, _ := initializeTestServer()
+	server, readMock, _ := initializeTestServer()
 	defer server.Close()
 	expectedOperations := []string{"", "get"}
-	mock.On("GetOperations", "abc/trifle").Return(expectedOperations, nil).Once()
+	readMock.On("GetOperations", mock.AnythingOfType("*context.valueCtx"), "abc/trifle").Return(expectedOperations, nil).Once()
 
 	var response structuredResponse
 	err := getJSON(server.URL+"/api/operations?service=abc%2Ftrifle", &response)
@@ -497,20 +497,20 @@ func TestGetOperationsNoServiceName(t *testing.T) {
 }
 
 func TestGetOperationsStorageFailure(t *testing.T) {
-	server, mock, _ := initializeTestServer()
+	server, reaMock, _ := initializeTestServer()
 	defer server.Close()
-	mock.On("GetOperations", "trifle").Return(nil, errStorage).Once()
+	reaMock.On("GetOperations", mock.AnythingOfType("*context.valueCtx"), "trifle").Return(nil, errStorage).Once()
 
 	var response structuredResponse
 	err := getJSON(server.URL+"/api/operations?service=trifle", &response)
 	assert.Error(t, err)
 }
 
 func TestGetOperationsLegacySuccess(t *testing.T) {
-	server, mock, _ := initializeTestServer()
+	server, readMock, _ := initializeTestServer()
 	defer server.Close()
 	expectedOperations := []string{"", "get"}
-	mock.On("GetOperations", "abc/trifle").Return(expectedOperations, nil).Once()
+	readMock.On("GetOperations", mock.AnythingOfType("*context.valueCtx"), "abc/trifle").Return(expectedOperations, nil).Once()
 
 	var response structuredResponse
 	err := getJSON(server.URL+"/api/services/abc%2Ftrifle/operations", &response)
@@ -523,9 +523,9 @@ func TestGetOperationsLegacySuccess(t *testing.T) {
 }
 
 func TestGetOperationsLegacyStorageFailure(t *testing.T) {
-	server, mock, _ := initializeTestServer()
+	server, readMock, _ := initializeTestServer()
 	defer server.Close()
-	mock.On("GetOperations", "trifle").Return(nil, errStorage).Once()
+	readMock.On("GetOperations", mock.AnythingOfType("*context.valueCtx"), "trifle").Return(nil, errStorage).Once()
 
 	var response structuredResponse
 	err := getJSON(server.URL+"/api/services/trifle/operations", &response)

plugin/storage/cassandra/spanstore/reader.go

@@ -232,29 +232,44 @@ func validateQuery(p *spanstore.TraceQueryParameters) error {
 
 // FindTraces retrieves traces that match the traceQuery
 func (s *SpanReader) FindTraces(ctx context.Context, traceQuery *spanstore.TraceQueryParameters) ([]*model.Trace, error) {
+	uniqueTraceIDs, err := s.FindTraceIDs(ctx, traceQuery)
+	if err != nil {
+		return nil, err
+	}
+	var retMe []*model.Trace
+	for _, traceID := range uniqueTraceIDs {
+		jTrace, err := s.GetTrace(ctx, traceID)
+		if err != nil {
+			s.logger.Error("Failure to read trace", zap.String("trace_id", traceID.String()), zap.Error(err))
+			continue
+		}
+		retMe = append(retMe, jTrace)
+	}
+	return retMe, nil
+}
+
+// FindTraceIDs retrieve traceIDs that match the traceQuery
+func (s *SpanReader) FindTraceIDs(ctx context.Context, traceQuery *spanstore.TraceQueryParameters) ([]model.TraceID, error) {
 	if err := validateQuery(traceQuery); err != nil {
 		return nil, err
 	}
 	if traceQuery.NumTraces == 0 {
 		traceQuery.NumTraces = defaultNumTraces
 	}
-	uniqueTraceIDs, err := s.findTraceIDs(ctx, traceQuery)
+
+	dbTraceIDs, err := s.findTraceIDs(ctx, traceQuery)
 	if err != nil {
 		return nil, err
 	}
-	var retMe []*model.Trace
-	for traceID := range uniqueTraceIDs {
-		if len(retMe) >= traceQuery.NumTraces {
+
+	var traceIDs []model.TraceID
+	for t := range dbTraceIDs {
+		if len(traceIDs) >= traceQuery.NumTraces {
 			break
 		}
-		jTrace, err := s.readTrace(ctx, traceID)
-		if err != nil {
-			s.logger.Error("Failure to read trace", zap.String("trace_id", traceID.String()), zap.Error(err))
-			continue
-		}
-		retMe = append(retMe, jTrace)
+		traceIDs = append(traceIDs, t.ToDomain())
 	}
-	return retMe, nil
+	return traceIDs, nil
 }
 
 func (s *SpanReader) findTraceIDs(ctx context.Context, traceQuery *spanstore.TraceQueryParameters) (dbmodel.UniqueTraceIDs, error) {

plugin/storage/es/spanstore/reader.go

@@ -234,6 +234,10 @@ func (s *SpanReader) FindTraces(ctx context.Context, traceQuery *spanstore.Trace
 	return s.multiRead(ctx, uniqueTraceIDs, traceQuery.StartTimeMin, traceQuery.StartTimeMax)
 }
 
+func (s *SpanReader) FindTraceIDs(ctx context.Context, traceQuery *spanstore.TraceQueryParameters) ([]model.TraceID, error) {
+	return nil, errors.New("not implemented") // TODO: Implement
+}
+
 func (s *SpanReader) multiRead(ctx context.Context, traceIDs []string, startTime, endTime time.Time) ([]*model.Trace, error) {
 
 	childSpan, _ := opentracing.StartSpanFromContext(ctx, "multiRead")

plugin/storage/es/spanstore/reader_test.go

@@ -668,6 +668,14 @@ func TestFindTraceIDs(t *testing.T) {
 	testGet(traceIDAggregation, t)
 }
 
+func TestFindTraceIDNotImplemented(t *testing.T) {
+	withSpanReader(func(r *spanReaderTest) {
+		traceIDs, err := r.reader.FindTraceIDs(context.Background(), nil)
+		assert.Nil(t, traceIDs)
+		assert.EqualError(t, err, "not implemented")
+	})
+}
+
 func mockMultiSearchService(r *spanReaderTest) *mock.Call {
 	multiSearchService := &mocks.MultiSearchService{}
 	multiSearchService.On("Add", mock.Anything, mock.Anything, mock.Anything).Return(multiSearchService)

plugin/storage/integration/kafka_test.go

@@ -113,6 +113,10 @@ func (r *ingester) FindTraces(ctx context.Context, query *spanstore.TraceQueryPa
 	return nil, nil
 }
 
+func (r *ingester) FindTraceIDs(ctx context.Context, query *spanstore.TraceQueryParameters) ([]model.TraceID, error) {
+	return nil, nil
+}
+
 func TestKafkaStorage(t *testing.T) {
 	if os.Getenv("STORAGE") != "kafka" {
 		t.Skip("Integration test against kafka skipped; set STORAGE env var to kafka to run this")

plugin/storage/memory/memory.go

@@ -198,6 +198,10 @@ func (m *Store) FindTraces(ctx context.Context, query *spanstore.TraceQueryParam
 	return retMe, nil
 }
 
+func (m *Store) FindTraceIDs(ctx context.Context, query *spanstore.TraceQueryParameters) ([]model.TraceID, error) {
+	return nil, errors.New("not implemented")
+}
+
 func (m *Store) validTrace(trace *model.Trace, query *spanstore.TraceQueryParameters) bool {
 	for _, span := range trace.Spans {
 		if m.validSpan(span, query) {

plugin/storage/memory/memory_test.go

@@ -326,3 +326,11 @@ func TestStoreGetTrace(t *testing.T) {
 		})
 	}
 }
+
+func TestStore_FindTraceIDs(t *testing.T) {
+	withMemoryStore(func(store *Store) {
+		traceIDs, err := store.FindTraceIDs(context.Background(), nil)
+		assert.Nil(t, traceIDs)
+		assert.EqualError(t, err, "not implemented")
+	})
+}

storage/spanstore/interface.go

@@ -38,6 +38,7 @@ type Reader interface {
 	GetServices(ctx context.Context) ([]string, error)
 	GetOperations(ctx context.Context, service string) ([]string, error)
 	FindTraces(ctx context.Context, query *TraceQueryParameters) ([]*model.Trace, error)
+	FindTraceIDs(ctx context.Context, query *TraceQueryParameters) ([]model.TraceID, error)
 }
 
 // TraceQueryParameters contains parameters of a trace query.

storage/spanstore/metrics/decorator.go

@@ -28,6 +28,7 @@ import (
 type ReadMetricsDecorator struct {
 	spanReader           spanstore.Reader
 	findTracesMetrics    *queryMetrics
+	findTraceIDsMetrics  *queryMetrics
 	getTraceMetrics      *queryMetrics
 	getServicesMetrics   *queryMetrics
 	getOperationsMetrics *queryMetrics
@@ -36,7 +37,7 @@ type ReadMetricsDecorator struct {
 type queryMetrics struct {
 	Errors     metrics.Counter `metric:"requests" tags:"result=err"`
 	Successes  metrics.Counter `metric:"requests" tags:"result=ok"`
-	Responses  metrics.Timer   `metric:"responses"` //used as a histogram, not necessary for GetTrace
+	Responses  metrics.Timer   `metric:"responses"` // used as a histogram, not necessary for GetTrace
 	ErrLatency metrics.Timer   `metric:"latency" tags:"result=err"`
 	OKLatency  metrics.Timer   `metric:"latency" tags:"result=ok"`
 }
@@ -57,6 +58,7 @@ func NewReadMetricsDecorator(spanReader spanstore.Reader, metricsFactory metrics
 	return &ReadMetricsDecorator{
 		spanReader:           spanReader,
 		findTracesMetrics:    buildQueryMetrics("find_traces", metricsFactory),
+		findTraceIDsMetrics:  buildQueryMetrics("find_trace_ids", metricsFactory),
 		getTraceMetrics:      buildQueryMetrics("get_trace", metricsFactory),
 		getServicesMetrics:   buildQueryMetrics("get_services", metricsFactory),
 		getOperationsMetrics: buildQueryMetrics("get_operations", metricsFactory),
@@ -78,6 +80,14 @@ func (m *ReadMetricsDecorator) FindTraces(ctx context.Context, traceQuery *spans
 	return retMe, err
 }
 
+// FindTraceIDs implements spanstore.Reader#FindTraceIDs
+func (m *ReadMetricsDecorator) FindTraceIDs(ctx context.Context, traceQuery *spanstore.TraceQueryParameters) ([]model.TraceID, error) {
+	start := time.Now()
+	retMe, err := m.spanReader.FindTraceIDs(ctx, traceQuery)
+	m.findTraceIDsMetrics.emit(err, time.Since(start), len(retMe))
+	return retMe, err
+}
+
 // GetTrace implements spanstore.Reader#GetTrace
 func (m *ReadMetricsDecorator) GetTrace(ctx context.Context, traceID model.TraceID) (*model.Trace, error) {
 	start := time.Now()

storage/spanstore/metrics/decorator_test.go

@@ -33,14 +33,16 @@ func TestSuccessfulUnderlyingCalls(t *testing.T) {
 
 	mockReader := mocks.Reader{}
 	mrs := NewReadMetricsDecorator(&mockReader, mf)
-	mockReader.On("GetServices").Return([]string{}, nil)
+	mockReader.On("GetServices", context.Background()).Return([]string{}, nil)
 	mrs.GetServices(context.Background())
-	mockReader.On("GetOperations", "something").Return([]string{}, nil)
+	mockReader.On("GetOperations", context.Background(), "something").Return([]string{}, nil)
 	mrs.GetOperations(context.Background(), "something")
-	mockReader.On("GetTrace", model.TraceID{}).Return(&model.Trace{}, nil)
+	mockReader.On("GetTrace", context.Background(), model.TraceID{}).Return(&model.Trace{}, nil)
 	mrs.GetTrace(context.Background(), model.TraceID{})
-	mockReader.On("FindTraces", &spanstore.TraceQueryParameters{}).Return([]*model.Trace{}, nil)
+	mockReader.On("FindTraces", context.Background(), &spanstore.TraceQueryParameters{}).Return([]*model.Trace{}, nil)
 	mrs.FindTraces(context.Background(), &spanstore.TraceQueryParameters{})
+	mockReader.On("FindTraceIDs", context.Background(), &spanstore.TraceQueryParameters{}).Return([]model.TraceID{}, nil)
+	mrs.FindTraceIDs(context.Background(), &spanstore.TraceQueryParameters{})
 	counters, gauges := mf.Snapshot()
 	expecteds := map[string]int64{
 		"requests|operation=get_operations|result=ok":  1,
@@ -49,6 +51,8 @@ func TestSuccessfulUnderlyingCalls(t *testing.T) {
 		"requests|operation=get_trace|result=err":      0,
 		"requests|operation=find_traces|result=ok":     1,
 		"requests|operation=find_traces|result=err":    0,
+		"requests|operation=find_trace_ids|result=ok":  1,
+		"requests|operation=find_trace_ids|result=err": 0,
 		"requests|operation=get_services|result=ok":    1,
 		"requests|operation=get_services|result=err":   0,
 	}
@@ -86,14 +90,16 @@ func TestFailingUnderlyingCalls(t *testing.T) {
 
 	mockReader := mocks.Reader{}
 	mrs := NewReadMetricsDecorator(&mockReader, mf)
-	mockReader.On("GetServices").Return(nil, errors.New("Failure"))
+	mockReader.On("GetServices", context.Background()).Return(nil, errors.New("Failure"))
 	mrs.GetServices(context.Background())
-	mockReader.On("GetOperations", "something").Return(nil, errors.New("Failure"))
+	mockReader.On("GetOperations", context.Background(), "something").Return(nil, errors.New("Failure"))
 	mrs.GetOperations(context.Background(), "something")
-	mockReader.On("GetTrace", model.TraceID{}).Return(nil, errors.New("Failure"))
+	mockReader.On("GetTrace", context.Background(), model.TraceID{}).Return(nil, errors.New("Failure"))
 	mrs.GetTrace(context.Background(), model.TraceID{})
-	mockReader.On("FindTraces", &spanstore.TraceQueryParameters{}).Return(nil, errors.New("Failure"))
+	mockReader.On("FindTraces", context.Background(), &spanstore.TraceQueryParameters{}).Return(nil, errors.New("Failure"))
 	mrs.FindTraces(context.Background(), &spanstore.TraceQueryParameters{})
+	mockReader.On("FindTraceIDs", context.Background(), &spanstore.TraceQueryParameters{}).Return(nil, errors.New("Failure"))
+	mrs.FindTraceIDs(context.Background(), &spanstore.TraceQueryParameters{})
 	counters, gauges := mf.Snapshot()
 	expecteds := map[string]int64{
 		"requests|operation=get_operations|result=ok":  0,
@@ -102,6 +108,8 @@ func TestFailingUnderlyingCalls(t *testing.T) {
 		"requests|operation=get_trace|result=err":      1,
 		"requests|operation=find_traces|result=ok":     0,
 		"requests|operation=find_traces|result=err":    1,
+		"requests|operation=find_trace_ids|result=ok":   0,
+		"requests|operation=find_trace_ids|result=err":  1,
 		"requests|operation=get_services|result=ok":    0,
 		"requests|operation=get_services|result=err":   1,
 	}

storage/spanstore/mocks/Reader.go

@@ -14,110 +14,127 @@
 
 package mocks
 
-import (
-	"context"
-
-	mock "github.com/stretchr/testify/mock"
-
-	model "github.com/jaegertracing/jaeger/model"
-	spanstore "github.com/jaegertracing/jaeger/storage/spanstore"
-)
+import context "context"
+import mock "github.com/stretchr/testify/mock"
+import model "github.com/jaegertracing/jaeger/model"
+import spanstore "github.com/jaegertracing/jaeger/storage/spanstore"
 
 // Reader is an autogenerated mock type for the Reader type
 type Reader struct {
 	mock.Mock
 }
 
-// FindTraces provides a mock function with given fields: query
+// FindTraceIDs provides a mock function with given fields: ctx, query
+func (_m *Reader) FindTraceIDs(ctx context.Context, query *spanstore.TraceQueryParameters) ([]model.TraceID, error) {
+	ret := _m.Called(ctx, query)
+
+	var r0 []model.TraceID
+	if rf, ok := ret.Get(0).(func(context.Context, *spanstore.TraceQueryParameters) []model.TraceID); ok {
+		r0 = rf(ctx, query)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]model.TraceID)
+		}
+	}
+
+	var r1 error
+	if rf, ok := ret.Get(1).(func(context.Context, *spanstore.TraceQueryParameters) error); ok {
+		r1 = rf(ctx, query)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// FindTraces provides a mock function with given fields: ctx, query
 func (_m *Reader) FindTraces(ctx context.Context, query *spanstore.TraceQueryParameters) ([]*model.Trace, error) {
-	ret := _m.Called(query)
+	ret := _m.Called(ctx, query)
 
 	var r0 []*model.Trace
-	if rf, ok := ret.Get(0).(func(*spanstore.TraceQueryParameters) []*model.Trace); ok {
-		r0 = rf(query)
+	if rf, ok := ret.Get(0).(func(context.Context, *spanstore.TraceQueryParameters) []*model.Trace); ok {
+		r0 = rf(ctx, query)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]*model.Trace)
 		}
 	}
 
 	var r1 error
-	if rf, ok := ret.Get(1).(func(*spanstore.TraceQueryParameters) error); ok {
-		r1 = rf(query)
+	if rf, ok := ret.Get(1).(func(context.Context, *spanstore.TraceQueryParameters) error); ok {
+		r1 = rf(ctx, query)
 	} else {
 		r1 = ret.Error(1)
 	}
 
 	return r0, r1
 }
 
-// GetOperations provides a mock function with given fields: service
+// GetOperations provides a mock function with given fields: ctx, service
 func (_m *Reader) GetOperations(ctx context.Context, service string) ([]string, error) {
-	ret := _m.Called(service)
+	ret := _m.Called(ctx, service)
 
 	var r0 []string
-	if rf, ok := ret.Get(0).(func(string) []string); ok {
-		r0 = rf(service)
+	if rf, ok := ret.Get(0).(func(context.Context, string) []string); ok {
+		r0 = rf(ctx, service)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]string)
 		}
 	}
 
 	var r1 error
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(service)
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, service)
 	} else {
 		r1 = ret.Error(1)
 	}
 
 	return r0, r1
 }
 
-// GetServices provides a mock function with given fields:
+// GetServices provides a mock function with given fields: ctx
 func (_m *Reader) GetServices(ctx context.Context) ([]string, error) {
-	ret := _m.Called()
+	ret := _m.Called(ctx)
 
 	var r0 []string
-	if rf, ok := ret.Get(0).(func() []string); ok {
-		r0 = rf()
+	if rf, ok := ret.Get(0).(func(context.Context) []string); ok {
+		r0 = rf(ctx)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]string)
 		}
 	}
 
 	var r1 error
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
+	if rf, ok := ret.Get(1).(func(context.Context) error); ok {
+		r1 = rf(ctx)
 	} else {
 		r1 = ret.Error(1)
 	}
 
 	return r0, r1
 }
 
-// GetTrace provides a mock function with given fields: traceID
+// 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(traceID)
+	ret := _m.Called(ctx, traceID)
 
 	var r0 *model.Trace
-	if rf, ok := ret.Get(0).(func(model.TraceID) *model.Trace); ok {
-		r0 = rf(traceID)
+	if rf, ok := ret.Get(0).(func(context.Context, model.TraceID) *model.Trace); ok {
+		r0 = rf(ctx, traceID)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(*model.Trace)
 		}
 	}
 
 	var r1 error
-	if rf, ok := ret.Get(1).(func(model.TraceID) error); ok {
-		r1 = rf(traceID)
+	if rf, ok := ret.Get(1).(func(context.Context, model.TraceID) error); ok {
+		r1 = rf(ctx, traceID)
 	} else {
 		r1 = ret.Error(1)
 	}
 
 	return r0, r1
 }
-
-var _ spanstore.Reader = (*Reader)(nil)