Implement gRPC server zipkin handler. (#96)

Implement server zipkin handler.

Author: anuraaga

middleware/grpc/client.go

@@ -1,15 +1,10 @@
-// +build go1.9
-
 package grpc
 
 import (
 	"context"
-	"strings"
 
-	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/stats"
-	"google.golang.org/grpc/status"
 
 	"github.com/openzipkin/zipkin-go"
 	"github.com/openzipkin/zipkin-go/model"
@@ -18,28 +13,23 @@ import (
 
 type clientHandler struct {
 	tracer            *zipkin.Tracer
-	rpcHandlers       []RPCHandler
 	remoteServiceName string
 }
 
 // A ClientOption can be passed to NewClientHandler to customize the returned handler.
 type ClientOption func(*clientHandler)
 
-// A RPCHandler can be registered using WithRPCHandler to intercept calls to HandleRPC of a
-// handler for additional span customization.
-type RPCHandler func(span zipkin.Span, rpcStats stats.RPCStats)
-
-// WithRPCHandler allows one to add custom logic for handling a stats.RPCStats, e.g.,
-// to add additional tags.
-func WithRPCHandler(handler RPCHandler) ClientOption {
+// WithRemoteServiceName will set the value for the remote endpoint's service name on
+// all spans.
+func WithRemoteServiceName(name string) ClientOption {
 	return func(c *clientHandler) {
-		c.rpcHandlers = append(c.rpcHandlers, handler)
+		c.remoteServiceName = name
 	}
 }
 
 // NewClientHandler returns a stats.Handler which can be used with grpc.WithStatsHandler to add
 // tracing to a gRPC client. The gRPC method name is used as the span name and by default the only
-// tags are the gRPC status code if the call fails. Use WithRPCHandler to add additional tags.
+// tags are the gRPC status code if the call fails.
 func NewClientHandler(tracer *zipkin.Tracer, options ...ClientOption) stats.Handler {
 	c := &clientHandler{
 		tracer: tracer,
@@ -63,36 +53,18 @@ func (c *clientHandler) TagConn(ctx context.Context, cti *stats.ConnTagInfo) con
 
 // HandleRPC implements per-RPC tracing and stats instrumentation.
 func (c *clientHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
-	span := zipkin.SpanFromContext(ctx)
-
-	for _, h := range c.rpcHandlers {
-		h(span, rs)
-	}
-
-	switch rs := rs.(type) {
-	case *stats.End:
-		s, ok := status.FromError(rs.Error)
-		// rs.Error should always be convertable to a status, this is just a defensive check.
-		if ok {
-			if s.Code() != codes.OK {
-				// Uppercase for consistency with Brave
-				c := strings.ToUpper(s.Code().String())
-				span.Tag("grpc.status_code", c)
-				zipkin.TagError.Set(span, c)
-			}
-		} else {
-			zipkin.TagError.Set(span, rs.Error.Error())
-		}
-		span.Finish()
-	}
+	handleRPC(ctx, rs)
 }
 
 // TagRPC implements per-RPC context management.
 func (c *clientHandler) TagRPC(ctx context.Context, rti *stats.RPCTagInfo) context.Context {
 	var span zipkin.Span
 
+	ep := remoteEndpointFromContext(ctx, c.remoteServiceName)
+
 	name := spanName(rti)
-	span, ctx = c.tracer.StartSpanFromContext(ctx, name, zipkin.Kind(model.Client))
+	span, ctx = c.tracer.StartSpanFromContext(ctx, name, zipkin.Kind(model.Client), zipkin.RemoteEndpoint(ep))
+
 	md, ok := metadata.FromOutgoingContext(ctx)
 	if ok {
 		md = md.Copy()

middleware/grpc/client_test.go

@@ -1,5 +1,3 @@
-// +build go1.9
-
 package grpc_test
 
 import (
@@ -9,10 +7,10 @@ import (
 	"github.com/onsi/gomega"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
-	"google.golang.org/grpc/stats"
 
 	"github.com/openzipkin/zipkin-go"
 	zipkingrpc "github.com/openzipkin/zipkin-go/middleware/grpc"
+	"github.com/openzipkin/zipkin-go/model"
 	"github.com/openzipkin/zipkin-go/propagation/b3"
 	service "github.com/openzipkin/zipkin-go/proto/testing"
 	"github.com/openzipkin/zipkin-go/reporter/recorder"
@@ -32,7 +30,7 @@ var _ = ginkgo.Describe("gRPC Client", func() {
 		reporter = recorder.NewReporter()
 		ep, _ := zipkin.NewEndpoint("grpcClient", "")
 		tracer, err = zipkin.NewTracer(
-			reporter, zipkin.WithLocalEndpoint(ep), zipkin.WithIDGenerator(newSequentialIdGenerator()))
+			reporter, zipkin.WithLocalEndpoint(ep), zipkin.WithIDGenerator(newSequentialIdGenerator(1)))
 		gomega.Expect(tracer, err).ToNot(gomega.BeNil())
 	})
 
@@ -56,8 +54,12 @@ var _ = ginkgo.Describe("gRPC Client", func() {
 
 			spans := reporter.Flush()
 			gomega.Expect(spans).To(gomega.HaveLen(1))
-			gomega.Expect(spans[0].Name).To(gomega.Equal("zipkin.testing.HelloService.Hello"))
-			gomega.Expect(spans[0].Tags).To(gomega.BeEmpty())
+
+			span := spans[0]
+			gomega.Expect(span.Kind).To(gomega.Equal(model.Client))
+			gomega.Expect(span.Name).To(gomega.Equal("zipkin.testing.HelloService.Hello"))
+			gomega.Expect(span.RemoteEndpoint).To(gomega.BeNil())
+			gomega.Expect(span.Tags).To(gomega.BeEmpty())
 		})
 
 		ginkgo.It("propagates trace context", func() {
@@ -94,7 +96,7 @@ var _ = ginkgo.Describe("gRPC Client", func() {
 		})
 	})
 
-	ginkgo.Context("with custom RPCHandler", func() {
+	ginkgo.Context("with remote service name", func() {
 		ginkgo.BeforeEach(func() {
 			var err error
 
@@ -103,21 +105,18 @@ var _ = ginkgo.Describe("gRPC Client", func() {
 				grpc.WithInsecure(),
 				grpc.WithStatsHandler(zipkingrpc.NewClientHandler(
 					tracer,
-					zipkingrpc.WithRPCHandler(func(span zipkin.Span, rpcStats stats.RPCStats) {
-						span.Tag("custom", "tag")
-					}))))
+					zipkingrpc.WithRemoteServiceName("remoteService"))))
 			gomega.Expect(conn, err).ToNot(gomega.BeNil())
 			client = service.NewHelloServiceClient(conn)
 		})
 
-		ginkgo.It("calls custom RPCHandler", func() {
+		ginkgo.It("has remote service name", func() {
 			resp, err := client.Hello(context.Background(), &service.HelloRequest{Payload: "Hello"})
 			gomega.Expect(resp, err).ToNot(gomega.BeNil())
 
 			spans := reporter.Flush()
 			gomega.Expect(spans).To(gomega.HaveLen(1))
-			gomega.Expect(spans[0].Tags).To(gomega.HaveLen(1))
-			gomega.Expect(spans[0].Tags).To(gomega.HaveKeyWithValue("custom", "tag"))
+			gomega.Expect(spans[0].RemoteEndpoint.ServiceName).To(gomega.Equal("remoteService"))
 		})
 	})
 })

middleware/grpc/doc.go

@@ -1,5 +1,3 @@
-// +build go1.9
-
 /*
 Package grpc contains several gRPC handlers which can be used for instrumenting calls with Zipkin.
 */

middleware/grpc/grpc_suite_test.go

@@ -1,5 +1,3 @@
-// +build go1.9
-
 package grpc_test
 
 import (
@@ -15,13 +13,23 @@ import (
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 
+	"github.com/openzipkin/zipkin-go"
+	zipkingrpc "github.com/openzipkin/zipkin-go/middleware/grpc"
 	"github.com/openzipkin/zipkin-go/model"
+	"github.com/openzipkin/zipkin-go/propagation/b3"
 	service "github.com/openzipkin/zipkin-go/proto/testing"
+	"github.com/openzipkin/zipkin-go/reporter/recorder"
 )
 
 var (
+	serverIdGenerator *sequentialIdGenerator
+	serverReporter    *recorder.ReporterRecorder
+
 	server     *grpc.Server
 	serverAddr string
+
+	customServer     *grpc.Server
+	customServerAddr string
 )
 
 func TestGrpc(t *testing.T) {
@@ -30,28 +38,53 @@ func TestGrpc(t *testing.T) {
 }
 
 var _ = ginkgo.BeforeSuite(func() {
+	var err error
+
+	serverReporter = recorder.NewReporter()
+	ep, _ := zipkin.NewEndpoint("grpcServer", "")
+	serverIdGenerator = newSequentialIdGenerator(0x1000000)
+	tracer, err := zipkin.NewTracer(
+		serverReporter, zipkin.WithLocalEndpoint(ep), zipkin.WithIDGenerator(serverIdGenerator), zipkin.WithSharedSpans(false))
+
 	lis, err := net.Listen("tcp", ":0")
 	gomega.Expect(lis, err).ToNot(gomega.BeNil(), "failed to listen to tcp port")
 
-	server = grpc.NewServer()
+	server = grpc.NewServer(grpc.StatsHandler(zipkingrpc.NewServerHandler(tracer)))
 	service.RegisterHelloServiceServer(server, &TestHelloService{})
 	go func() {
 		_ = server.Serve(lis)
 	}()
 	serverAddr = lis.Addr().String()
+
+	customLis, err := net.Listen("tcp", ":0")
+	gomega.Expect(customLis, err).ToNot(gomega.BeNil(), "failed to listen to tcp port")
+
+	tracer, err = zipkin.NewTracer(
+		serverReporter, zipkin.WithLocalEndpoint(ep), zipkin.WithIDGenerator(serverIdGenerator), zipkin.WithSharedSpans(true))
+	customServer = grpc.NewServer(grpc.StatsHandler(zipkingrpc.NewServerHandler(
+		tracer,
+		zipkingrpc.ServerTags(map[string]string{"default": "tag"}))))
+	service.RegisterHelloServiceServer(customServer, &TestHelloService{})
+	go func() {
+		_ = customServer.Serve(customLis)
+	}()
+	customServerAddr = customLis.Addr().String()
 })
 
 var _ = ginkgo.AfterSuite(func() {
 	server.Stop()
+	customServer.Stop()
+	_ = serverReporter.Close()
 })
 
 type sequentialIdGenerator struct {
 	nextTraceId uint64
 	nextSpanId  uint64
+	start       uint64
 }
 
-func newSequentialIdGenerator() *sequentialIdGenerator {
-	return &sequentialIdGenerator{1, 1}
+func newSequentialIdGenerator(start uint64) *sequentialIdGenerator {
+	return &sequentialIdGenerator{start, start, start}
 }
 
 func (g *sequentialIdGenerator) SpanID(traceID model.TraceID) model.ID {
@@ -69,27 +102,43 @@ func (g *sequentialIdGenerator) TraceID() model.TraceID {
 	return id
 }
 
+func (g *sequentialIdGenerator) reset() {
+	g.nextTraceId = g.start
+	g.nextSpanId = g.start
+}
+
 type TestHelloService struct{}
 
 func (s *TestHelloService) Hello(ctx context.Context, req *service.HelloRequest) (*service.HelloResponse, error) {
 	if req.Payload == "fail" {
 		return nil, status.Error(codes.Aborted, "fail")
 	}
 
+	resp := &service.HelloResponse{
+		Payload:     "World",
+		Metadata:    map[string]string{},
+		SpanContext: map[string]string{},
+	}
+
 	md, ok := metadata.FromIncomingContext(ctx)
 	if !ok {
 		return nil, errors.New("could not parse incoming metadata")
 	}
 
-	resp := &service.HelloResponse{
-		Payload:  "World",
-		Metadata: map[string]string{},
-	}
-
 	for k := range md {
 		// Just append the first value for a key for simplicity since we don't use multi-value headers.
 		resp.GetMetadata()[k] = md[k][0]
 	}
 
+	span := zipkin.SpanFromContext(ctx)
+	if span != nil {
+		spanCtx := span.Context()
+		resp.GetSpanContext()[b3.SpanID] = spanCtx.ID.String()
+		resp.GetSpanContext()[b3.TraceID] = spanCtx.TraceID.String()
+		if spanCtx.ParentID != nil {
+			resp.GetSpanContext()[b3.ParentSpanID] = spanCtx.ParentID.String()
+		}
+	}
+
 	return resp, nil
 }

middleware/grpc/server.go

@@ -0,0 +1,76 @@
+package grpc
+
+import (
+	"context"
+	"github.com/openzipkin/zipkin-go"
+	"github.com/openzipkin/zipkin-go/model"
+	"github.com/openzipkin/zipkin-go/propagation/b3"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/stats"
+)
+
+type serverHandler struct {
+	tracer      *zipkin.Tracer
+	defaultTags map[string]string
+}
+
+// A ServerOption can be passed to NewServerHandler to customize the returned handler.
+type ServerOption func(*serverHandler)
+
+// ServerTags adds default Tags to inject into server spans.
+func ServerTags(tags map[string]string) ServerOption {
+	return func(h *serverHandler) {
+		h.defaultTags = tags
+	}
+}
+
+// NewServerHandler returns a stats.Handler which can be used with grpc.WithStatsHandler to add
+// tracing to a gRPC server. The gRPC method name is used as the span name and by default the only
+// tags are the gRPC status code if the call fails. Use ServerTags to add additional tags that
+// should be applied to all spans.
+func NewServerHandler(tracer *zipkin.Tracer, options ...ServerOption) stats.Handler {
+	c := &serverHandler{
+		tracer: tracer,
+	}
+	for _, option := range options {
+		option(c)
+	}
+	return c
+}
+
+// HandleConn exists to satisfy gRPC stats.Handler.
+func (s *serverHandler) HandleConn(ctx context.Context, cs stats.ConnStats) {
+	// no-op
+}
+
+// TagConn exists to satisfy gRPC stats.Handler.
+func (s *serverHandler) TagConn(ctx context.Context, cti *stats.ConnTagInfo) context.Context {
+	// no-op
+	return ctx
+}
+
+// HandleRPC implements per-RPC tracing and stats instrumentation.
+func (s *serverHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
+	handleRPC(ctx, rs)
+}
+
+// TagRPC implements per-RPC context management.
+func (s *serverHandler) TagRPC(ctx context.Context, rti *stats.RPCTagInfo) context.Context {
+	md, ok := metadata.FromIncomingContext(ctx)
+	// In practice, ok never seems to be false but add a defensive check.
+	if !ok {
+		md = metadata.New(nil)
+	}
+
+	name := spanName(rti)
+
+	sc := s.tracer.Extract(b3.ExtractGRPC(&md))
+
+	span := s.tracer.StartSpan(name, zipkin.Kind(model.Server), zipkin.Parent(sc), zipkin.RemoteEndpoint(remoteEndpointFromContext(ctx, "")))
+
+	for k, v := range s.defaultTags {
+		span.Tag(k, v)
+	}
+
+	return zipkin.NewContext(ctx, span)
+}