chore: polish span profiler (#2811)

Author: kolesnikovae

examples/golang-push/rideshare/go.work.sum

@@ -16,6 +16,7 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/rogpeppe/fastuuid v1.2.0 h1:Ppwyp6VYCF1nvBTXL3trRso7mXMlRrw9ooo375wvi2s=
 github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=
 github.com/stretchr/objx v0.1.0 h1:4G4v2dO3VZwixGIRoQ5Lfboy6nUhCyYzaqnIAPPhYs4=
+github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c=
 golang.org/x/crypto v0.11.0 h1:6Ewdq3tDic1mg5xRO4milcWCfMVQhI4NkqWWvqejpuA=
 golang.org/x/crypto v0.14.0 h1:wBqGXzWJW6m1XrIKlAH0Hs1JJ7+9KBwnIO8v66Q9cHc=
 golang.org/x/crypto v0.15.0 h1:frVn1TEaCEaZcn3Tmd7Y2b5KKPaZ+I32Q2OA3kYp5TA=

pkg/util/spanprofiler/README.md

@@ -21,12 +21,15 @@ You can wrap the global tracer using `spanprofiler.NewTracer`:
 ```go
 func main() {
     // Initialize your OpenTracing tracer
-    tracer := // Your opentracing tracer initialization
-    wrappedTracer := spanprofiler.NewTracer(tracer) // Wrap the global tracer
-
+    tracer := opentracing.GlobalTracer()
+    // Wrap it with the tracer-profiler 
+    wrappedTracer := spanprofiler.NewTracer(tracer)
     // Use the wrapped tracer in your application
     opentracing.SetGlobalTracer(wrappedTracer)
 
+	// Or, as an oneliner:
+	// opentracing.SetGlobalTracer(spanprofiler.NewTracer(opentracing.GlobalTracer()))
+
     // Your application logic here
 }
 ``` 
@@ -49,7 +52,7 @@ The `spanprofiler.StartSpanFromContext` function allows you to granularly contro
 
 func main() {
     // Start a span and enable profiling for it
-    span, ctx := spanprofiler.StartSpanFromContext(context.Background(), "YourSpanOperation", tracer)
+    span, ctx := spanprofiler.StartSpanFromContext(context.Background(), "YourOperationName", tracer)
     defer span.Finish() // Finish the span when done
 
     // Use the span in your application logic
@@ -63,7 +66,7 @@ seamless integration, reserving explicit span profiling only for cases where spa
 
 ## Implementation details
 
-When a new trace span is created, and is eligible for profiling, the tracer sets a `span_id` and `span_name` [pprof labels](https://github.com/google/pprof/blob/master/doc/README.md#tag-filtering)
+When a new trace span is created, and is eligible for profiling, the tracer sets `span_id` and `span_name` [pprof labels](https://github.com/google/pprof/blob/master/doc/README.md#tag-filtering)
 that point to the respective span. These labels are stored in the goroutine's local storage and inherited by any
 subsequent child goroutines.
 
@@ -79,8 +82,8 @@ This allows to find such spans in the trace view (in the screenshot) and fetch p
 It's important to note that the presence of this attribute does not guarantee profile availability; stack trace samples
 might not be collected if the CPU time utilized falls below the sample interval (10ms).
 
-It is crucial to understand that this module doesn't directly control the pprof profiler; its initiation is still
-necessary for profile collection. This initiation can be achieved through the `runtime/pprof` package, or using the
+It is crucial to understand that this module doesn't directly control the pprof profiler; its initialization is still
+necessary for profile collection. This initialization can be achieved through the `runtime/pprof` package, or using the
 [Pyroscope client](https://github.com/grafana/pyroscope-go).
 
 Limitations:

pkg/util/spanprofiler/spanprofiler.go

@@ -12,6 +12,9 @@ import (
 // any Span found within `ctx` as a ChildOfRef. If no such parent could be
 // found, StartSpanFromContext creates a root (parentless) Span.
 //
+// The call sets `operationName` as `span_name` pprof label, and the new span
+// identifier as `span_id` pprof label, if the trace is sampled.
+//
 // The second return value is a context.Context object built around the
 // returned Span.
 //
@@ -31,6 +34,9 @@ func StartSpanFromContext(ctx context.Context, operationName string, opts ...ope
 // it creates a root span. It also returns a context.Context object built
 // around the returned span.
 //
+// The call sets `operationName` as `span_name` pprof label, and the new span
+// identifier as `span_id` pprof label, if the trace is sampled.
+//
 // It's behavior is identical to StartSpanFromContext except that it takes an explicit
 // tracer as opposed to using the global tracer.
 func StartSpanFromContextWithTracer(ctx context.Context, tracer opentracing.Tracer, operationName string, opts ...opentracing.StartSpanOption) (opentracing.Span, context.Context) {
@@ -52,26 +58,30 @@ func wrapJaegerSpanWithGoroutineLabels(
 	// storage and are always copied to child goroutines. This way, stack
 	// trace samples collected during execution of child spans will be taken
 	// into account at the root.
-	var labels []string
+	var ctx context.Context
 	if spanID != "" {
-		labels = []string{spanNameLabelName, operationName, spanIDLabelName, spanID}
+		ctx = pprof.WithLabels(parentCtx, pprof.Labels(
+			spanNameLabelName, operationName,
+			spanIDLabelName, spanID))
 	} else {
 		// Even if the trace has not been sampled, we still need to keep track
 		// of samples that belong to the span (all spans with the given name).
-		labels = []string{spanNameLabelName, operationName}
+		ctx = pprof.WithLabels(parentCtx, pprof.Labels(
+			spanNameLabelName, operationName))
 	}
+	// Goroutine labels should be set as early as possible,
+	// in order to capture the overhead of the function call.
+	pprof.SetGoroutineLabels(ctx)
 	// We create a span wrapper to ensure we remove the newly attached pprof
 	// labels when span finishes. The need of this wrapper is questioned:
 	// as we do not have the original context, we could leave the goroutine
 	// labels – normally, span is finished at the very end of the goroutine's
 	// lifetime, so no significant side effects should take place.
 	w := spanWrapper{
 		parentPprofCtx:  parentCtx,
-		currentPprofCtx: pprof.WithLabels(parentCtx, pprof.Labels(labels...)),
-		Span:            span,
+		currentPprofCtx: ctx,
 	}
-	pprof.SetGoroutineLabels(w.currentPprofCtx)
-	opentracing.Tag{Key: profileIDTagKey, Value: spanID}.Set(span)
+	w.Span = span.SetTag(profileIDTagKey, spanID)
 	return &w
 }
 

pkg/util/spanprofiler/spanprofiler_test.go

@@ -13,23 +13,23 @@ func TestSpanProfiler_pprof_labels_propagation(t *testing.T) {
 	defer func() { require.NoError(t, tt.Close()) }()
 
 	t.Run("pprof labels are not propagated to child spans", func(t *testing.T) {
-		spanR, ctx := StartSpanFromContext(context.Background(), "RootSpan")
-		defer spanR.Finish()
-		rootLabels := spanPprofLabels(spanR)
+		rootSpan, ctx := StartSpanFromContext(context.Background(), "RootSpan")
+		defer rootSpan.Finish()
+		rootLabels := spanPprofLabels(rootSpan)
 		require.Equal(t, rootLabels["span_name"], "RootSpan")
 		rootSpanID, err := jaeger.SpanIDFromString(rootLabels["span_id"])
 		require.NoError(t, err)
 
 		// Regardless of anything, pprof labels are attached to the current
 		// goroutine, and the "pyroscope.profile.id" tag is set.
-		spanA, _ := StartSpanFromContext(ctx, "ChildSpan")
-		defer spanA.Finish()
-		childLabels := spanPprofLabels(spanA)
+		childSpan, _ := StartSpanFromContext(ctx, "ChildSpan")
+		defer childSpan.Finish()
+		childLabels := spanPprofLabels(childSpan)
 		require.Equal(t, childLabels["span_name"], "ChildSpan")
 		childSpanID, err := jaeger.SpanIDFromString(childLabels["span_id"])
 		require.NoError(t, err)
 
 		require.NotEqual(t, rootSpanID, childSpanID)
-		require.Equal(t, childLabels["span_id"], spanTags(spanA)["pyroscope.profile.id"])
+		require.Equal(t, childSpanID.String(), spanTags(childSpan)["pyroscope.profile.id"])
 	})
 }

pkg/util/spanprofiler/tracer.go

@@ -17,6 +17,22 @@ const (
 
 type tracer struct{ opentracing.Tracer }
 
+// NewTracer creates a new opentracing.Tracer with the span profiler integrated.
+//
+// For efficiency, the tracer selectively records profiles for _root_ spans
+// — the initial _local_ span in a process — since a trace may encompass
+// thousands of spans. All stack trace samples accumulated during the execution
+// of their child spans contribute to the root span's profile. In practical
+// terms, this signifies that, for instance, an HTTP request results in a
+// singular profile, irrespective of the numerous spans within the trace. It's
+// important to note that these profiles don't extend beyond the boundaries of
+// a single process.
+//
+// The limitation of this approach is that only spans created within the same
+// goroutine, or its children, as the parent are taken into account.
+// Consequently, in scenarios involving asynchronous execution, where the parent
+// span context is passed to another goroutine, explicit profiling becomes
+// necessary using `spanprofiler.StartSpanFromContext`.
 func NewTracer(tr opentracing.Tracer) opentracing.Tracer { return &tracer{tr} }
 
 func (t *tracer) StartSpan(operationName string, opts ...opentracing.StartSpanOption) opentracing.Span {
@@ -25,46 +41,69 @@ func (t *tracer) StartSpan(operationName string, opts ...opentracing.StartSpanOp
 	if !ok {
 		return span
 	}
-	parent, ok := parentSpanContextFromRef(opts...)
-	if ok && !isRemoteSpan(parent) {
+	// pprof labels are attached only once, at the span root level.
+	if !isRootSpan(opts...) {
 		return span
 	}
 	// The pprof label API assumes that pairs of labels are passed through the
 	// context. Unfortunately, the opentracing Tracer API doesn't match this
-	// concept. This makes it impossible to save an existing pprof context and
+	// concept: this makes it impossible to save an existing pprof context and
 	// all the original pprof labels associated with the goroutine.
-	return wrapJaegerSpanWithGoroutineLabels(context.Background(), span, operationName, sampledSpanID(spanCtx))
+	ctx := context.Background()
+	return wrapJaegerSpanWithGoroutineLabels(ctx, span, operationName, sampledSpanID(spanCtx))
+}
+
+// isRootSpan reports whether the span is a root span.
+//
+// There are only two valid cases: if the span is the first span in the trace,
+// or is the first _local_ span in the trace.
+//
+// An exception is made for FollowsFrom reference: spans without an explicit
+// parent are considered as root ones.
+func isRootSpan(opts ...opentracing.StartSpanOption) bool {
+	parent, ok := parentSpanContextFromRef(opts...)
+	return !ok || isRemoteSpan(parent)
 }
 
-func parentSpanContextFromRef(options ...opentracing.StartSpanOption) (jaeger.SpanContext, bool) {
+// parentSpanContextFromRef returns the first parent reference.
+func parentSpanContextFromRef(options ...opentracing.StartSpanOption) (sc jaeger.SpanContext, ok bool) {
 	var sso opentracing.StartSpanOptions
 	for _, option := range options {
 		option.Apply(&sso)
 	}
 	for _, ref := range sso.References {
 		if ref.Type == opentracing.ChildOfRef && ref.ReferencedContext != nil {
-			c, ok := ref.ReferencedContext.(jaeger.SpanContext)
-			return c, ok
+			sc, ok = ref.ReferencedContext.(jaeger.SpanContext)
+			return sc, ok
 		}
 	}
-	return jaeger.SpanContext{}, false
+	return sc, ok
 }
 
+// isRemoteSpan reports whether the span context represents a remote parent.
+//
+// NOTE(kolesnikovae): this is ugly, but the only reliable method I found.
+// The opentracing-go package and Jaeger client are not meant to change as
+// both are deprecated.
 func isRemoteSpan(c jaeger.SpanContext) bool {
-	// This is ugly and unsafe, but is the only reliable way to get to know which
-	// spans should be profiled. The opentracing-go package and Jaeger client
-	// are not meant to change as both are deprecated.
 	jaegerCtx := *(*jaegerSpanCtx)(unsafe.Pointer(&c))
 	return jaegerCtx.remote
 }
 
+// jaegerSpanCtx represents memory layout of the jaeger.SpanContext type.
 type jaegerSpanCtx struct {
-	traceID       [16]byte   // TraceID
-	spanID        [8]byte    // SpanID
-	parentID      [8]byte    // SpanID
-	baggage       uintptr    // map[string]string
-	debugID       [2]uintptr // string
+	traceID  [16]byte   // TraceID
+	spanID   [8]byte    // SpanID
+	parentID [8]byte    // SpanID
+	baggage  uintptr    // map[string]string
+	debugID  [2]uintptr // string
+
+	// samplingState is a pointer to a struct that has "localRootSpan" member,
+	// which we could probably use: that would allow omitting quite expensive
+	// parentSpanContextFromRef call. However, interpreting the pointer and
+	// the complex struct memory layout is more complicated and dangerous.
 	samplingState uintptr
+
 	// remote indicates that span context represents a remote parent
 	remote bool
 }

pkg/util/spanprofiler/tracer_test.go

@@ -17,63 +17,63 @@ func TestTracer_pprof_labels_propagation(t *testing.T) {
 	defer func() { require.NoError(t, tt.Close()) }()
 
 	t.Run("root span name and ID are propagated as pprof labels", func(t *testing.T) {
-		spanR, _ := opentracing.StartSpanFromContext(context.Background(), "RootSpan")
-		defer spanR.Finish()
-		pprofLabels := spanPprofLabels(spanR)
+		rootSpan, _ := opentracing.StartSpanFromContext(context.Background(), "RootSpan")
+		defer rootSpan.Finish()
+		pprofLabels := spanPprofLabels(rootSpan)
 		// Label / tag names should be specified explicitly.
 		require.Equal(t, pprofLabels["span_name"], "RootSpan")
 		_, err := jaeger.SpanIDFromString(pprofLabels["span_id"])
 		require.NoError(t, err)
 		// Make sure the root span has "pyroscope.profile.id" attribute,
 		// and it matches the corresponding pprof label.
-		require.Equal(t, pprofLabels["span_id"], spanTags(spanR)["pyroscope.profile.id"])
+		require.Equal(t, pprofLabels["span_id"], spanTags(rootSpan)["pyroscope.profile.id"])
 	})
 
 	t.Run("pprof labels are propagated to child spans", func(t *testing.T) {
-		spanR, ctx := opentracing.StartSpanFromContext(context.Background(), "RootSpan")
-		defer spanR.Finish()
-		// Nested child span has the same labels, but
-		spanA, _ := opentracing.StartSpanFromContext(ctx, "SpanA")
-		defer spanA.Finish()
+		rootSpan, ctx := opentracing.StartSpanFromContext(context.Background(), "RootSpan")
+		defer rootSpan.Finish()
+		childSpan, _ := opentracing.StartSpanFromContext(ctx, "ChildSpan")
+		defer childSpan.Finish()
 		// Goroutine labels are inherited from the parent,
 		// we do not set them repeatedly for the child spans.
-		require.Empty(t, spanPprofLabels(spanA))
+		require.Empty(t, spanPprofLabels(childSpan))
 		// Only the root span is annotated with the profile ID tag.
-		require.Nil(t, spanTags(spanA)["pyroscope.profile.id"])
+		require.Nil(t, spanTags(childSpan)["pyroscope.profile.id"])
 	})
 
 	t.Run("pprof labels are not propagated to child spans after parent is finished", func(t *testing.T) {
-		spanR, ctx := opentracing.StartSpanFromContext(context.Background(), "RootSpan")
-		rootLabels := spanPprofLabels(spanR)
-		// Finalize the span, which removes the labels from the goroutine's
-		// storage. Note that we can't access them (Go runtime does not provide
-		// public methods) but we rely on SetGoroutineLabels implementation:
+		rootSpan, ctx := opentracing.StartSpanFromContext(context.Background(), "RootSpan")
+		rootLabels := spanPprofLabels(rootSpan)
+		require.NotEmpty(t, rootLabels)
+		// This removes the labels from the goroutine's storage.
+		// Note that we can't access them (Go runtime does not provide public
+		// methods) but we rely on SetGoroutineLabels implementation:
 		// tracer alters currentPprofCtx, so that it actually points to the
 		// parentPprofCtx – the state prior StartSpanFromContext call.
-		spanR.Finish()
-		spanA, _ := opentracing.StartSpanFromContext(ctx, "SpanA")
-		defer spanA.Finish()
-		childLabels := spanPprofLabels(spanA)
-		require.NotEqual(t, rootLabels, childLabels)
+		rootSpan.Finish()
+		childSpan, _ := opentracing.StartSpanFromContext(ctx, "ChildSpan")
+		defer childSpan.Finish()
+		childLabels := spanPprofLabels(childSpan)
+		require.Empty(t, childLabels)
 	})
 
-	t.Run("pprof labels are not propagated to child spans if those are spawn in a separate goroutine hierarchy", func(t *testing.T) {
+	t.Run("pprof labels are not propagated to child spans if they are created in a separate goroutine hierarchy", func(t *testing.T) {
 		c := make(chan opentracing.SpanContext)
 		done := make(chan struct{})
 		go func() {
 			defer close(done)
 			// Normally, we assume that pprof labels are propagated to child
 			// goroutines, and we do not have to set pprof labels for each nested
-			// span repeatedly. However, if the child span is created in e.g.,
-			// sibling goroutine, no pprof labels will be attached.
-			// The exact relation (child of / follows from) does not matter.
-			span := opentracing.StartSpan("Detached", opentracing.ChildOf(<-c))
+			// span repeatedly. However, if the child span is created in a sibling
+			// goroutine, no pprof labels will be attached.
+			// The exact span relation (child of / follows from) does not matter.
+			span := opentracing.StartSpan("ChildSpan", opentracing.ChildOf(<-c))
 			require.Empty(t, spanPprofLabels(span))
 			span.Finish()
 		}()
-		spanR, _ := opentracing.StartSpanFromContext(context.Background(), "RootSpan")
-		defer spanR.Finish()
-		c <- spanR.Context()
+		rootSpan, _ := opentracing.StartSpanFromContext(context.Background(), "RootSpan")
+		defer rootSpan.Finish()
+		c <- rootSpan.Context()
 		<-done
 	})
 }