.

go/internal/xtracing/tracetest/example_test.go

@@ -61,6 +61,7 @@ func setupTracing(t *tracetest.T) *tracing.ProbeGroup {
 		t.RxEvent(eventHello(who))
 	})
 	tracing.Unlock()
+	// NOTE pg.Done must be invoked by caller when setup tracing is no longer needed.
 	return pg
 }
 
@@ -86,18 +87,24 @@ func routeEvent(event interface{}) (stream string) {
 	return strings.ToLower(who[:i])
 }
 
+// setupTraceEvents sets up tracing and routing.
+// it is used in the prologue of tests.
+func setupTraceEvents(t *tracetest.T) {
+	// setup tracing to deliver trace events to t.
+	pg := setupTracing(t)
+	t.Cleanup(pg.Done)
+	// tell t to which stream an event should go.
+	t.SetEventRouter(routeEvent)
+}
+
 
 // TestExample demonstrates how to use tracetest to verify concurrent system with 2 threads.
 func TestExample(t *testing.T) {
 	tracetest.Verify(t, tracetestExample)
 }
 func tracetestExample(t *tracetest.T) {
-	// setup tracing to deliver trace events to t.
-	pg := setupTracing(t)
-	defer pg.Done()
-
-	// tell tracetest to which stream an event should go.
-	t.SetEventRouter(routeEvent)
+	// setup tracing to deliver trace events to t and route them appropriately.
+	setupTraceEvents(t)
 
 	// run the workload
 	var wg sync.WaitGroup