Remove NoTrace from achttp and ocgrpc integrations (#585)

NoTrace was previously used to disable trace integration. This can
now mostly be achieved by customizing the sampler that each integration
takes as an argument.

The one thing that can no longer be achieved after this PR is
disabling trace propagation while still using the plugin. I think
this is really a good thing: trace propagation should not be optional.

plugin/ocgrpc/client.go

@@ -15,6 +15,7 @@
 package ocgrpc
 
 import (
+	"go.opencensus.io/trace"
 	"golang.org/x/net/context"
 
 	"google.golang.org/grpc/stats"
@@ -23,13 +24,12 @@ import (
 // ClientHandler implements a gRPC stats.Handler for recording OpenCensus stats and
 // traces. Use with gRPC clients only.
 type ClientHandler struct {
-	// NoTrace may be set to disable recording OpenCensus Spans around
-	// gRPC methods.
-	NoTrace bool
-
 	// NoStats may be set to disable recording OpenCensus Stats around each
 	// gRPC method.
 	NoStats bool
+
+	// StartOptions allows configuring the StartOptions used to create new spans.
+	StartOptions trace.StartOptions
 }
 
 func (c *ClientHandler) HandleConn(ctx context.Context, cs stats.ConnStats) {
@@ -44,19 +44,15 @@ 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) {
-	if !c.NoTrace {
-		traceHandleRPC(ctx, rs)
-	}
+	traceHandleRPC(ctx, rs)
 	if !c.NoStats {
 		c.statsHandleRPC(ctx, rs)
 	}
 }
 
 // TagRPC implements per-RPC context management.
 func (c *ClientHandler) TagRPC(ctx context.Context, rti *stats.RPCTagInfo) context.Context {
-	if !c.NoTrace {
-		ctx = c.traceTagRPC(ctx, rti)
-	}
+	ctx = c.traceTagRPC(ctx, rti)
 	if !c.NoStats {
 		ctx = c.statsTagRPC(ctx, rti)
 	}

plugin/ocgrpc/client_stats_handler_test.go

@@ -18,6 +18,7 @@ package ocgrpc
 import (
 	"testing"
 
+	"go.opencensus.io/trace"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
@@ -302,7 +303,8 @@ func TestClientDefaultCollections(t *testing.T) {
 			t.Error(err)
 		}
 
-		h := &ClientHandler{NoTrace: true}
+		h := &ClientHandler{}
+		h.StartOptions.Sampler = trace.NeverSample()
 		for _, rpc := range tc.rpcs {
 			mods := []tag.Mutator{}
 			for _, t := range rpc.tags {

plugin/ocgrpc/server.go

@@ -29,11 +29,6 @@ import (
 // present but the SpanContext isn't sampled, then a new trace may be started
 // (as determined by Sampler).
 type ServerHandler struct {
-	// NoTrace may be set to true to disable OpenCensus tracing integration.
-	// If set to true, no trace metadata will be read from inbound RPCs and no
-	// new Spans will be created.
-	NoTrace bool
-
 	// NoStats may be set to true to disable recording OpenCensus stats for RPCs.
 	NoStats bool
 
@@ -72,19 +67,15 @@ func (s *ServerHandler) TagConn(ctx context.Context, cti *stats.ConnTagInfo) con
 
 // HandleRPC implements per-RPC tracing and stats instrumentation.
 func (s *ServerHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
-	if !s.NoTrace {
-		traceHandleRPC(ctx, rs)
-	}
+	traceHandleRPC(ctx, rs)
 	if !s.NoStats {
 		s.statsHandleRPC(ctx, rs)
 	}
 }
 
 // TagRPC implements per-RPC context management.
 func (s *ServerHandler) TagRPC(ctx context.Context, rti *stats.RPCTagInfo) context.Context {
-	if !s.NoTrace {
-		ctx = s.traceTagRPC(ctx, rti)
-	}
+	ctx = s.traceTagRPC(ctx, rti)
 	if !s.NoStats {
 		ctx = s.statsTagRPC(ctx, rti)
 	}

plugin/ocgrpc/server_stats_handler_test.go

@@ -18,6 +18,7 @@ package ocgrpc
 import (
 	"testing"
 
+	"go.opencensus.io/trace"
 	"golang.org/x/net/context"
 
 	"go.opencensus.io/stats/view"
@@ -301,7 +302,8 @@ func TestServerDefaultCollections(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		h := &ServerHandler{NoTrace: true}
+		h := &ServerHandler{}
+		h.StartOptions.Sampler = trace.NeverSample()
 		for _, rpc := range tc.rpcs {
 			mods := []tag.Mutator{}
 			for _, t := range rpc.tags {

plugin/ocgrpc/trace_common.go

@@ -35,11 +35,9 @@ const traceContextKey = "grpc-trace-bin"
 // SpanContext added to the outgoing gRPC metadata.
 func (c *ClientHandler) traceTagRPC(ctx context.Context, rti *stats.RPCTagInfo) context.Context {
 	name := "Sent" + strings.Replace(rti.FullMethodName, "/", ".", -1)
-	ctx, _ = trace.StartSpan(ctx, name)
-	traceContextBinary := propagation.Binary(trace.FromContext(ctx).SpanContext())
-	if len(traceContextBinary) == 0 {
-		return ctx
-	}
+	span := trace.NewSpan(name, trace.FromContext(ctx), c.StartOptions) // span is ended by traceHandleRPC
+	ctx = trace.WithSpan(ctx, span)
+	traceContextBinary := propagation.Binary(span.SpanContext())
 	return metadata.AppendToOutgoingContext(ctx, traceContextKey, string(traceContextBinary))
 }
 

plugin/ocgrpc/trace_common_test.go

@@ -25,6 +25,8 @@ import (
 	"go.opencensus.io/trace"
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/stats"
 )
 
 type testServer struct{}
@@ -87,6 +89,28 @@ func (t *testExporter) ExportSpan(s *trace.SpanData) {
 	go func() { t.ch <- s }()
 }
 
+func TestClientHandler_traceTagRPC(t *testing.T) {
+	ch := &ClientHandler{}
+	ch.StartOptions.Sampler = trace.AlwaysSample()
+	rti := &stats.RPCTagInfo{
+		FullMethodName: "xxx",
+	}
+	ctx := context.Background()
+	ctx = ch.traceTagRPC(ctx, rti)
+
+	span := trace.FromContext(ctx)
+	if span == nil {
+		t.Fatal("expected span, got nil")
+	}
+	if !span.IsRecordingEvents() {
+		t.Errorf("span should be sampled")
+	}
+	md, ok := metadata.FromOutgoingContext(ctx)
+	if !ok || len(md) == 0 || len(md[traceContextKey]) == 0 {
+		t.Fatal("no metadata")
+	}
+}
+
 func TestStreaming(t *testing.T) {
 	trace.SetDefaultSampler(trace.AlwaysSample())
 	te := testExporter{make(chan *trace.SpanData)}

plugin/ochttp/client.go

@@ -35,9 +35,6 @@ type Transport struct {
 	// NoStats may be set to disable recording of stats.
 	NoStats bool
 
-	// NoTrace may be set to disable recording of traces.
-	NoTrace bool
-
 	// Propagation defines how traces are propagated. If unspecified, a default
 	// (currently B3 format) will be used.
 	Propagation propagation.HTTPFormat
@@ -53,16 +50,14 @@ type Transport struct {
 func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
 	rt := t.base()
 	// TODO: remove excessive nesting of http.RoundTrippers here.
-	if !t.NoTrace {
-		format := t.Propagation
-		if format == nil {
-			format = defaultFormat
-		}
-		rt = &traceTransport{
-			base:         rt,
-			format:       format,
-			startOptions: t.StartOptions,
-		}
+	format := t.Propagation
+	if format == nil {
+		format = defaultFormat
+	}
+	rt = &traceTransport{
+		base:         rt,
+		format:       format,
+		startOptions: t.StartOptions,
 	}
 	if !t.NoStats {
 		rt = statsTransport{

plugin/ochttp/client_test.go

@@ -123,16 +123,18 @@ func TestClient(t *testing.T) {
 	}
 }
 
+var noTrace = trace.StartOptions{Sampler: trace.NeverSample()}
+
 func BenchmarkTransportNoInstrumentation(b *testing.B) {
-	benchmarkClientServer(b, &ochttp.Transport{NoStats: true, NoTrace: true})
+	benchmarkClientServer(b, &ochttp.Transport{NoStats: true, StartOptions: noTrace})
 }
 
 func BenchmarkTransportTraceOnly(b *testing.B) {
 	benchmarkClientServer(b, &ochttp.Transport{NoStats: true})
 }
 
 func BenchmarkTransportStatsOnly(b *testing.B) {
-	benchmarkClientServer(b, &ochttp.Transport{NoTrace: true})
+	benchmarkClientServer(b, &ochttp.Transport{StartOptions: noTrace})
 }
 
 func BenchmarkTransportAllInstrumentation(b *testing.B) {
@@ -168,15 +170,15 @@ func benchmarkClientServer(b *testing.B, transport *ochttp.Transport) {
 }
 
 func BenchmarkTransportParallel64NoInstrumentation(b *testing.B) {
-	benchmarkClientServerParallel(b, 64, &ochttp.Transport{NoTrace: true, NoStats: true})
+	benchmarkClientServerParallel(b, 64, &ochttp.Transport{StartOptions: noTrace, NoStats: true})
 }
 
 func BenchmarkTransportParallel64TraceOnly(b *testing.B) {
 	benchmarkClientServerParallel(b, 64, &ochttp.Transport{NoStats: true})
 }
 
 func BenchmarkTransportParallel64StatsOnly(b *testing.B) {
-	benchmarkClientServerParallel(b, 64, &ochttp.Transport{NoTrace: true})
+	benchmarkClientServerParallel(b, 64, &ochttp.Transport{StartOptions: noTrace})
 }
 
 func BenchmarkTransportParallel64AllInstrumentation(b *testing.B) {

plugin/ochttp/server.go

@@ -41,9 +41,6 @@ type Handler struct {
 	// NoStats may be set to disable recording of stats.
 	NoStats bool
 
-	// NoTrace may be set to disable recording of traces.
-	NoTrace bool
-
 	// Propagation defines how traces are propagated. If unspecified,
 	// B3 propagation will be used.
 	Propagation propagation.HTTPFormat
@@ -57,11 +54,9 @@ type Handler struct {
 }
 
 func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	if !h.NoTrace {
-		var end func()
-		r, end = h.startTrace(w, r)
-		defer end()
-	}
+	var end func()
+	r, end = h.startTrace(w, r)
+	defer end()
 	if !h.NoStats {
 		var end func()
 		w, end = h.startStats(w, r)

plugin/ochttp/server_test.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 
 	"go.opencensus.io/stats/view"
+	"go.opencensus.io/trace"
 )
 
 func httpHandler(statusCode, respSize int) http.Handler {
@@ -53,9 +54,9 @@ func TestHandlerStatsCollection(t *testing.T) {
 			r := httptest.NewRequest(test.method, test.target, body)
 			w := httptest.NewRecorder()
 			h := &Handler{
-				NoTrace: true,
 				Handler: httpHandler(test.statusCode, test.respSize),
 			}
+			h.StartOptions.Sampler = trace.NeverSample()
 
 			for i := 0; i < test.count; i++ {
 				h.ServeHTTP(w, r)