Add if the trace was sampled to the ctxTags

tracing/opentracing/id_extract.go

@@ -1,50 +1,67 @@
-// Copyright 2017 Michal Witkowski. All Rights Reserved.
-// See LICENSE for licensing terms.
-
 package grpc_opentracing
 
 import (
 	"strings"
 
-	"github.com/grpc-ecosystem/go-grpc-middleware/tags"
-	"github.com/opentracing/opentracing-go"
+	grpc_ctxtags "github.com/grpc-ecosystem/go-grpc-middleware/tags"
+	opentracing "github.com/opentracing/opentracing-go"
 	"google.golang.org/grpc/grpclog"
 )
 
 const (
-	TagTraceId = "trace.traceid"
-	TagSpanId  = "trace.spanid"
+	TagTraceId           = "trace.traceid"
+	TagSpanId            = "trace.spanid"
+	TagSampled           = "trace.sampled"
+	jaegerNotSampledFlag = "0"
 )
 
-// hackyInjectOpentracingIdsToTags writes the given context to the ctxtags.
+// injectOpentracingIdsToTags writes trace data to ctxtags.
 // This is done in an incredibly hacky way, because the public-facing interface of opentracing doesn't give access to
 // the TraceId and SpanId of the SpanContext. Only the Tracer's Inject/Extract methods know what these are.
 // Most tracers have them encoded as keys with 'traceid' and 'spanid':
 // https://github.com/openzipkin/zipkin-go-opentracing/blob/594640b9ef7e5c994e8d9499359d693c032d738c/propagation_ot.go#L29
 // https://github.com/opentracing/basictracer-go/blob/1b32af207119a14b1b231d451df3ed04a72efebf/propagation_ot.go#L26
 // Jaeger from Uber use one-key schema with next format '{trace-id}:{span-id}:{parent-span-id}:{flags}'
 // https://www.jaegertracing.io/docs/client-libraries/#trace-span-identity
-func hackyInjectOpentracingIdsToTags(span opentracing.Span, tags grpc_ctxtags.Tags) {
-	if err := span.Tracer().Inject(span.Context(), opentracing.HTTPHeaders, &hackyTagsCarrier{tags}); err != nil {
-		grpclog.Printf("grpc_opentracing: failed extracting trace info into ctx %v", err)
+func injectOpentracingIdsToTags(span opentracing.Span, tags grpc_ctxtags.Tags) {
+	if err := span.Tracer().Inject(span.Context(), opentracing.HTTPHeaders, &tagsCarrier{tags}); err != nil {
+		grpclog.Infof("grpc_opentracing: failed extracting trace info into ctx %v", err)
 	}
 }
 
-// hackyTagsCarrier is a really hacky way of
-type hackyTagsCarrier struct {
+// tagsCarrier is a really hacky way of
+type tagsCarrier struct {
 	grpc_ctxtags.Tags
 }
 
-func (t *hackyTagsCarrier) Set(key, val string) {
-	if strings.Contains(key, "traceid") || strings.Contains(strings.ToLower(key), "traceid") {
+func (t *tagsCarrier) Set(key, val string) {
+	key = strings.ToLower(key)
+	if strings.Contains(key, "traceid") {
 		t.Tags.Set(TagTraceId, val) // this will most likely be base-16 (hex) encoded
-	} else if (strings.Contains(key, "spanid") && !strings.Contains(key, "parent")) || (strings.Contains(strings.ToLower(key), "spanid") && !strings.Contains(strings.ToLower(key), "parent")) {
+	}
+
+	if strings.Contains(key, "spanid") && !strings.Contains(strings.ToLower(key), "parent") {
 		t.Tags.Set(TagSpanId, val) // this will most likely be base-16 (hex) encoded
-	} else if key == "uber-trace-id" {
+	}
+
+	if strings.Contains(key, "sampled") {
+		switch val {
+		case "true", "false":
+			t.Tags.Set(TagSampled, val)
+		}
+	}
+
+	if key == "uber-trace-id" {
 		parts := strings.Split(val, ":")
-		if len(parts) >= 2 {
+		if len(parts) == 4 {
 			t.Tags.Set(TagTraceId, parts[0])
 			t.Tags.Set(TagSpanId, parts[1])
+
+			if parts[3] != jaegerNotSampledFlag {
+				t.Tags.Set(TagSampled, "true")
+			} else {
+				t.Tags.Set(TagSampled, "false")
+			}
 		}
 	}
 }

tracing/opentracing/interceptors_test.go

@@ -4,7 +4,6 @@
 package grpc_opentracing_test
 
 import (
-	"encoding/json"
 	"errors"
 	"strconv"
 	"strings"
@@ -17,10 +16,10 @@ import (
 
 	"github.com/grpc-ecosystem/go-grpc-middleware"
 	"github.com/grpc-ecosystem/go-grpc-middleware/tags"
-	grpc_testing "github.com/grpc-ecosystem/go-grpc-middleware/testing"
+	"github.com/grpc-ecosystem/go-grpc-middleware/testing"
 	pb_testproto "github.com/grpc-ecosystem/go-grpc-middleware/testing/testproto"
-	grpc_opentracing "github.com/grpc-ecosystem/go-grpc-middleware/tracing/opentracing"
-	opentracing "github.com/opentracing/opentracing-go"
+	"github.com/grpc-ecosystem/go-grpc-middleware/tracing/opentracing"
+	"github.com/opentracing/opentracing-go"
 	"github.com/opentracing/opentracing-go/mocktracer"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -36,20 +35,6 @@ var (
 	fakeInboundSpanId  = 999
 )
 
-func tagsToJson(value map[string]interface{}) string {
-	str, _ := json.Marshal(value)
-	return string(str)
-}
-
-func tagsFromJson(t *testing.T, jstring string) map[string]interface{} {
-	var msgMapTemplate interface{}
-	err := json.Unmarshal([]byte(jstring), &msgMapTemplate)
-	if err != nil {
-		t.Fatalf("failed unmarshaling tags from response %v", err)
-	}
-	return msgMapTemplate.(map[string]interface{})
-}
-
 type tracingAssertService struct {
 	pb_testproto.TestServiceServer
 	T *testing.T
@@ -59,7 +44,9 @@ func (s *tracingAssertService) Ping(ctx context.Context, ping *pb_testproto.Ping
 	assert.NotNil(s.T, opentracing.SpanFromContext(ctx), "handlers must have the spancontext in their context, otherwise propagation will fail")
 	tags := grpc_ctxtags.Extract(ctx)
 	assert.True(s.T, tags.Has(grpc_opentracing.TagTraceId), "tags must contain traceid")
-	assert.True(s.T, tags.Has(grpc_opentracing.TagSpanId), "tags must contain traceid")
+	assert.True(s.T, tags.Has(grpc_opentracing.TagSpanId), "tags must contain spanid")
+	assert.True(s.T, tags.Has(grpc_opentracing.TagSampled), "tags must contain sampled")
+	assert.Equal(s.T, tags.Values()[grpc_opentracing.TagSampled], "true", "sampled must be set to true")
 	return s.TestServiceServer.Ping(ctx, ping)
 }
 
@@ -72,12 +59,19 @@ func (s *tracingAssertService) PingList(ping *pb_testproto.PingRequest, stream p
 	assert.NotNil(s.T, opentracing.SpanFromContext(stream.Context()), "handlers must have the spancontext in their context, otherwise propagation will fail")
 	tags := grpc_ctxtags.Extract(stream.Context())
 	assert.True(s.T, tags.Has(grpc_opentracing.TagTraceId), "tags must contain traceid")
-	assert.True(s.T, tags.Has(grpc_opentracing.TagSpanId), "tags must contain traceid")
+	assert.True(s.T, tags.Has(grpc_opentracing.TagSpanId), "tags must contain spanid")
+	assert.True(s.T, tags.Has(grpc_opentracing.TagSampled), "tags must contain sampled")
+	assert.Equal(s.T, tags.Values()[grpc_opentracing.TagSampled], "true", "sampled must be set to true")
 	return s.TestServiceServer.PingList(ping, stream)
 }
 
 func (s *tracingAssertService) PingEmpty(ctx context.Context, empty *pb_testproto.Empty) (*pb_testproto.PingResponse, error) {
 	assert.NotNil(s.T, opentracing.SpanFromContext(ctx), "handlers must have the spancontext in their context, otherwise propagation will fail")
+	tags := grpc_ctxtags.Extract(ctx)
+	assert.True(s.T, tags.Has(grpc_opentracing.TagTraceId), "tags must contain traceid")
+	assert.True(s.T, tags.Has(grpc_opentracing.TagSpanId), "tags must contain spanid")
+	assert.True(s.T, tags.Has(grpc_opentracing.TagSampled), "tags must contain sampled")
+	assert.Equal(s.T, tags.Values()[grpc_opentracing.TagSampled], "false", "sampled must be set to false")
 	return s.TestServiceServer.PingEmpty(ctx, empty)
 }
 
@@ -135,12 +129,17 @@ func (s *OpentracingSuite) SetupTest() {
 	s.mockTracer.Reset()
 }
 
-func (s *OpentracingSuite) createContextFromFakeHttpRequestParent(ctx context.Context) context.Context {
+func (s *OpentracingSuite) createContextFromFakeHttpRequestParent(ctx context.Context, sampled bool) context.Context {
+	jFlag := 0
+	if sampled {
+		jFlag = 1
+	}
+
 	hdr := http.Header{}
-	hdr.Set("uber-trace-id", fmt.Sprintf("%d:%d:%d:1", fakeInboundTraceId, fakeInboundSpanId, fakeInboundSpanId))
+	hdr.Set("uber-trace-id", fmt.Sprintf("%d:%d:%d:%d", fakeInboundTraceId, fakeInboundSpanId, fakeInboundSpanId, jFlag))
 	hdr.Set("mockpfx-ids-traceid", fmt.Sprint(fakeInboundTraceId))
 	hdr.Set("mockpfx-ids-spanid", fmt.Sprint(fakeInboundSpanId))
-	hdr.Set("mockpfx-ids-sampled", fmt.Sprint(true))
+	hdr.Set("mockpfx-ids-sampled", fmt.Sprint(sampled))
 
 	parentSpanContext, err := s.mockTracer.Extract(opentracing.HTTPHeaders, opentracing.HTTPHeadersCarrier(hdr))
 	require.NoError(s.T(), err, "parsing a fake HTTP request headers shouldn't fail, ever")
@@ -179,7 +178,7 @@ func (s *OpentracingSuite) assertTracesCreated(methodName string) (clientSpan *m
 }
 
 func (s *OpentracingSuite) TestPing_PropagatesTraces() {
-	ctx := s.createContextFromFakeHttpRequestParent(s.SimpleCtx())
+	ctx := s.createContextFromFakeHttpRequestParent(s.SimpleCtx(), true)
 	_, err := s.Client.Ping(ctx, goodPing)
 	require.NoError(s.T(), err, "there must be not be an on a successful call")
 	s.assertTracesCreated("/mwitkow.testproto.TestService/Ping")
@@ -188,7 +187,7 @@ func (s *OpentracingSuite) TestPing_PropagatesTraces() {
 func (s *OpentracingSuite) TestPing_ClientContextTags() {
 	const name = "opentracing.custom"
 	ctx := grpc_opentracing.ClientAddContextTags(
-		s.createContextFromFakeHttpRequestParent(s.SimpleCtx()),
+		s.createContextFromFakeHttpRequestParent(s.SimpleCtx(), true),
 		opentracing.Tags{name: ""},
 	)
 
@@ -206,7 +205,7 @@ func (s *OpentracingSuite) TestPing_ClientContextTags() {
 }
 
 func (s *OpentracingSuite) TestPingList_PropagatesTraces() {
-	ctx := s.createContextFromFakeHttpRequestParent(s.SimpleCtx())
+	ctx := s.createContextFromFakeHttpRequestParent(s.SimpleCtx(), true)
 	stream, err := s.Client.PingList(ctx, goodPing)
 	require.NoError(s.T(), err, "should not fail on establishing the stream")
 	for {
@@ -220,7 +219,7 @@ func (s *OpentracingSuite) TestPingList_PropagatesTraces() {
 }
 
 func (s *OpentracingSuite) TestPingError_PropagatesTraces() {
-	ctx := s.createContextFromFakeHttpRequestParent(s.SimpleCtx())
+	ctx := s.createContextFromFakeHttpRequestParent(s.SimpleCtx(), true)
 	erroringPing := &pb_testproto.PingRequest{Value: "something", ErrorCodeReturned: uint32(codes.OutOfRange)}
 	_, err := s.Client.PingError(ctx, erroringPing)
 	require.Error(s.T(), err, "there must be an error returned here")
@@ -229,6 +228,12 @@ func (s *OpentracingSuite) TestPingError_PropagatesTraces() {
 	assert.Equal(s.T(), true, serverSpan.Tag("error"), "server span needs to be marked as an error")
 }
 
+func (s *OpentracingSuite) TestPingEmpty_NotSampleTraces() {
+	ctx := s.createContextFromFakeHttpRequestParent(s.SimpleCtx(), false)
+	_, err := s.Client.PingEmpty(ctx, &pb_testproto.Empty{})
+	require.NoError(s.T(), err, "there must be not be an on a successful call")
+}
+
 type jaegerFormatInjector struct{}
 
 func (jaegerFormatInjector) Inject(ctx mocktracer.MockSpanContext, carrier interface{}) error {

tracing/opentracing/server_interceptors.go

@@ -62,7 +62,8 @@ func newServerSpanFromInbound(ctx context.Context, tracer opentracing.Tracer, fu
 		ext.RPCServerOption(parentSpanContext),
 		grpcTag,
 	)
-	hackyInjectOpentracingIdsToTags(serverSpan, grpc_ctxtags.Extract(ctx))
+
+	injectOpentracingIdsToTags(serverSpan, grpc_ctxtags.Extract(ctx))
 	return opentracing.ContextWithSpan(ctx, serverSpan), serverSpan
 }