Simplify the trace package (#287)

- NewSpan now can create children; (*Span).StartSpan is removed.
- s/StartSpanOptions/StartOptions to reduce verbosity.
- Only connivance functions works against the context object, Span type provides the full API.
- Introspection functions are prefixed with Report* that make them distinguishable.

Author: JBD

examples/trace/helloworld/main.go

@@ -32,19 +32,16 @@ func main() {
 
 	trace.SetDefaultSampler(trace.AlwaysSample())
 
-	span := trace.NewSpan("/foo", trace.StartSpanOptions{})
-
-	ctx = trace.WithSpan(ctx, span)
+	ctx, span := trace.StartSpan(ctx, "/foo")
 	bar(ctx)
-
 	span.End()
 
 	time.Sleep(1 * time.Second) // Wait enough for the exporter to report.
 }
 
 func bar(ctx context.Context) {
-	ctx = trace.StartSpan(ctx, "/bar")
-	defer trace.EndSpan(ctx)
+	ctx, span := trace.StartSpan(ctx, "/bar")
+	defer span.End()
 
 	// Do bar...
 }

examples/trace/jaeger/main.go

@@ -41,16 +41,16 @@ func main() {
 	// For demoing purposes, always sample.
 	trace.SetDefaultSampler(trace.AlwaysSample())
 
-	ctx = trace.StartSpan(ctx, "/foo")
+	ctx, span := trace.StartSpan(ctx, "/foo")
 	bar(ctx)
-	trace.EndSpan(ctx)
+	span.End()
 
 	exporter.Flush()
 }
 
 func bar(ctx context.Context) {
-	ctx = trace.StartSpan(ctx, "/bar")
-	defer trace.EndSpan(ctx)
+	ctx, span := trace.StartSpan(ctx, "/bar")
+	defer span.End()
 
 	// Do bar...
 }

exporter/stackdriver/trace_proto_test.go

@@ -61,50 +61,50 @@ func TestExportTrace(t *testing.T) {
 	trace.RegisterExporter(&te)
 	defer trace.UnregisterExporter(&te)
 
-	ctx := trace.StartSpanWithRemoteParent(context.Background(), "span0",
+	ctx, span0 := trace.StartSpanWithRemoteParent(context.Background(), "span0",
 		trace.SpanContext{
 			TraceID:      traceID,
 			SpanID:       spanID,
 			TraceOptions: 1,
 		},
-		trace.StartSpanOptions{})
+		trace.StartOptions{})
 	{
-		ctx1 := trace.StartSpan(ctx, "span1")
+		ctx1, span1 := trace.StartSpan(ctx, "span1")
 		{
-			ctx2 := trace.StartSpan(ctx1, "span2")
-			trace.AddMessageSendEvent(ctx2, 0x123, 1024, 512)
-			trace.Annotatef(ctx2, nil, "in span%d", 2)
-			trace.Annotate(ctx2, nil, big.NewRat(2, 4).String())
-			trace.SetSpanAttributes(ctx2,
+			_, span2 := trace.StartSpan(ctx1, "span2")
+			span2.AddMessageSendEvent(0x123, 1024, 512)
+			span2.Annotatef(nil, "in span%d", 2)
+			span2.Annotate(nil, big.NewRat(2, 4).String())
+			span2.SetAttributes(
 				trace.StringAttribute{Key: "key1", Value: "value1"},
 				trace.StringAttribute{Key: "key2", Value: "value2"})
-			trace.SetSpanAttributes(ctx2, trace.Int64Attribute{Key: "key1", Value: 100})
-			trace.EndSpan(ctx2)
+			span2.SetAttributes(trace.Int64Attribute{Key: "key1", Value: 100})
+			span2.End()
 		}
 		{
-			ctx3 := trace.StartSpan(ctx1, "span3")
-			trace.Annotate(ctx3, nil, "in span3")
-			trace.AddMessageReceiveEvent(ctx3, 0x456, 2048, 1536)
-			trace.SetSpanStatus(ctx3, trace.Status{Code: int32(codepb.Code_UNAVAILABLE)})
-			trace.EndSpan(ctx3)
+			ctx3, span3 := trace.StartSpan(ctx1, "span3")
+			span3.Annotate(nil, "in span3")
+			span3.AddMessageReceiveEvent(0x456, 2048, 1536)
+			span3.SetStatus(trace.Status{Code: int32(codepb.Code_UNAVAILABLE)})
+			span3.End()
 			{
-				ctx4 := trace.StartSpan(ctx3, "span4")
+				_, span4 := trace.StartSpan(ctx3, "span4")
 				x := 42
 				a1 := []trace.Attribute{trace.StringAttribute{Key: "k1", Value: "v1"}}
 				a2 := []trace.Attribute{trace.StringAttribute{Key: "k2", Value: "v2"}}
 				a3 := []trace.Attribute{trace.StringAttribute{Key: "k3", Value: "v3"}}
 				a4 := map[string]interface{}{"k4": "v4"}
 				r := big.NewRat(2, 4)
-				trace.Annotate(ctx4, a1, r.String())
-				trace.Annotatef(ctx4, a2, "foo %d", x)
-				trace.Annotate(ctx4, a3, "in span4")
-				trace.AddLink(ctx4, trace.Link{TraceID: trace.TraceID{1, 2}, SpanID: trace.SpanID{3}, Type: trace.LinkTypeParent, Attributes: a4})
-				trace.EndSpan(ctx4)
+				span4.Annotate(a1, r.String())
+				span4.Annotatef(a2, "foo %d", x)
+				span4.Annotate(a3, "in span4")
+				span4.AddLink(trace.Link{TraceID: trace.TraceID{1, 2}, SpanID: trace.SpanID{3}, Type: trace.LinkTypeParent, Attributes: a4})
+				span4.End()
 			}
 		}
-		trace.EndSpan(ctx1)
+		span1.End()
 	}
-	trace.EndSpan(ctx)
+	span0.End()
 	if len(te.spans) != 5 {
 		t.Errorf("got %d exported spans, want 5", len(te.spans))
 	}

exporter/stackdriver/trace_test.go

@@ -37,8 +37,8 @@ func TestBundling(t *testing.T) {
 
 	trace.SetDefaultSampler(trace.AlwaysSample())
 	for i := 0; i < 35; i++ {
-		ctx := trace.StartSpan(context.Background(), "span")
-		trace.EndSpan(ctx)
+		_, span := trace.StartSpan(context.Background(), "span")
+		span.End()
 	}
 
 	// Read the first three bundles.

internal/readme/trace.go

@@ -24,7 +24,7 @@ func traceExamples() {
 	ctx := context.Background()
 
 	// START startend
-	ctx = trace.StartSpan(ctx, "your choice of name")
-	defer trace.EndSpan(ctx)
+	ctx, span := trace.StartSpan(ctx, "your choice of name")
+	defer span.End()
 	// END startend
 }

plugin/grpc/grpc_test.go

@@ -37,7 +37,7 @@ func TestNewClientStatsHandler(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	span := trace.NewSpan("/foo", trace.StartSpanOptions{
+	span := trace.NewSpan("/foo", nil, trace.StartOptions{
 		Sampler: trace.AlwaysSample(),
 	})
 	ctx = trace.WithSpan(ctx, span)
@@ -84,7 +84,7 @@ func TestNewServerStatsHandler(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	span := trace.NewSpan("/foo", trace.StartSpanOptions{
+	span := trace.NewSpan("/foo", nil, trace.StartOptions{
 		Sampler: trace.AlwaysSample(),
 	})
 	ctx = trace.WithSpan(ctx, span)

plugin/grpc/grpctrace/grpc.go

@@ -67,7 +67,7 @@ const traceContextKey = "grpc-trace-bin"
 // SpanContext added to the outgoing gRPC metadata.
 func (c *ClientStatsHandler) TagRPC(ctx context.Context, rti *stats.RPCTagInfo) context.Context {
 	name := "Sent" + strings.Replace(rti.FullMethodName, "/", ".", -1)
-	ctx = trace.StartSpanWithOptions(ctx, name, trace.StartSpanOptions{RecordEvents: true, RegisterNameForLocalSpanStore: true})
+	ctx, _ = trace.StartSpanWithOptions(ctx, name, trace.StartOptions{RecordEvents: true, RegisterNameForLocalSpanStore: true})
 	traceContextBinary := propagation.Binary(trace.FromContext(ctx).SpanContext())
 	if len(traceContextBinary) == 0 {
 		return ctx
@@ -88,13 +88,15 @@ func (c *ClientStatsHandler) TagRPC(ctx context.Context, rti *stats.RPCTagInfo)
 func (s *ServerStatsHandler) TagRPC(ctx context.Context, rti *stats.RPCTagInfo) context.Context {
 	md, _ := metadata.FromIncomingContext(ctx)
 	name := "Recv" + strings.Replace(rti.FullMethodName, "/", ".", -1)
-	opt := trace.StartSpanOptions{RecordEvents: true, RegisterNameForLocalSpanStore: true}
+	opt := trace.StartOptions{RecordEvents: true, RegisterNameForLocalSpanStore: true}
 	if s := md[traceContextKey]; len(s) > 0 {
 		if parent, ok := propagation.FromBinary([]byte(s[0])); ok {
-			return trace.StartSpanWithRemoteParent(ctx, name, parent, opt)
+			ctx, _ = trace.StartSpanWithRemoteParent(ctx, name, parent, opt)
+			return ctx
 		}
 	}
-	return trace.StartSpanWithOptions(ctx, name, opt)
+	ctx, _ = trace.StartSpanWithOptions(ctx, name, opt)
+	return ctx
 }
 
 // HandleRPC processes the RPC stats, adding information to the current trace span.
@@ -108,30 +110,31 @@ func (s *ServerStatsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
 }
 
 func handleRPC(ctx context.Context, rs stats.RPCStats) {
+	span := trace.FromContext(ctx)
 	// TODO: compressed and uncompressed sizes are not populated in every message.
 	switch rs := rs.(type) {
 	case *stats.Begin:
-		trace.SetSpanAttributes(ctx,
+		span.SetAttributes(
 			trace.BoolAttribute{Key: "Client", Value: rs.Client},
 			trace.BoolAttribute{Key: "FailFast", Value: rs.FailFast})
 	case *stats.InPayload:
-		trace.AddMessageReceiveEvent(ctx, 0 /* TODO: messageID */, int64(rs.Length), int64(rs.WireLength))
+		span.AddMessageReceiveEvent(0 /* TODO: messageID */, int64(rs.Length), int64(rs.WireLength))
 	case *stats.InHeader:
-		trace.AddMessageReceiveEvent(ctx, 0, int64(rs.WireLength), int64(rs.WireLength))
+		span.AddMessageReceiveEvent(0, int64(rs.WireLength), int64(rs.WireLength))
 	case *stats.InTrailer:
-		trace.AddMessageReceiveEvent(ctx, 0, int64(rs.WireLength), int64(rs.WireLength))
+		span.AddMessageReceiveEvent(0, int64(rs.WireLength), int64(rs.WireLength))
 	case *stats.OutPayload:
-		trace.AddMessageSendEvent(ctx, 0, int64(rs.Length), int64(rs.WireLength))
+		span.AddMessageSendEvent(0, int64(rs.Length), int64(rs.WireLength))
 	case *stats.OutHeader:
-		trace.AddMessageSendEvent(ctx, 0, 0, 0)
+		span.AddMessageSendEvent(0, 0, 0)
 	case *stats.OutTrailer:
-		trace.AddMessageSendEvent(ctx, 0, int64(rs.WireLength), int64(rs.WireLength))
+		span.AddMessageSendEvent(0, int64(rs.WireLength), int64(rs.WireLength))
 	case *stats.End:
 		if rs.Error != nil {
 			code, desc := grpc.Code(rs.Error), grpc.ErrorDesc(rs.Error)
-			trace.SetSpanStatus(ctx, trace.Status{Code: int32(code), Message: desc})
+			span.SetStatus(trace.Status{Code: int32(code), Message: desc})
 		}
-		trace.EndSpan(ctx)
+		span.End()
 	}
 }
 

plugin/http/httptrace/httptrace.go

@@ -16,7 +16,6 @@
 package httptrace
 
 import (
-	"context"
 	"io"
 	"net/http"
 	"net/url"
@@ -51,10 +50,9 @@ func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
 	name := spanNameFromURL("Sent", req.URL)
 	// TODO(jbd): Discuss whether we want to prefix
 	// outgoing requests with Sent.
-	ctx := trace.StartSpan(req.Context(), name)
+	ctx, span := trace.StartSpan(req.Context(), name)
 	req = req.WithContext(ctx)
 
-	span := trace.FromContext(ctx)
 	for _, f := range t.Formats {
 		f.ToRequest(span.SpanContext(), req)
 	}
@@ -63,23 +61,23 @@ func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
 
 	// TODO(jbd): Add status and attributes.
 	if err != nil {
-		trace.EndSpan(ctx)
+		span.End()
 		return resp, err
 	}
 
 	// trace.EndSpan(ctx) will be invoked after
 	// a read from resp.Body returns io.EOF or when
 	// resp.Body.Close() is invoked.
-	resp.Body = &spanEndBody{rc: resp.Body, spanCtx: ctx}
+	resp.Body = &spanEndBody{rc: resp.Body, span: span}
 	return resp, err
 }
 
 // spanEndBody wraps a response.Body and invokes
 // trace.EndSpan on encountering io.EOF on reading
 // the body of the original response.
 type spanEndBody struct {
-	rc      io.ReadCloser
-	spanCtx context.Context
+	rc   io.ReadCloser
+	span *trace.Span
 
 	endSpanOnce sync.Once
 }
@@ -96,7 +94,7 @@ func (seb *spanEndBody) Read(b []byte) (int, error) {
 		seb.endSpan()
 	default:
 		// For all other errors, set the span status
-		trace.SetSpanStatus(seb.spanCtx, trace.Status{
+		seb.span.SetStatus(trace.Status{
 			// Code 2 is the error code for Internal server error.
 			Code:    2,
 			Message: err.Error(),
@@ -108,7 +106,7 @@ func (seb *spanEndBody) Read(b []byte) (int, error) {
 // endSpan invokes trace.EndSpan exactly once
 func (seb *spanEndBody) endSpan() {
 	seb.endSpanOnce.Do(func() {
-		trace.EndSpan(seb.spanCtx)
+		seb.span.End()
 	})
 }
 
@@ -178,12 +176,13 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			break
 		}
 	}
+	var span *trace.Span
 	if ok {
-		ctx = trace.StartSpanWithRemoteParent(ctx, name, sc, trace.StartSpanOptions{})
+		ctx, span = trace.StartSpanWithRemoteParent(ctx, name, sc, trace.StartOptions{})
 	} else {
-		ctx = trace.StartSpan(ctx, name)
+		ctx, span = trace.StartSpan(ctx, name)
 	}
-	defer trace.EndSpan(ctx)
+	defer span.End()
 
 	// TODO(jbd): Add status and attributes.
 	r = r.WithContext(ctx)

plugin/http/httptrace/httptrace_test.go

@@ -71,7 +71,7 @@ func (t testPropagator) ToRequest(sc trace.SpanContext, req *http.Request) {
 }
 
 func TestTransport_RoundTrip(t *testing.T) {
-	parent := trace.NewSpan("parent", trace.StartSpanOptions{})
+	parent := trace.NewSpan("parent", nil, trace.StartOptions{})
 	tests := []struct {
 		name       string
 		parent     *trace.Span
@@ -171,9 +171,9 @@ func TestEndToEnd(t *testing.T) {
 	trace.RegisterExporter(&spans)
 	defer trace.UnregisterExporter(&spans)
 
-	ctx := trace.StartSpanWithOptions(context.Background(),
+	ctx, _ := trace.StartSpanWithOptions(context.Background(),
 		"top-level",
-		trace.StartSpanOptions{
+		trace.StartOptions{
 			RecordEvents: true,
 			Sampler:      trace.AlwaysSample(),
 		})

trace/doc.go

@@ -43,8 +43,8 @@ It is common to want to capture all the activity of a function call in a span. F
 this to work, the function must take a context.Context as a parameter. Add these two
 lines to the top of the function:
 
-    ctx = trace.StartSpan(ctx, "your choice of name")
-    defer trace.EndSpan(ctx)
+    ctx, span := trace.StartSpan(ctx, "your choice of name")
+    defer span.End()
 
 StartSpan will create a new top-level span if the context
 doesn't contain another span, otherwise it will create a child span.