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")
+ }
+}