fix: streaming span writer is not working in grpc based remote storag…

…e plugin (#3887)

* plugin/storage: streaming span writer is not assigned for grpc remote plugin

Signed-off-by: Arunprasad Rajkumar <ar.arunprasad@gmail.com>

* storage/grpc: propagate grpc library errors

grpc_handler.go handles only io.EOF and skipping other errors causes crash.

Signed-off-by: Arunprasad Rajkumar <ar.arunprasad@gmail.com>

Signed-off-by: Arunprasad Rajkumar <ar.arunprasad@gmail.com>

plugin/storage/grpc/config/config.go

@@ -107,8 +107,9 @@ func (c *Configuration) buildRemote(logger *zap.Logger) (*ClientPluginServices,
 	grpcClient := shared.NewGRPCClient(conn)
 	return &ClientPluginServices{
 		PluginServices: shared.PluginServices{
-			Store:        grpcClient,
-			ArchiveStore: grpcClient,
+			Store:               grpcClient,
+			ArchiveStore:        grpcClient,
+			StreamingSpanWriter: grpcClient,
 		},
 		Capabilities: grpcClient,
 	}, nil

plugin/storage/grpc/shared/grpc_handler.go

@@ -115,6 +115,9 @@ func (s *GRPCHandler) WriteSpanStream(stream storage_v1.StreamingSpanWriterPlugi
 		if err == io.EOF {
 			break
 		}
+		if err != nil {
+			return err
+		}
 		err = writer.WriteSpan(stream.Context(), in.Span)
 		if err != nil {
 			return err

plugin/storage/grpc/shared/grpc_handler_test.go

@@ -245,6 +245,20 @@ func TestGRPCServerWriteSpanStream(t *testing.T) {
 	})
 }
 
+func TestGRPCServerWriteSpanStreamWithGRPCError(t *testing.T) {
+	withGRPCServer(func(r *grpcServerTest) {
+		stream := new(grpcMocks.StreamingSpanWriterPlugin_WriteSpanStreamServer)
+		stream.On("Recv").Return(&storage_v1.WriteSpanRequest{Span: &mockTraceSpans[0]}, nil).Twice().
+			On("Recv").Return(nil, context.DeadlineExceeded).Once()
+		stream.On("SendAndClose", &storage_v1.WriteSpanResponse{}).Return(nil)
+		stream.On("Context").Return(context.Background())
+		r.impl.streamWriter.On("WriteSpan", context.Background(), &mockTraceSpans[0]).Return(nil)
+
+		err := r.server.WriteSpanStream(stream)
+		assert.ErrorContains(t, err, context.DeadlineExceeded.Error())
+	})
+}
+
 func TestGRPCServerGetDependencies(t *testing.T) {
 	withGRPCServer(func(r *grpcServerTest) {
 		lookback := time.Duration(1 * time.Second)