[grpc][v2] Implement FindTraces Call in gRPC reader for remote storage api v2 (jaegertracing#6962)

<!--
!! Please DELETE this comment before posting.
We appreciate your contribution to the Jaeger project! 👋🎉
-->

## Which problem is this PR solving?
- Towards jaegertracing#6789

## Description of the changes
- This PR implements the `FindTraces` call in the gRPC v2 API client

## How was this change tested?
- Added unit tests

## Checklist
- [x] I have read
https://github.com/jaegertracing/jaeger/blob/master/CONTRIBUTING_GUIDELINES.md
- [x] I have signed all commits
- [x] I have added unit tests for the new functionality
- [x] I have run lint and test steps successfully
  - for `jaeger`: `make lint test`
  - for `jaeger-ui`: `npm run lint` and `npm run test`

---------

Signed-off-by: Mahad Zaryab <[email protected]>
Signed-off-by: amol-verma-allen <[email protected]>

Author: mahadzaryab1

internal/storage/v2/grpc/tracereader.go

@@ -50,7 +50,7 @@ func (tr *TraceReader) GetTraces(
 			Query: query,
 		})
 		if err != nil {
-			yield(nil, fmt.Errorf("received error from grpc reader client: %w", err))
+			yield(nil, fmt.Errorf("failed to execute GetTraces: %w", err))
 			return
 		}
 		for received, err := stream.Recv(); !errors.Is(err, io.EOF); received, err = stream.Recv() {
@@ -68,7 +68,7 @@ func (tr *TraceReader) GetTraces(
 func (tr *TraceReader) GetServices(ctx context.Context) ([]string, error) {
 	resp, err := tr.client.GetServices(ctx, &storage.GetServicesRequest{})
 	if err != nil {
-		return nil, fmt.Errorf("failed to get services: %w", err)
+		return nil, fmt.Errorf("failed to execute GetServices: %w", err)
 	}
 	return resp.Services, nil
 }
@@ -82,7 +82,7 @@ func (tr *TraceReader) GetOperations(
 		SpanKind: params.SpanKind,
 	})
 	if err != nil {
-		return nil, fmt.Errorf("failed to get operations: %w", err)
+		return nil, fmt.Errorf("failed to execute GetOperations: %w", err)
 	}
 	operations := make([]tracestore.Operation, len(resp.Operations))
 	for i, op := range resp.Operations {
@@ -94,11 +94,28 @@ func (tr *TraceReader) GetOperations(
 	return operations, nil
 }
 
-func (*TraceReader) FindTraces(
-	context.Context,
-	tracestore.TraceQueryParams,
+func (tr *TraceReader) FindTraces(
+	ctx context.Context,
+	params tracestore.TraceQueryParams,
 ) iter.Seq2[[]ptrace.Traces, error] {
-	panic("not implemented")
+	return func(yield func([]ptrace.Traces, error) bool) {
+		stream, err := tr.client.FindTraces(ctx, &storage.FindTracesRequest{
+			Query: toProtoQueryParameters(params),
+		})
+		if err != nil {
+			yield(nil, fmt.Errorf("failed to execute FindTraces: %w", err))
+			return
+		}
+		for received, err := stream.Recv(); !errors.Is(err, io.EOF); received, err = stream.Recv() {
+			if err != nil {
+				yield(nil, fmt.Errorf("received error from grpc stream: %w", err))
+				return
+			}
+			if !yield([]ptrace.Traces{received.ToTraces()}, nil) {
+				return
+			}
+		}
+	}
 }
 
 func (tr *TraceReader) FindTraceIDs(

internal/storage/v2/grpc/tracereader_test.go

@@ -59,6 +59,16 @@ func (ts *testServer) GetOperations(
 	}, ts.err
 }
 
+func (ts *testServer) FindTraces(
+	_ *storage.FindTracesRequest,
+	s storage.TraceReader_FindTracesServer,
+) error {
+	for _, trace := range ts.traces {
+		s.Send(trace)
+	}
+	return ts.err
+}
+
 func (ts *testServer) FindTraceIDs(
 	context.Context,
 	*storage.FindTracesRequest,
@@ -215,7 +225,7 @@ func TestTraceReader_GetTraces_GRPCClientError(t *testing.T) {
 	reader := NewTraceReader(conn)
 	getTracesIter := reader.GetTraces(context.Background(), tracestore.GetTraceParams{})
 	_, err = jiter.FlattenWithErrors(getTracesIter)
-	require.ErrorContains(t, err, "received error from grpc reader client")
+	require.ErrorContains(t, err, "failed to execute GetTraces")
 }
 
 func TestTraceReader_GetServices(t *testing.T) {
@@ -237,7 +247,7 @@ func TestTraceReader_GetServices(t *testing.T) {
 			testServer: &testServer{
 				err: assert.AnError,
 			},
-			expectedError: "failed to get services",
+			expectedError: "failed to execute GetServices",
 		},
 	}
 
@@ -282,7 +292,7 @@ func TestTraceReader_GetOperations(t *testing.T) {
 			testServer: &testServer{
 				err: assert.AnError,
 			},
-			expectedError: "failed to get operations",
+			expectedError: "failed to execute GetOperations",
 		},
 	}
 
@@ -306,11 +316,121 @@ func TestTraceReader_GetOperations(t *testing.T) {
 }
 
 func TestTraceReader_FindTraces(t *testing.T) {
-	tr := &TraceReader{}
+	queryParams := tracestore.TraceQueryParams{
+		ServiceName:   "service-a",
+		OperationName: "operation-a",
+		Attributes:    pcommon.NewMap(),
+	}
+	tests := []struct {
+		name           string
+		testServer     *testServer
+		traces         []*jptrace.TracesData
+		expectedTraces []ptrace.Traces
+		expectedError  string
+	}{
+		{
+			name: "single trace",
+			testServer: &testServer{
+				traces: func() []*jptrace.TracesData {
+					trace := makeTestTrace()
+					traces := []*jptrace.TracesData{(*jptrace.TracesData)(&trace)}
+					return traces
+				}(),
+			},
+			expectedTraces: []ptrace.Traces{makeTestTrace()},
+		},
+		{
+			name: "multiple traces",
+			testServer: &testServer{
+				traces: func() []*jptrace.TracesData {
+					traceA := makeTestTrace()
+					traceB := makeTestTrace()
+					traces := []*jptrace.TracesData{
+						(*jptrace.TracesData)(&traceA),
+						(*jptrace.TracesData)(&traceB),
+					}
+					return traces
+				}(),
+			},
+			expectedTraces: []ptrace.Traces{makeTestTrace(), makeTestTrace()},
+		},
+		{
+			name: "error",
+			testServer: &testServer{
+				traces: func() []*jptrace.TracesData {
+					trace := ptrace.NewTraces()
+					traces := []*jptrace.TracesData{(*jptrace.TracesData)(&trace)}
+					return traces
+				}(),
+				err: assert.AnError,
+			},
+			expectedError: "received error from grpc stream",
+		},
+	}
 
-	require.Panics(t, func() {
-		tr.FindTraces(context.Background(), tracestore.TraceQueryParams{})
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			conn := startTestServer(t, test.testServer)
+
+			reader := NewTraceReader(conn)
+			getTracesIter := reader.FindTraces(context.Background(), queryParams)
+			traces, err := jiter.FlattenWithErrors(getTracesIter)
+
+			if test.expectedError != "" {
+				require.ErrorContains(t, err, test.expectedError)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, test.expectedTraces, traces)
+			}
+		})
+	}
+}
+
+func TestTraceReader_FindTraces_YieldStopsIteration(t *testing.T) {
+	queryParams := tracestore.TraceQueryParams{
+		ServiceName:   "service-a",
+		OperationName: "operation-a",
+		Attributes:    pcommon.NewMap(),
+	}
+	traceA := makeTestTrace()
+	traceB := makeTestTrace()
+	testServer := &testServer{
+		traces: []*jptrace.TracesData{
+			(*jptrace.TracesData)(&traceA),
+			(*jptrace.TracesData)(&traceB),
+		},
+	}
+
+	conn := startTestServer(t, testServer)
+	reader := NewTraceReader(conn)
+
+	getTracesIter := reader.FindTraces(context.Background(), queryParams)
+	var gotTraces []ptrace.Traces
+	getTracesIter(func(traces []ptrace.Traces, _ error) bool {
+		gotTraces = append(gotTraces, traces...)
+		return false
+	})
+
+	require.Len(t, gotTraces, 1)
+}
+
+func TestTraceReader_FindTraces_GRPCClientError(t *testing.T) {
+	queryParams := tracestore.TraceQueryParams{
+		ServiceName:   "service-a",
+		OperationName: "operation-a",
+		Attributes:    pcommon.NewMap(),
+	}
+	conn, err := grpc.NewClient(":0",
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	) // create client without a started server
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		conn.Close()
 	})
+	reader := NewTraceReader(conn)
+	getTracesIter := reader.FindTraces(context.Background(), queryParams)
+	_, err = jiter.FlattenWithErrors(getTracesIter)
+	require.ErrorContains(t, err, "failed to execute FindTraces")
 }
 
 func TestTraceReader_FindTraceIDs(t *testing.T) {