Support graceful shutdown in grpc plugin (#3249)

* Support graceful shutdown in grpc plugin

Signed-off-by: Fedor Korotkiy <dartslon@gmail.com>

* Update plugin/storage/grpc/shared/grpc_client.go

Co-authored-by: Yuri Shkuro <yurishkuro@users.noreply.github.com>
Signed-off-by: Fedor Korotkiy <dartslon@gmail.com>

* Return Unimplemented status

Signed-off-by: Fedor Korotkiy <dartslon@gmail.com>

Co-authored-by: Yuri Shkuro <yurishkuro@users.noreply.github.com>

Makefile

@@ -415,6 +415,7 @@ install-mockery:
 generate-mocks: install-mockery
 	$(MOCKERY) -all -dir ./pkg/es/ -output ./pkg/es/mocks && rm pkg/es/mocks/ClientBuilder.go
 	$(MOCKERY) -all -dir ./storage/spanstore/ -output ./storage/spanstore/mocks
+	$(MOCKERY) -all -dir ./proto-gen/storage_v1/ -output ./proto-gen/storage_v1/mocks
 
 .PHONY: echo-version
 echo-version:

plugin/storage/grpc/proto/storage.proto

@@ -58,6 +58,14 @@ message WriteSpanResponse {
 
 }
 
+// empty; extensible in the future
+message CloseWriterRequest {
+}
+
+// empty; extensible in the future
+message CloseWriterResponse {
+}
+
 message GetTraceRequest {
     bytes trace_id = 1 [
       (gogoproto.nullable) = false,
@@ -135,6 +143,7 @@ message FindTraceIDsResponse {
 service SpanWriterPlugin {
     // spanstore/Writer
     rpc WriteSpan(WriteSpanRequest) returns (WriteSpanResponse);
+    rpc Close(CloseWriterRequest) returns (CloseWriterResponse);
 }
 
 service SpanReaderPlugin {

plugin/storage/grpc/shared/grpc_client.go

@@ -222,13 +222,23 @@ func (c *grpcClient) WriteSpan(ctx context.Context, span *model.Span) error {
 	_, err := c.writerClient.WriteSpan(ctx, &storage_v1.WriteSpanRequest{
 		Span: span,
 	})
+
 	if err != nil {
 		return fmt.Errorf("plugin error: %w", err)
 	}
 
 	return nil
 }
 
+func (c *grpcClient) Close() error {
+	_, err := c.writerClient.Close(context.Background(), &storage_v1.CloseWriterRequest{})
+	if err != nil && status.Code(err) != codes.Unimplemented {
+		return fmt.Errorf("plugin error: %w", err)
+	}
+
+	return nil
+}
+
 // GetDependencies returns all interservice dependencies
 func (c *grpcClient) GetDependencies(ctx context.Context, endTs time.Time, lookback time.Duration) ([]model.DependencyLink, error) {
 	resp, err := c.depsReaderClient.GetDependencies(ctx, &storage_v1.GetDependenciesRequest{

plugin/storage/grpc/shared/grpc_client_test.go

@@ -298,6 +298,25 @@ func TestGRPCClientWriteSpan(t *testing.T) {
 	})
 }
 
+func TestGRPCClientCloseWriter(t *testing.T) {
+	withGRPCClient(func(r *grpcClientTest) {
+		r.spanWriter.On("Close", mock.Anything, &storage_v1.CloseWriterRequest{}).Return(&storage_v1.CloseWriterResponse{}, nil)
+
+		err := r.client.Close()
+		assert.NoError(t, err)
+	})
+}
+
+func TestGRPCClientCloseNotSupported(t *testing.T) {
+	withGRPCClient(func(r *grpcClientTest) {
+		r.spanWriter.On("Close", mock.Anything, &storage_v1.CloseWriterRequest{}).Return(
+			nil, status.Errorf(codes.Unimplemented, "method not implemented"))
+
+		err := r.client.Close()
+		assert.NoError(t, err)
+	})
+}
+
 func TestGRPCClientGetDependencies(t *testing.T) {
 	withGRPCClient(func(r *grpcClientTest) {
 		lookback := time.Duration(1 * time.Second)

plugin/storage/grpc/shared/grpc_server.go

@@ -17,6 +17,7 @@ package shared
 import (
 	"context"
 	"fmt"
+	"io"
 
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -54,6 +55,18 @@ func (s *grpcServer) WriteSpan(ctx context.Context, r *storage_v1.WriteSpanReque
 	return &storage_v1.WriteSpanResponse{}, nil
 }
 
+func (s *grpcServer) Close(ctx context.Context, r *storage_v1.CloseWriterRequest) (*storage_v1.CloseWriterResponse, error) {
+	if closer, ok := s.Impl.SpanWriter().(io.Closer); ok {
+		if err := closer.Close(); err != nil {
+			return nil, err
+		}
+
+		return &storage_v1.CloseWriterResponse{}, nil
+	} else {
+		return nil, status.Error(codes.Unimplemented, "span writer does not support graceful shutdown")
+	}
+}
+
 // GetTrace takes a traceID and streams a Trace associated with that traceID
 func (s *grpcServer) GetTrace(r *storage_v1.GetTraceRequest, stream storage_v1.SpanReaderPlugin_GetTraceServer) error {
 	trace, err := s.Impl.SpanReader().GetTrace(stream.Context(), r.TraceID)