Rework trace.IDGenerator interface

* NewTraceID() -> NewIDs(ctx) ** Returns both TraceID and SpanID * NewSpanID() -> NewSpanID(ctx, traceID) ** Returns only SpanID, has access to TraceID * Both methods now receive a context, from which they may extract information * startSpanInternal() updated to receive a context to pass to the ID generator
open-telemetry · Dec 8, 2020 · 30e6437 · 30e6437
1 parent 4f95b6e
commit 30e6437
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 14 deletions.
diff --git a/sdk/trace/id_generator.go b/sdk/trace/id_generator.go
@@ -15,6 +15,7 @@
 package trace // import "go.opentelemetry.io/otel/sdk/trace"
 
 import (
+	"context"
 	crand "crypto/rand"
 	"encoding/binary"
 	"math/rand"
@@ -25,8 +26,8 @@ import (
 
 // IDGenerator allows custom generators for TraceID and SpanID.
 type IDGenerator interface {
-	NewTraceID() trace.TraceID
-	NewSpanID() trace.SpanID
+	NewIDs(ctx context.Context) (trace.TraceID, trace.SpanID)
+	NewSpanID(ctx context.Context, traceID trace.TraceID) trace.SpanID
 }
 
 type randomIDGenerator struct {
@@ -37,22 +38,24 @@ type randomIDGenerator struct {
 var _ IDGenerator = &randomIDGenerator{}
 
 // NewSpanID returns a non-zero span ID from a randomly-chosen sequence.
-func (gen *randomIDGenerator) NewSpanID() trace.SpanID {
+func (gen *randomIDGenerator) NewSpanID(ctx context.Context, traceID trace.TraceID) trace.SpanID {
 	gen.Lock()
 	defer gen.Unlock()
 	sid := trace.SpanID{}
 	gen.randSource.Read(sid[:])
 	return sid
 }
 
-// NewTraceID returns a non-zero trace ID from a randomly-chosen sequence.
-// mu should be held while this function is called.
-func (gen *randomIDGenerator) NewTraceID() trace.TraceID {
+// NewIDs returns a non-zero trace ID and a non-zero span ID from a
+// randomly-chosen sequence. mu should be held while this function is called.
+func (gen *randomIDGenerator) NewIDs(ctx context.Context) (trace.TraceID, trace.SpanID) {
 	gen.Lock()
 	defer gen.Unlock()
 	tid := trace.TraceID{}
 	gen.randSource.Read(tid[:])
-	return tid
+	sid := trace.SpanID{}
+	gen.randSource.Read(sid[:])
+	return tid, sid
 }
 
 func defaultIDGenerator() IDGenerator {

diff --git a/sdk/trace/sampling_test.go b/sdk/trace/sampling_test.go
@@ -15,6 +15,7 @@
 package trace
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 	"testing"
@@ -178,7 +179,7 @@ func TestTraceIdRatioSamplesInclusively(t *testing.T) {
 		samplerHi := TraceIDRatioBased(ratioHi)
 		samplerLo := TraceIDRatioBased(ratioLo)
 		for j := 0; j < numTraces; j++ {
-			traceID := idg.NewTraceID()
+			traceID, _ := idg.NewIDs(context.Background())
 
 			params := SamplingParameters{TraceID: traceID}
 			if samplerLo.ShouldSample(params).Decision == RecordAndSample {

diff --git a/sdk/trace/span.go b/sdk/trace/span.go
@@ -15,6 +15,7 @@
 package trace // import "go.opentelemetry.io/otel/sdk/trace"
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"reflect"
@@ -306,18 +307,21 @@ func (s *span) addChild() {
 	s.mu.Unlock()
 }
 
-func startSpanInternal(tr *tracer, name string, parent trace.SpanContext, remoteParent bool, o *trace.SpanConfig) *span {
+func startSpanInternal(ctx context.Context, tr *tracer, name string, parent trace.SpanContext, remoteParent bool, o *trace.SpanConfig) *span {
 	var noParent bool
 	span := &span{}
 	span.spanContext = parent
 
 	cfg := tr.provider.config.Load().(*Config)
 
 	if parent == emptySpanContext {
-		span.spanContext.TraceID = cfg.IDGenerator.NewTraceID()
+		// Generate both TraceID and SpanID
+		span.spanContext.TraceID, span.spanContext.SpanID = cfg.IDGenerator.NewIDs(ctx)
 		noParent = true
+	} else {
+		// TraceID already exists, just generate a SpanID
+		span.spanContext.SpanID = cfg.IDGenerator.NewSpanID(ctx, parent.TraceID)
 	}
-	span.spanContext.SpanID = cfg.IDGenerator.NewSpanID()
 	data := samplingData{
 		noParent:     noParent,
 		remoteParent: remoteParent,

diff --git a/sdk/trace/trace_test.go b/sdk/trace/trace_test.go
@@ -263,9 +263,10 @@ func TestSampling(t *testing.T) {
 			for i := 0; i < total; i++ {
 				ctx := context.Background()
 				if tc.parent {
+					tid, sid := idg.NewIDs(ctx)
 					psc := trace.SpanContext{
-						TraceID: idg.NewTraceID(),
-						SpanID:  idg.NewSpanID(),
+						TraceID: tid,
+						SpanID:  sid,
 					}
 					if tc.sampledParent {
 						psc.TraceFlags = trace.FlagsSampled

diff --git a/sdk/trace/tracer.go b/sdk/trace/tracer.go
@@ -47,7 +47,7 @@ func (tr *tracer) Start(ctx context.Context, name string, options ...trace.SpanO
 		}
 	}
 
-	span := startSpanInternal(tr, name, parentSpanContext, remoteParent, config)
+	span := startSpanInternal(ctx, tr, name, parentSpanContext, remoteParent, config)
 	for _, l := range links {
 		span.addLink(l)
 	}