octtrpc: Fix span status defer, add tests

It turns out for years that the autogenerated TTRPC spans have not been
marked correctly if the call failed. This is because defers evaluate
their arguments immediately, rather than at the deferred execution time.
Fix this by changing err from an argument to the defer, to a variable
evaluated inside the defer.

Also adds tests for octtrpc client and server interceptors.

Signed-off-by: Kevin Parsons <kevpar@microsoft.com>

pkg/octtrpc/interceptor.go

@@ -77,7 +77,7 @@ func ClientInterceptor(opts ...Option) ttrpc.UnaryClientInterceptor {
 			trace.WithSampler(o.sampler),
 			oc.WithClientSpanKind)
 		defer span.End()
-		defer setSpanStatus(span, err)
+		defer func() { setSpanStatus(span, err) }()
 
 		spanContextBinary := propagation.Binary(span.SpanContext())
 		b64 := base64.StdEncoding.EncodeToString(spanContextBinary)
@@ -110,7 +110,7 @@ func ServerInterceptor(opts ...Option) ttrpc.UnaryServerInterceptor {
 			ctx, span = oc.StartSpan(ctx, name, opts...)
 		}
 		defer span.End()
-		defer setSpanStatus(span, err)
+		defer func() { setSpanStatus(span, err) }()
 
 		return method(ctx, unmarshal)
 	}

pkg/octtrpc/interceptor_test.go

@@ -0,0 +1,203 @@
+package octtrpc
+
+import (
+	"context"
+	"encoding/base64"
+	"fmt"
+	"testing"
+
+	"github.com/containerd/ttrpc"
+	"go.opencensus.io/trace"
+	"go.opencensus.io/trace/propagation"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+type spanExporter struct {
+	spans []*trace.SpanData
+}
+
+func (e *spanExporter) ExportSpan(s *trace.SpanData) {
+	e.spans = append(e.spans, s)
+}
+
+func TestClientInterceptor(t *testing.T) {
+	for name, tc := range map[string]struct {
+		methodName       string
+		expectedSpanName string
+		invoker          func(ctx context.Context, req *ttrpc.Request, resp *ttrpc.Response) error
+		expectedStatus   trace.Status
+	}{
+		"callWithMethodName": {
+			methodName:       "TestMethod",
+			expectedSpanName: "TestMethod",
+		},
+		"callWithWeirdSpanName": {
+			methodName:       "/Test/Method/Foo",
+			expectedSpanName: "Test.Method.Foo",
+		},
+		"callFailsWithGenericError": {
+			invoker: func(ctx context.Context, req *ttrpc.Request, resp *ttrpc.Response) error {
+				return fmt.Errorf("generic error")
+			},
+			expectedStatus: trace.Status{Code: 13, Message: "generic error"},
+		},
+		"callFailsWithGRPCError": {
+			invoker: func(ctx context.Context, req *ttrpc.Request, resp *ttrpc.Response) error {
+				return status.Error(codes.AlreadyExists, "already exists")
+			},
+			expectedStatus: trace.Status{Code: int32(codes.AlreadyExists), Message: "already exists"},
+		},
+	} {
+		t.Run(name, func(t *testing.T) {
+			exporter := &spanExporter{}
+			trace.RegisterExporter(exporter)
+			interceptor := ClientInterceptor(WithSampler(trace.AlwaysSample()))
+
+			methodName := tc.methodName
+			if methodName == "" {
+				methodName = "TestMethod"
+			}
+			invoker := tc.invoker
+			if invoker == nil {
+				invoker = func(ctx context.Context, req *ttrpc.Request, resp *ttrpc.Response) error { return nil }
+			}
+
+			var md []*ttrpc.KeyValue
+
+			interceptor(
+				context.Background(),
+				&ttrpc.Request{},
+				&ttrpc.Response{},
+				&ttrpc.UnaryClientInfo{
+					FullMethod: methodName,
+				},
+				func(ctx context.Context, req *ttrpc.Request, resp *ttrpc.Response) error {
+					md = req.Metadata
+					return invoker(ctx, req, resp)
+				},
+			)
+
+			if len(exporter.spans) != 1 {
+				t.Fatalf("expected exporter to have 1 span but got %d", len(exporter.spans))
+			}
+			span := exporter.spans[0]
+			if tc.expectedSpanName != "" && span.Name != tc.expectedSpanName {
+				t.Errorf("expected span name %s but got %s", tc.expectedSpanName, span.Name)
+			}
+			if span.SpanKind != trace.SpanKindClient {
+				t.Errorf("expected client span kind but got %v", span.SpanKind)
+			}
+			var spanMD string
+			for _, kvp := range md {
+				if kvp.Key == metadataTraceContextKey {
+					spanMD = kvp.Value
+					break
+				}
+			}
+			if spanMD == "" {
+				t.Error("expected span metadata in the request")
+			} else {
+				expectedSpanMD := base64.StdEncoding.EncodeToString(propagation.Binary(span.SpanContext))
+				if spanMD != expectedSpanMD {
+					t.Errorf("expected span metadata %s but got %s", expectedSpanMD, spanMD)
+				}
+			}
+			if span.Status != tc.expectedStatus {
+				t.Errorf("expected status %+v but got %+v", tc.expectedStatus, span.Status)
+			}
+		})
+	}
+}
+
+func TestServerInterceptor(t *testing.T) {
+	for name, tc := range map[string]struct {
+		methodName       string
+		expectedSpanName string
+		method           func(ctx context.Context, unmarshal func(interface{}) error) (interface{}, error)
+		expectedStatus   trace.Status
+		parentSpan       *trace.SpanContext
+	}{
+		"callWithMethodName": {
+			methodName:       "TestMethod",
+			expectedSpanName: "TestMethod",
+		},
+		"callWithWeirdSpanName": {
+			methodName:       "/Test/Method/Foo",
+			expectedSpanName: "Test.Method.Foo",
+		},
+		"callWithRemoteSpanParent": {
+			parentSpan: &trace.SpanContext{
+				TraceID: trace.TraceID{0xa0, 0xa1, 0xa2, 0xa3, 0xa4, 0xa5, 0xa6, 0xa7, 0xa8, 0xa9, 0xaa, 0xab, 0xac, 0xad, 0xae, 0xaf},
+				SpanID:  trace.SpanID{0xb0, 0xb1, 0xb2, 0xb3, 0xb4, 0xb5, 0xb6, 0xb7},
+			},
+		},
+		"callFailsWithGenericError": {
+			method: func(ctx context.Context, unmarshal func(interface{}) error) (interface{}, error) {
+				return nil, fmt.Errorf("generic error")
+			},
+			expectedStatus: trace.Status{Code: 13, Message: "generic error"},
+		},
+		"callFailsWithGRPCError": {
+			method: func(ctx context.Context, unmarshal func(interface{}) error) (interface{}, error) {
+				return nil, status.Error(codes.AlreadyExists, "already exists")
+			},
+			expectedStatus: trace.Status{Code: int32(codes.AlreadyExists), Message: "already exists"},
+		},
+	} {
+		t.Run(name, func(t *testing.T) {
+			exporter := &spanExporter{}
+			trace.RegisterExporter(exporter)
+			interceptor := ServerInterceptor(WithSampler(trace.AlwaysSample()))
+
+			ctx := context.Background()
+			if tc.parentSpan != nil {
+				ctx = ttrpc.WithMetadata(ctx, ttrpc.MD{metadataTraceContextKey: []string{base64.StdEncoding.EncodeToString(propagation.Binary(*tc.parentSpan))}})
+			}
+			methodName := tc.methodName
+			if methodName == "" {
+				methodName = "TestMethod"
+			}
+			method := tc.method
+			if method == nil {
+				method = func(ctx context.Context, unmarshal func(interface{}) error) (interface{}, error) { return nil, nil }
+			}
+
+			interceptor(
+				ctx,
+				nil,
+				&ttrpc.UnaryServerInfo{
+					FullMethod: methodName,
+				},
+				func(ctx context.Context, unmarshal func(interface{}) error) (interface{}, error) {
+					return method(ctx, unmarshal)
+				},
+			)
+
+			if len(exporter.spans) != 1 {
+				t.Fatalf("expected exporter to have 1 span but got %d", len(exporter.spans))
+			}
+			span := exporter.spans[0]
+			if tc.expectedSpanName != "" && span.Name != tc.expectedSpanName {
+				t.Errorf("expected span name %s but got %s", tc.expectedSpanName, span.Name)
+			}
+			if span.SpanKind != trace.SpanKindServer {
+				t.Errorf("expected server span kind but got %v", span.SpanKind)
+			}
+			if tc.parentSpan != nil {
+				if span.TraceID != tc.parentSpan.TraceID {
+					t.Errorf("expected trace id %v but got %v", tc.parentSpan.TraceID, span.TraceID)
+				}
+				if span.ParentSpanID != tc.parentSpan.SpanID {
+					t.Errorf("expected parent span id %v but got %v", tc.parentSpan.SpanID, span.ParentSpanID)
+				}
+				if !span.HasRemoteParent {
+					t.Error("expected span to have remote parent")
+				}
+			}
+			if span.Status != tc.expectedStatus {
+				t.Errorf("expected status %+v but got %+v", tc.expectedStatus, span.Status)
+			}
+		})
+	}
+}