Handling of expected error codes coming from grpc storage plugins #1741

plugin/storage/grpc/shared/grpc_client.go

@@ -20,6 +20,7 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
+	"google.golang.org/grpc/status"
 
 	"github.com/jaegertracing/jaeger/model"
 	"github.com/jaegertracing/jaeger/proto-gen/storage_v1"
@@ -61,6 +62,11 @@ func (c *grpcClient) GetTrace(ctx context.Context, traceID model.TraceID) (*mode
 	trace := model.Trace{}
 	for received, err := stream.Recv(); err != io.EOF; received, err = stream.Recv() {
 		if err != nil {
+			if e, ok := status.FromError(err); !ok {
+				if e.Message() == spanstore.ErrTraceNotFound.Error() {
+					return nil, spanstore.ErrTraceNotFound
+				}
+			}
 			return nil, errors.Wrap(err, "grpc stream error")
 		}
 

plugin/storage/grpc/shared/grpc_client_test.go

@@ -166,6 +166,20 @@ func TestGRPCClientGetTrace_NoTrace(t *testing.T) {
 	})
 }
 
+func TestGRPCClientGetTrace_StreamErrorTraceNotFound(t *testing.T) {
+	withGRPCClient(func(r *grpcClientTest) {
+		traceClient := new(grpcMocks.SpanReaderPlugin_GetTraceClient)
+		traceClient.On("Recv").Return(nil, spanstore.ErrTraceNotFound)
+		r.spanReader.On("GetTrace", mock.Anything, &storage_v1.GetTraceRequest{
+			TraceID: mockTraceID,
+		}).Return(traceClient, nil)
+
+		s, err := r.client.GetTrace(context.Background(), mockTraceID)
+		assert.Error(t, err)
+		assert.Nil(t, s)
+	})
+}
+
 func TestGRPCClientFindTraces(t *testing.T) {
 	withGRPCClient(func(r *grpcClientTest) {
 		traceClient := new(grpcMocks.SpanReaderPlugin_FindTracesClient)

plugin/storage/memory/memory.go

@@ -28,8 +28,6 @@ import (
 	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
-var errTraceNotFound = errors.New("trace was not found")
-
 // Store is an in-memory store of traces
 type Store struct {
 	sync.RWMutex
@@ -154,7 +152,7 @@ func (m *Store) GetTrace(ctx context.Context, traceID model.TraceID) (*model.Tra
 	defer m.RUnlock()
 	retMe := m.traces[traceID]
 	if retMe == nil {
-		return nil, errTraceNotFound
+		return nil, spanstore.ErrTraceNotFound
 	}
 	return retMe, nil
 }

plugin/storage/memory/memory_test.go

@@ -210,7 +210,7 @@ func TestStoreGetTraceSuccess(t *testing.T) {
 func TestStoreGetTraceFailure(t *testing.T) {
 	withPopulatedMemoryStore(func(store *Store) {
 		trace, err := store.GetTrace(context.Background(), model.TraceID{})
-		assert.EqualError(t, err, errTraceNotFound.Error())
+		assert.EqualError(t, err, spanstore.ErrTraceNotFound.Error())
 		assert.Nil(t, trace)
 	})
 }