Merge pull request #1340 from circleci/dm/bpg-shim

Make it easier to use a custom provider

Author: danmux

config/o11y/otel.go

@@ -75,6 +75,9 @@ type OtelConfig struct {
 
 	// SpanExporters allows you explicitly provide a set of exporters, as an advanced use-case.
 	SpanExporters []sdktrace.SpanExporter
+
+	// ProviderFunc is used to provide a custom provider.
+	ProviderFunc func(conf otel.Config) (o11y.Provider, error)
 }
 
 // Otel is the primary entrypoint to initialize the o11y system for otel.
@@ -90,7 +93,10 @@ func Otel(ctx context.Context, o OtelConfig) (context.Context, func(context.Cont
 	}
 	cfg.Metrics = mProv
 
-	o11yProvider, err := otel.New(cfg)
+	if o.ProviderFunc == nil {
+		o.ProviderFunc = otel.New
+	}
+	o11yProvider, err := o.ProviderFunc(cfg)
 	if err != nil {
 		return ctx, nil, err
 	}

o11y/otel/custom.go

@@ -9,6 +9,14 @@ import (
 
 var globalFields = Annotator{}
 
+// GlobalFieldsAnnotator exposes the global fields annotator. When used in conjunction with a Provider, such as the
+// MetricsOnly Provider a custom provider can implement the global fields behaviour.
+func GlobalFieldsAnnotator() *Annotator {
+	return &globalFields
+}
+
+var _ sdktrace.SpanProcessor = &Annotator{}
+
 // Annotator is a SpanProcessor that adds attributes to all started spans.
 type Annotator struct {
 	attrs []attribute.KeyValue

o11y/otel/otel.go

@@ -22,6 +22,7 @@ import (
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	semconv "go.opentelemetry.io/otel/semconv/v1.12.0"
 	"go.opentelemetry.io/otel/trace"
+	"go.opentelemetry.io/otel/trace/noop"
 
 	"github.com/circleci/ex/config/secret"
 	"github.com/circleci/ex/o11y"
@@ -82,11 +83,11 @@ func New(conf Config) (o11y.Provider, error) {
 		exporters = append(exporters, http)
 	}
 
-	var sampler *deterministicSampler
+	var sampler *DeterministicSampler
 	if conf.SampleTraces {
-		sampler = &deterministicSampler{
-			sampleKeyFunc: conf.SampleKeyFunc,
-			sampleRates:   conf.SampleRates,
+		sampler = &DeterministicSampler{
+			SampleKeyFunc: conf.SampleKeyFunc,
+			SampleRates:   conf.SampleRates,
 		}
 	}
 
@@ -104,9 +105,9 @@ func New(conf Config) (o11y.Provider, error) {
 		exporters = append(exporters, text)
 	}
 
-	tp := traceProvider(multipleExporter{
-		exporters: exporters,
-		sampler:   sampler,
+	tp := traceProvider(MultipleExporter{
+		Exporters: exporters,
+		Sampler:   sampler,
 	}, conf)
 
 	// set the global options
@@ -121,6 +122,15 @@ func New(conf Config) (o11y.Provider, error) {
 	}, nil
 }
 
+// NewMetricsOnly returns a metrics only provider, to capture the span metrics behavior.
+// This can be used to compose with a custom provider to access span based metrics.
+func NewMetricsOnly(metrics o11y.ClosableMetricsProvider) o11y.Provider {
+	return &Provider{
+		metricsProvider: metrics,
+		tracer:          noop.NewTracerProvider().Tracer(""),
+	}
+}
+
 func NewHttpExporter(conf Config) (*otlptrace.Exporter, error) {
 	var serviceName, serviceVersion string
 	for _, a := range conf.ResourceAttributes {
@@ -326,7 +336,9 @@ func (o Provider) Log(ctx context.Context, name string, fields ...o11y.Pair) {
 
 func (o Provider) Close(ctx context.Context) {
 	// TODO Handle these errors in a sensible manner where possible
-	_ = o.tp.Shutdown(ctx)
+	if o.tp != nil {
+		_ = o.tp.Shutdown(ctx)
+	}
 	if o.metricsProvider != nil {
 		_ = o.metricsProvider.Close()
 	}
@@ -543,37 +555,37 @@ func (s *span) snapshotFields() map[string]any {
 	return res
 }
 
-type multipleExporter struct {
-	exporters []sdktrace.SpanExporter
-	sampler   *deterministicSampler
+type MultipleExporter struct {
+	Exporters []sdktrace.SpanExporter
+	Sampler   *DeterministicSampler
 }
 
-func (m multipleExporter) ExportSpans(ctx context.Context, spans []sdktrace.ReadOnlySpan) error {
+func (m MultipleExporter) ExportSpans(ctx context.Context, spans []sdktrace.ReadOnlySpan) error {
 	spans = m.sampleSpans(spans)
-	for _, e := range m.exporters {
+	for _, e := range m.Exporters {
 		if err := e.ExportSpans(ctx, spans); err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
-func (m multipleExporter) Shutdown(ctx context.Context) error {
-	for _, e := range m.exporters {
+func (m MultipleExporter) Shutdown(ctx context.Context) error {
+	for _, e := range m.Exporters {
 		if err := e.Shutdown(ctx); err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
-func (m multipleExporter) sampleSpans(spans []sdktrace.ReadOnlySpan) []sdktrace.ReadOnlySpan {
-	if m.sampler == nil {
+func (m MultipleExporter) sampleSpans(spans []sdktrace.ReadOnlySpan) []sdktrace.ReadOnlySpan {
+	if m.Sampler == nil {
 		return spans
 	}
 	ss := make([]sdktrace.ReadOnlySpan, 0, len(spans))
 	for _, s := range spans {
-		if ok, rate := m.sampler.shouldSample(s); ok {
+		if ok, rate := m.Sampler.shouldSample(s); ok {
 			ss = append(ss, sampleRateSpan{ReadOnlySpan: s, rate: rate})
 		}
 	}

o11y/otel/sampler.go

@@ -7,13 +7,13 @@ import (
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 )
 
-type deterministicSampler struct {
-	sampleKeyFunc func(map[string]any) string
-	sampleRates   map[string]uint
+type DeterministicSampler struct {
+	SampleKeyFunc func(map[string]any) string
+	SampleRates   map[string]uint
 }
 
 // shouldSample means should sample in, returning true if the span should be sampled in (kept)
-func (s deterministicSampler) shouldSample(p sdktrace.ReadOnlySpan) (bool, uint) {
+func (s DeterministicSampler) shouldSample(p sdktrace.ReadOnlySpan) (bool, uint) {
 	fields := map[string]any{}
 	for _, attr := range p.Attributes() {
 		fields[string(attr.Key)] = attr.Value.AsInterface()
@@ -30,8 +30,8 @@ func (s deterministicSampler) shouldSample(p sdktrace.ReadOnlySpan) (bool, uint)
 	fields["duration_ms"] = int(p.EndTime().Sub(p.StartTime()).Milliseconds())
 	fields["name"] = p.Name()
 
-	key := s.sampleKeyFunc(fields)
-	rate, ok := s.sampleRates[key] // no rate found means keep
+	key := s.SampleKeyFunc(fields)
+	rate, ok := s.SampleRates[key] // no rate found means keep
 	if !ok {
 		return true, 1 // and is a sample rate of 1/1
 	}