trace: Add NewContext and FromContext functions.

These will be the standard way for inserting a Trace into a Context,
and retrieving it later. NewContext will typically only be used by
top-level context creators (e.g. the gRPC package).

Change-Id: Ifb6fc81f46a71bb965196f664714d3042a9aea18

1 files changed, 17 insertions, 0 deletions

diff --git a/trace/trace.go b/trace/trace.go
index b469f68..0bd8d1e 100644
--- a/trace/trace.go
+++ b/trace/trace.go
@@ -36,6 +36,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"golang.org/x/net/context"
 	"golang.org/x/net/internal/timeseries"
 )
 
@@ -207,6 +208,22 @@ func lookupBucket(fam string, b int) *traceBucket {
 	return f.Buckets[b]
 }
 
+type contextKeyT string
+
+var contextKey = contextKeyT("golang.org/x/net/trace.Trace")
+
+// NewContext returns a copy of the parent context
+// and associates it with a Trace.
+func NewContext(ctx context.Context, tr Trace) context.Context {
+	return context.WithValue(ctx, contextKey, tr)
+}
+
+// FromContext returns the Trace bound to the context, if any.
+func FromContext(ctx context.Context) (tr Trace, ok bool) {
+	tr, ok = ctx.Value(contextKey).(Trace)
+	return
+}
+
 // Trace represents an active request.
 type Trace interface {
 	// LazyLog adds x to the event log. It will be evaluated each time the