From 16ba2b6308d7b6e6ae7fac01564589f540856c4d Mon Sep 17 00:00:00 2001
From: Krzesimir Nowak <qdlacz@gmail.com>
Date: Fri, 13 Sep 2019 00:27:53 +0200
Subject: [PATCH] Add an extension to Tracer interface for custom go context
 creation (#220)

* Add an extension interface for Tracer implementations

The new TracerContextWithSpanExtension interface provides a way to
hook into the ContextWithSpan function, so the implementation can put
some extra information to the context.

The opentracing to opentelemetry bridge needs the context to set the
current opentelemetry span, so the opentelemetry API in the layer
below the one using opentracing can still get the right parent span.

* Call the hook in the ContextWithSpan function if possible

So the implementation can still affect the way the go context is set
up.

* Add a test for the TracerContextWithSpanExtension interface
---
 ext.go            | 24 ++++++++++++++++++++++++
 gocontext.go      |  3 +++
 gocontext_test.go | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+)
 create mode 100644 ext.go

diff --git a/ext.go b/ext.go
new file mode 100644
index 0000000..e11977e
--- /dev/null
+++ b/ext.go
@@ -0,0 +1,24 @@
+package opentracing
+
+import (
+	"context"
+)
+
+// TracerContextWithSpanExtension is an extension interface that the
+// implementation of the Tracer interface may want to implement. It
+// allows to have some control over the go context when the
+// ContextWithSpan is invoked.
+//
+// The primary purpose of this extension are adapters from opentracing
+// API to some other tracing API.
+type TracerContextWithSpanExtension interface {
+	// ContextWithSpanHook gets called by the ContextWithSpan
+	// function, when the Tracer implementation also implements
+	// this interface. It allows to put extra information into the
+	// context and make it available to the callers of the
+	// ContextWithSpan.
+	//
+	// This hook is invoked before the ContextWithSpan function
+	// actually puts the span into the context.
+	ContextWithSpanHook(ctx context.Context, span Span) context.Context
+}
diff --git a/gocontext.go b/gocontext.go
index 08c00c0..8865e75 100644
--- a/gocontext.go
+++ b/gocontext.go
@@ -9,6 +9,9 @@ var activeSpanKey = contextKey{}
 // ContextWithSpan returns a new `context.Context` that holds a reference to
 // `span`'s SpanContext.
 func ContextWithSpan(ctx context.Context, span Span) context.Context {
+	if tracerWithHook, ok := span.Tracer().(TracerContextWithSpanExtension); ok {
+		ctx = tracerWithHook.ContextWithSpanHook(ctx, span)
+	}
 	return context.WithValue(ctx, activeSpanKey, span)
 }
 
diff --git a/gocontext_test.go b/gocontext_test.go
index 65ca1c4..a5b1aed 100644
--- a/gocontext_test.go
+++ b/gocontext_test.go
@@ -29,6 +29,41 @@ func TestContextWithSpan(t *testing.T) {
 	}
 }
 
+type noopExtTracer struct {
+	NoopTracer
+}
+
+type noopExtTracerCtxType struct{}
+
+func (noopExtTracer) ContextWithSpanHook(ctx context.Context, span Span) context.Context {
+	return context.WithValue(ctx, noopExtTracerCtxType{}, noopExtTracerCtxType{})
+}
+
+var _ Tracer = noopExtTracer{}
+var _ TracerContextWithSpanExtension = noopExtTracer{}
+
+type noopExtSpan struct {
+	noopSpan
+}
+
+func (noopExtSpan) Tracer() Tracer {
+	return noopExtTracer{}
+}
+
+var _ Span = noopExtSpan{}
+
+func TestContextWithSpanWithExtension(t *testing.T) {
+	span := &noopExtSpan{}
+	ctx := ContextWithSpan(context.Background(), span)
+	span2 := SpanFromContext(ctx)
+	if span != span2 {
+		t.Errorf("Not the same span returned from context, expected=%+v, actual=%+v", span, span2)
+	}
+	if _, ok := ctx.Value(noopExtTracerCtxType{}).(noopExtTracerCtxType); !ok {
+		t.Error("ContextWithSpanHook was not called")
+	}
+}
+
 func TestStartSpanFromContext(t *testing.T) {
 	testTracer := testTracer{}