Add Configuration support for custom Injector and Extractor (jaegertr…

…acing#263)

Signed-off-by: Martin Liu <[email protected]>

config/config.go

@@ -175,6 +175,14 @@ func (c Configuration) New(
 		tracerOptions = append(tracerOptions, jaeger.TracerOptions.ContribObserver(cobs))
 	}
 
+	for format, injector := range opts.injectors {
+		tracerOptions = append(tracerOptions, jaeger.TracerOptions.Injector(format, injector))
+	}
+
+	for format, extractor := range opts.extractors {
+		tracerOptions = append(tracerOptions, jaeger.TracerOptions.Extractor(format, extractor))
+	}
+
 	if c.BaggageRestrictions != nil {
 		mgr := remote.NewRestrictionManager(
 			serviceName,
@@ -193,7 +201,8 @@ func (c Configuration) New(
 		serviceName,
 		sampler,
 		reporter,
-		tracerOptions...)
+		tracerOptions...,
+	)
 
 	return tracer, closer, nil
 }

config/config_test.go

@@ -257,3 +257,32 @@ func TestConfigWithGen128Bit(t *testing.T) {
 	require.True(t, traceID.High != 0)
 	require.True(t, traceID.Low != 0)
 }
+
+func TestConfigWithInjector(t *testing.T) {
+	c := Configuration{}
+	tracer, closer, err := c.New("test", Injector("custom.format", fakeInjector{}))
+	require.NoError(t, err)
+	defer closer.Close()
+
+	span := tracer.StartSpan("test")
+	defer span.Finish()
+
+	err = tracer.Inject(span.Context(), "unknown.format", nil)
+	require.Error(t, err)
+
+	err = tracer.Inject(span.Context(), "custom.format", nil)
+	require.NoError(t, err)
+}
+
+func TestConfigWithExtractor(t *testing.T) {
+	c := Configuration{}
+	tracer, closer, err := c.New("test", Extractor("custom.format", fakeExtractor{}))
+	require.NoError(t, err)
+	defer closer.Close()
+
+	_, err = tracer.Extract("unknown.format", nil)
+	require.Error(t, err)
+
+	_, err = tracer.Extract("custom.format", nil)
+	require.NoError(t, err)
+}

config/options.go

@@ -34,6 +34,8 @@ type Options struct {
 	gen128Bit           bool
 	zipkinSharedRPCSpan bool
 	tags                []opentracing.Tag
+	injectors           map[interface{}]jaeger.Injector
+	extractors          map[interface{}]jaeger.Extractor
 }
 
 // Metrics creates an Option that initializes Metrics in the tracer,
@@ -98,8 +100,25 @@ func Tag(key string, value interface{}) Option {
 	}
 }
 
+// Injector registers an Injector with the given format.
+func Injector(format interface{}, injector jaeger.Injector) Option {
+	return func(c *Options) {
+		c.injectors[format] = injector
+	}
+}
+
+// Extractor registers an Extractor with the given format.
+func Extractor(format interface{}, extractor jaeger.Extractor) Option {
+	return func(c *Options) {
+		c.extractors[format] = extractor
+	}
+}
+
 func applyOptions(options ...Option) Options {
-	opts := Options{}
+	opts := Options{
+		injectors:  make(map[interface{}]jaeger.Injector),
+		extractors: make(map[interface{}]jaeger.Extractor),
+	}
 	for _, option := range options {
 		option(&opts)
 	}

config/options_test.go

@@ -61,12 +61,24 @@ func TestApplyOptionsDefaults(t *testing.T) {
 
 type fakeObserver struct{}
 
-func (o fakeObserver) OnStartSpan(operationName string, options opentracing.StartSpanOptions) jaeger.SpanObserver {
+func (fakeObserver) OnStartSpan(operationName string, options opentracing.StartSpanOptions) jaeger.SpanObserver {
 	return nil
 }
 
 type fakeContribObserver struct{}
 
-func (o fakeContribObserver) OnStartSpan(span opentracing.Span, operationName string, options opentracing.StartSpanOptions) (jaeger.ContribSpanObserver, bool) {
+func (fakeContribObserver) OnStartSpan(span opentracing.Span, operationName string, options opentracing.StartSpanOptions) (jaeger.ContribSpanObserver, bool) {
 	return nil, false
 }
+
+type fakeInjector struct{}
+
+func (fakeInjector) Inject(ctx jaeger.SpanContext, carrier interface{}) error {
+	return nil
+}
+
+type fakeExtractor struct{}
+
+func (fakeExtractor) Extract(carrier interface{}) (jaeger.SpanContext, error) {
+	return jaeger.SpanContext{}, nil
+}