event: propagate trace context

TraceBuilder.Start returns the context that the trace handler returns.
It also makes sure that Builder.End is called with that context.

Change-Id: I5c6a02aab08f89147f5b40b4f3556a682e2f8697
Reviewed-on: https://go-review.googlesource.com/c/exp/+/320089
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Ian Cottrell <iancottrell@google.com>
diff --git a/event/builder.go b/event/builder.go
index e636acc..03486fa 100644
--- a/event/builder.go
+++ b/event/builder.go
@@ -257,7 +257,8 @@
 		// and now deliver the start event
 		b.data.Event.Message = name
 		ctx = newContext(ctx, b.data.exporter, b.data.Event.ID)
-		b.data.exporter.trace.Start(ctx, &b.data.Event)
+		ctx = b.data.exporter.trace.Start(ctx, &b.data.Event)
+		eb.data.ctx = ctx
 	}
 	b.done()
 	return ctx, end
diff --git a/event/builder_test.go b/event/builder_test.go
index 6cf7d69..ddab645 100644
--- a/event/builder_test.go
+++ b/event/builder_test.go
@@ -61,3 +61,30 @@
 func valueEqual(l1, l2 event.Value) bool {
 	return fmt.Sprint(l1) == fmt.Sprint(l2)
 }
+
+func TestTraceBuilder(t *testing.T) {
+	// Verify that the context returned from the handler is also returned from Start,
+	// and is the context passed to End.
+	ctx := event.WithExporter(context.Background(), event.NewExporter(&testTraceHandler{t}))
+	ctx, end := event.Trace(ctx).Start("s")
+	val := ctx.Value("x")
+	if val != 1 {
+		t.Fatal("context not returned from Start")
+	}
+	end()
+}
+
+type testTraceHandler struct {
+	t *testing.T
+}
+
+func (*testTraceHandler) Start(ctx context.Context, _ *event.Event) context.Context {
+	return context.WithValue(ctx, "x", 1)
+}
+
+func (t *testTraceHandler) End(ctx context.Context, _ *event.Event) {
+	val := ctx.Value("x")
+	if val != 1 {
+		t.t.Fatal("Start context not passed to End")
+	}
+}