Ensure spans created by httptrace client tracer reflect operation structure (#618)

* Ensure spans created by httptrace client tracer reflect operation structure

* Cleanup (clientTracer).start based on PR feedback

* Ensure a span is recorded even if end() is called before start()

* Ensure start attributes for spans started by (clientTracer).end() are recorded

Co-authored-by: Rahul Patel <rahulpa@google.com>

plugin/httptrace/clienttrace.go

@@ -37,12 +37,27 @@ var (
 	HTTPLocalAddr  = key.New("http.local")
 )
 
+var (
+	hookMap = map[string]string{
+		"http.dns":     "http.getconn",
+		"http.connect": "http.getconn",
+		"http.tls":     "http.getconn",
+	}
+)
+
+func parentHook(hook string) string {
+	if strings.HasPrefix(hook, "http.connect") {
+		return hookMap["http.connect"]
+	}
+	return hookMap[hook]
+}
+
 type clientTracer struct {
 	context.Context
 
 	tr trace.Tracer
 
-	activeHooks map[string]trace.Span
+	activeHooks map[string]context.Context
 	root        trace.Span
 	mtx         sync.Mutex
 }
@@ -50,7 +65,7 @@ type clientTracer struct {
 func NewClientTrace(ctx context.Context) *httptrace.ClientTrace {
 	ct := &clientTracer{
 		Context:     ctx,
-		activeHooks: make(map[string]trace.Span),
+		activeHooks: make(map[string]context.Context),
 	}
 
 	ct.tr = global.Tracer("go.opentelemetry.io/otel/plugin/httptrace")
@@ -76,25 +91,30 @@ func NewClientTrace(ctx context.Context) *httptrace.ClientTrace {
 }
 
 func (ct *clientTracer) start(hook, spanName string, attrs ...core.KeyValue) {
-	_, sp := ct.tr.Start(ct.Context, spanName, trace.WithAttributes(attrs...), trace.WithSpanKind(trace.SpanKindClient))
 	ct.mtx.Lock()
 	defer ct.mtx.Unlock()
+
+	if hookCtx, found := ct.activeHooks[hook]; !found {
+		var sp trace.Span
+		ct.activeHooks[hook], sp = ct.tr.Start(ct.getParentContext(hook), spanName, trace.WithAttributes(attrs...), trace.WithSpanKind(trace.SpanKindClient))
 		if ct.root == nil {
 			ct.root = sp
 		}
-	if _, ok := ct.activeHooks[hook]; ok {
-		// end was called before start is handled.
-		sp.End()
-		delete(ct.activeHooks, hook)
 	} else {
-		ct.activeHooks[hook] = sp
+		// end was called before start finished, add the start attributes and end the span here
+		span := trace.SpanFromContext(hookCtx)
+		span.SetAttributes(attrs...)
+		span.End()
+
+		delete(ct.activeHooks, hook)
 	}
 }
 
 func (ct *clientTracer) end(hook string, err error, attrs ...core.KeyValue) {
 	ct.mtx.Lock()
 	defer ct.mtx.Unlock()
-	if span, ok := ct.activeHooks[hook]; ok {
+	if ctx, ok := ct.activeHooks[hook]; ok {
+		span := trace.SpanFromContext(ctx)
 		if err != nil {
 			span.SetStatus(codes.Unknown, err.Error())
 		}
@@ -103,14 +123,31 @@ func (ct *clientTracer) end(hook string, err error, attrs ...core.KeyValue) {
 		delete(ct.activeHooks, hook)
 	} else {
 		// start is not finished before end is called.
-		ct.activeHooks[hook] = trace.NoopSpan{}
+		// Start a span here with the ending attributes that will be finished when start finishes.
+		// Yes, it's backwards. v0v
+		ctx, span := ct.tr.Start(ct.getParentContext(hook), hook, trace.WithAttributes(attrs...), trace.WithSpanKind(trace.SpanKindClient))
+		if err != nil {
+			span.SetStatus(codes.Unknown, err.Error())
 		}
+		ct.activeHooks[hook] = ctx
+	}
+}
+
+func (ct *clientTracer) getParentContext(hook string) context.Context {
+	ctx, ok := ct.activeHooks[parentHook(hook)]
+	if !ok {
+		return ct.Context
+	}
+	return ctx
 }
 
 func (ct *clientTracer) span(hook string) trace.Span {
 	ct.mtx.Lock()
 	defer ct.mtx.Unlock()
-	return ct.activeHooks[hook]
+	if ctx, ok := ct.activeHooks[hook]; ok {
+		return trace.SpanFromContext(ctx)
+	}
+	return nil
 }
 
 func (ct *clientTracer) getConn(host string) {

plugin/httptrace/clienttrace_test.go

@@ -17,6 +17,7 @@ import (
 	"context"
 	"net/http"
 	"net/http/httptest"
+	nhtrace "net/http/httptrace"
 	"sync"
 	"testing"
 
@@ -86,13 +87,28 @@ func TestHTTPRequestWithClientTrace(t *testing.T) {
 		panic("unexpected error in http request: " + err.Error())
 	}
 
+	getSpan := func(name string) *export.SpanData {
+		spans, ok := exp.spanMap[name]
+		if !ok {
+			t.Fatalf("no spans found with the name %s, %v", name, exp.spanMap)
+		}
+
+		if len(spans) != 1 {
+			t.Fatalf("Expected exactly one span for %s but found %d", name, len(spans))
+		}
+
+		return spans[0]
+	}
+
 	testLen := []struct {
 		name       string
 		attributes []core.KeyValue
+		parent     string
 	}{
 		{
 			name:       "http.connect",
 			attributes: []core.KeyValue{key.String("http.remote", address.String())},
+			parent:     "http.getconn",
 		},
 		{
 			name: "http.getconn",
@@ -100,27 +116,30 @@ func TestHTTPRequestWithClientTrace(t *testing.T) {
 				key.String("http.remote", address.String()),
 				key.String("http.host", address.String()),
 			},
+			parent: "test",
 		},
 		{
 			name:   "http.receive",
+			parent: "test",
 		},
 		{
 			name:   "http.send",
+			parent: "test",
 		},
 		{
 			name: "test",
 		},
 	}
 	for _, tl := range testLen {
-		spans, ok := exp.spanMap[tl.name]
-		if !ok {
-			t.Fatalf("no spans found with the name %s, %v", tl.name, exp.spanMap)
-		}
+		span := getSpan(tl.name)
 
-		if len(spans) != 1 {
-			t.Fatalf("Expected exactly one span for %s but found %d", tl.name, len(spans))
+		if tl.parent != "" {
+			parentSpan := getSpan(tl.parent)
+
+			if span.ParentSpanID != parentSpan.SpanContext.SpanID {
+				t.Fatalf("[span %s] does not have expected parent span %s", tl.name, tl.parent)
+			}
 		}
-		span := spans[0]
 
 		actualAttrs := make(map[core.Key]string)
 		for _, attr := range span.Attributes {
@@ -263,3 +282,36 @@ func TestConcurrentConnectionStart(t *testing.T) {
 		})
 	}
 }
+
+func TestEndBeforeStartCreatesSpan(t *testing.T) {
+	exp := &testExporter{
+		spanMap: make(map[string][]*export.SpanData),
+	}
+	tp, _ := sdktrace.NewProvider(sdktrace.WithSyncer(exp), sdktrace.WithConfig(sdktrace.Config{DefaultSampler: sdktrace.AlwaysSample()}))
+	global.SetTraceProvider(tp)
+
+	tr := tp.Tracer("httptrace/client")
+	ctx, span := tr.Start(context.Background(), "test")
+	defer span.End()
+
+	ct := httptrace.NewClientTrace(ctx)
+
+	ct.DNSDone(nhtrace.DNSDoneInfo{})
+	ct.DNSStart(nhtrace.DNSStartInfo{Host: "example.com"})
+
+	getSpan := func(name string) *export.SpanData {
+		spans, ok := exp.spanMap[name]
+		if !ok {
+			t.Fatalf("no spans found with the name %s, %v", name, exp.spanMap)
+		}
+
+		if len(spans) != 1 {
+			t.Fatalf("Expected exactly one span for %s but found %d", name, len(spans))
+		}
+
+		return spans[0]
+	}
+
+	getSpan("http.dns")
+
+}