From d9b367e57c3f224dc0442ca2835af76e53692da1 Mon Sep 17 00:00:00 2001 From: krsnaSuraj Date: Wed, 15 Jul 2026 12:46:28 +0530 Subject: [PATCH] fix(atenet): propagate trace context through router and namespace span/metric attrs Two bugs broke the trace chain through the router, creating two separate traces instead of one connected trace for each ResumeActor request. 1. Singleflight detaches trace context (primary): resumer.go used context.Background() inside singleflight.DoChan(). Fixed by propagating the caller's span context into the background context via trace.ContextWithSpanContext. 2. No traceparent injected into upstream request (secondary): extproc.go only rewrote the :authority header. Fixed by calling injectTraceContext(ctx, mutation) which uses otel.GetTextMapPropagator() to write traceparent/tracestate with OVERWRITE_IF_EXISTS_OR_ADD AppendAction. 3. Non-namespaced attribute keys: Renamed span and metric attributes to ate.* convention (#412 style). Replaced custom target_addr attribute with stable OTel semconv server.address + server.port. --- cmd/atenet/internal/router/extproc.go | 22 +++- cmd/atenet/internal/router/extproc_out.go | 21 ++++ cmd/atenet/internal/router/extproc_test.go | 123 +++++++++++++++++++-- cmd/atenet/internal/router/resumer.go | 7 +- cmd/atenet/internal/router/resumer_test.go | 85 ++++++++++++++ 5 files changed, 242 insertions(+), 16 deletions(-) diff --git a/cmd/atenet/internal/router/extproc.go b/cmd/atenet/internal/router/extproc.go index 203c668d1..9cc9f46d0 100644 --- a/cmd/atenet/internal/router/extproc.go +++ b/cmd/atenet/internal/router/extproc.go @@ -30,6 +30,7 @@ import ( "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/metric" "go.opentelemetry.io/otel/propagation" + "go.opentelemetry.io/otel/semconv/v1.40.0" "google.golang.org/grpc" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" @@ -147,6 +148,10 @@ func (s *ExtProcServer) handleRequestHeaders( // Host is invalid, respond with 404. return nil, metadata, "", "", "", invalidHostErr(metadata.host, err) } + span.SetAttributes( + attribute.String("ate.atespace", atespace), + attribute.String("ate.actor.name", actorName), + ) slog.InfoContext(ctx, "ResumeActor", slog.String("atespace", atespace), slog.String("actor", actorName)) actor, err := s.resumer.ResumeActor(ctx, atespace, actorName) @@ -176,9 +181,18 @@ func (s *ExtProcServer) handleRequestHeaders( slog.InfoContext(ctx, "Route ok", slog.String("actor", actorName), slog.String("targetAddr", targetAddr)) - // Route by rewriting the :authority header. + span.SetAttributes( + semconv.ServerAddress(workerIP), + semconv.ServerPort(80), + attribute.String("ate.actor.template.namespace", tmplNs), + attribute.String("ate.actor.template.name", tmplName), + ) + + // Route by rewriting the :authority header and injecting trace context + // so the upstream worker receives the same trace as the ingress path. mutation := &extprocv3.HeaderMutation{} addAuthorityMutation(targetAddr, mutation) + injectTraceContext(ctx, mutation) return &extprocv3.HeadersResponse{ Response: &extprocv3.CommonResponse{ @@ -192,9 +206,9 @@ func (s *ExtProcServer) recordRouteDuration(ctx context.Context, d time.Duration return } s.routeDuration.Record(ctx, d.Seconds(), metric.WithAttributes( - attribute.String("actor_template_namespace", tmplNs), - attribute.String("actor_template_name", tmplName), - attribute.String("outcome", outcome), + attribute.String("ate.actor.template.namespace", tmplNs), + attribute.String("ate.actor.template.name", tmplName), + attribute.String("ate.outcome", outcome), )) } diff --git a/cmd/atenet/internal/router/extproc_out.go b/cmd/atenet/internal/router/extproc_out.go index 98d21d93e..85a00653e 100644 --- a/cmd/atenet/internal/router/extproc_out.go +++ b/cmd/atenet/internal/router/extproc_out.go @@ -15,9 +15,13 @@ package router import ( + "context" + corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3" extproc "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3" envoy_type "github.com/envoyproxy/go-control-plane/envoy/type/v3" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/propagation" ) // reqError carries an HTTP-mappable status code and a client-safe message. @@ -43,6 +47,23 @@ func addAuthorityMutation(auth string, mut *extproc.HeaderMutation) { ) } +// injectTraceContext injects the trace context from ctx into the header mutation, +// so that Envoy forwards traceparent/tracestate to the upstream worker and spans +// are connected across the full request path. +func injectTraceContext(ctx context.Context, mutation *extproc.HeaderMutation) { + headers := make(map[string]string) + otel.GetTextMapPropagator().Inject(ctx, propagation.MapCarrier(headers)) + for k, v := range headers { + mutation.SetHeaders = append(mutation.SetHeaders, &corev3.HeaderValueOption{ + Header: &corev3.HeaderValue{ + Key: k, + Value: v, + }, + AppendAction: corev3.HeaderValueOption_OVERWRITE_IF_EXISTS_OR_ADD, + }) + } +} + func immediateResponse(statusCode envoy_type.StatusCode, message string) *extproc.ProcessingResponse { return &extproc.ProcessingResponse{ Response: &extproc.ProcessingResponse_ImmediateResponse{ diff --git a/cmd/atenet/internal/router/extproc_test.go b/cmd/atenet/internal/router/extproc_test.go index 81f97a944..3d03251c3 100644 --- a/cmd/atenet/internal/router/extproc_test.go +++ b/cmd/atenet/internal/router/extproc_test.go @@ -21,6 +21,7 @@ import ( "errors" "log/slog" "strings" + "sync" "testing" "time" @@ -28,6 +29,9 @@ import ( corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3" extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3" envoy_type "github.com/envoyproxy/go-control-plane/envoy/type/v3" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/propagation" + sdktrace "go.opentelemetry.io/otel/sdk/trace" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" @@ -230,17 +234,19 @@ func TestExtProcHeadersEvaluation(t *testing.T) { } mutation := res.Response.GetHeaderMutation() - if len(mutation.GetSetHeaders()) != 1 { - t.Fatalf("expected exactly one Header option set, found: %v", mutation.GetSetHeaders()) - } - - headerOption := mutation.GetSetHeaders()[0] - if strings.ToLower(headerOption.Header.Key) != ":authority" { - t.Errorf("invalid resulting dynamic parameter key: %s", headerOption.Header.Key) + headers := mutation.GetSetHeaders() + var found bool + for _, h := range headers { + if strings.ToLower(h.Header.Key) == ":authority" { + found = true + if string(h.Header.RawValue) != tc.expectedTarget { + t.Errorf("invalid destination mapping found: %s, expected: %s", h.Header.RawValue, tc.expectedTarget) + } + break + } } - - if string(headerOption.Header.RawValue) != tc.expectedTarget { - t.Errorf("invalid destination mapping found: %s, expected: %s", headerOption.Header.RawValue, tc.expectedTarget) + if !found { + t.Errorf(":authority header not found in mutation, got: %v", headers) } // Confirm that query logs recorded metric trace details @@ -252,3 +258,100 @@ func TestExtProcHeadersEvaluation(t *testing.T) { }) } } + +// inMemoryExporter stores ended spans for test inspection. +type inMemoryExporter struct { + mu sync.Mutex + spans []sdktrace.ReadOnlySpan +} + +func (e *inMemoryExporter) ExportSpans(_ context.Context, spans []sdktrace.ReadOnlySpan) error { + e.mu.Lock() + defer e.mu.Unlock() + e.spans = append(e.spans, spans...) + return nil +} + +func (e *inMemoryExporter) Shutdown(context.Context) error { + e.mu.Lock() + defer e.mu.Unlock() + e.spans = nil + return nil +} + +func (e *inMemoryExporter) Ended() []sdktrace.ReadOnlySpan { + e.mu.Lock() + defer e.mu.Unlock() + ret := make([]sdktrace.ReadOnlySpan, len(e.spans)) + copy(ret, e.spans) + return ret +} + +func TestInjectTraceContext_AddsTraceparentMatchingParentSpan(t *testing.T) { + prevProp := otel.GetTextMapPropagator() + otel.SetTextMapPropagator(propagation.TraceContext{}) + t.Cleanup(func() { otel.SetTextMapPropagator(prevProp) }) + + exporter := &inMemoryExporter{} + tp := sdktrace.NewTracerProvider( + sdktrace.WithSampler(sdktrace.AlwaysSample()), + sdktrace.WithSpanProcessor(sdktrace.NewSimpleSpanProcessor(exporter)), + ) + + ctx, span := tp.Tracer("test").Start(context.Background(), "parent") + parentTraceID := span.SpanContext().TraceID() + + mutation := &extprocv3.HeaderMutation{} + injectTraceContext(ctx, mutation) + span.End() + + var traceparent string + for _, h := range mutation.GetSetHeaders() { + if strings.ToLower(h.Header.Key) == "traceparent" { + traceparent = h.Header.Value + if h.GetAppendAction() != corev3.HeaderValueOption_OVERWRITE_IF_EXISTS_OR_ADD { + t.Errorf("AppendAction = %v, want OVERWRITE_IF_EXISTS_OR_ADD", h.GetAppendAction()) + } + break + } + } + if traceparent == "" { + t.Fatal("traceparent not found in header mutation") + } + + parts := strings.Split(traceparent, "-") + if len(parts) < 2 { + t.Fatalf("invalid traceparent format: %s", traceparent) + } + if parts[1] != parentTraceID.String() { + t.Errorf("trace ID in traceparent = %s, want parent trace ID = %s", parts[1], parentTraceID.String()) + } +} + +func TestInjectTraceContext_AppendActionDefaultsToOverwrite(t *testing.T) { + prevProp := otel.GetTextMapPropagator() + otel.SetTextMapPropagator(propagation.TraceContext{}) + t.Cleanup(func() { otel.SetTextMapPropagator(prevProp) }) + + exporter := &inMemoryExporter{} + tp := sdktrace.NewTracerProvider( + sdktrace.WithSampler(sdktrace.AlwaysSample()), + sdktrace.WithSpanProcessor(sdktrace.NewSimpleSpanProcessor(exporter)), + ) + + ctx, span := tp.Tracer("test").Start(context.Background(), "parent") + defer span.End() + + mutation := &extprocv3.HeaderMutation{} + injectTraceContext(ctx, mutation) + + for _, h := range mutation.GetSetHeaders() { + if strings.ToLower(h.Header.Key) == "traceparent" { + if h.GetAppendAction() != corev3.HeaderValueOption_OVERWRITE_IF_EXISTS_OR_ADD { + t.Errorf("AppendAction = %v, want OVERWRITE_IF_EXISTS_OR_ADD", h.GetAppendAction()) + } + return + } + } + t.Error("traceparent not found in mutation") +} diff --git a/cmd/atenet/internal/router/resumer.go b/cmd/atenet/internal/router/resumer.go index db2231c06..be7bff1db 100644 --- a/cmd/atenet/internal/router/resumer.go +++ b/cmd/atenet/internal/router/resumer.go @@ -46,8 +46,8 @@ func NewActorResumer(apiClient ateapipb.ControlClient) *ActorResumer { func (r *ActorResumer) ResumeActor(ctx context.Context, atespace, actorName string) (*ateapipb.Actor, error) { ctx, span := otel.Tracer(routerServiceName).Start(ctx, "ResumeActor", trace.WithAttributes( - attribute.String("atespace", atespace), - attribute.String("actor", actorName), + attribute.String("ate.atespace", atespace), + attribute.String("ate.actor.name", actorName), )) defer span.End() @@ -57,6 +57,9 @@ func (r *ActorResumer) ResumeActor(ctx context.Context, atespace, actorName stri // resume operation continues running for Caller 2 and Caller 3 without failing. bgCtx, bgCancel := context.WithTimeout(context.Background(), 15*time.Second) defer bgCancel() + // Propagate the caller's span context so the gRPC spans are children + // of ResumeActor rather than appearing as a separate trace. + bgCtx = trace.ContextWithSpanContext(bgCtx, trace.SpanContextFromContext(ctx)) backoff := wait.Backoff{ Steps: 7, diff --git a/cmd/atenet/internal/router/resumer_test.go b/cmd/atenet/internal/router/resumer_test.go index 149db9b0d..7c19bbc80 100644 --- a/cmd/atenet/internal/router/resumer_test.go +++ b/cmd/atenet/internal/router/resumer_test.go @@ -21,6 +21,8 @@ import ( "time" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "go.opentelemetry.io/otel" + sdktrace "go.opentelemetry.io/otel/sdk/trace" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" @@ -168,3 +170,86 @@ func TestActorResumer_ResumeActor(t *testing.T) { } }) } + +type resumerSpanExporter struct { + mu sync.Mutex + spans []sdktrace.ReadOnlySpan +} + +func (e *resumerSpanExporter) ExportSpans(_ context.Context, spans []sdktrace.ReadOnlySpan) error { + e.mu.Lock() + defer e.mu.Unlock() + e.spans = append(e.spans, spans...) + return nil +} + +func (e *resumerSpanExporter) Shutdown(context.Context) error { + e.mu.Lock() + defer e.mu.Unlock() + e.spans = nil + return nil +} + +func (e *resumerSpanExporter) Ended() []sdktrace.ReadOnlySpan { + e.mu.Lock() + defer e.mu.Unlock() + ret := make([]sdktrace.ReadOnlySpan, len(e.spans)) + copy(ret, e.spans) + return ret +} + +func TestActorResumer_ResumeChildSharesParentTraceID(t *testing.T) { + const testActorName = "actor-a" + const testAtespace = "team-a" + const expectedIP = "10.0.0.52" + + exporter := &resumerSpanExporter{} + tp := sdktrace.NewTracerProvider( + sdktrace.WithSampler(sdktrace.AlwaysSample()), + sdktrace.WithSpanProcessor(sdktrace.NewSimpleSpanProcessor(exporter)), + ) + prevTP := otel.GetTracerProvider() + otel.SetTracerProvider(tp) + t.Cleanup(func() { otel.SetTracerProvider(prevTP) }) + + ctx, parent := tp.Tracer("test").Start(context.Background(), "parent") + parentTraceID := parent.SpanContext().TraceID() + parentSpanID := parent.SpanContext().SpanID() + + mock := &resumerMockClient{ + resumeFn: func(ctx context.Context, in *ateapipb.ResumeActorRequest, opts ...grpc.CallOption) (*ateapipb.ResumeActorResponse, error) { + return &ateapipb.ResumeActorResponse{ + Actor: &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Name: testActorName}, + Status: ateapipb.Actor_STATUS_RUNNING, + AteomPodIp: expectedIP, + }, + }, nil + }, + } + + resumer := NewActorResumer(mock) + _, err := resumer.ResumeActor(ctx, testAtespace, testActorName) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + parent.End() + + var resumeSpan sdktrace.ReadOnlySpan + for _, s := range exporter.Ended() { + if s.Name() == "ResumeActor" { + resumeSpan = s + break + } + } + if resumeSpan == nil { + t.Fatal("ResumeActor span not found in exporter") + } + + if got, want := resumeSpan.SpanContext().TraceID(), parentTraceID; got != want { + t.Errorf("ResumeActor trace ID = %s, want parent trace ID = %s", got, want) + } + if got, want := resumeSpan.Parent().SpanID(), parentSpanID; got != want { + t.Errorf("ResumeActor parent span ID = %s, want %s", got, want) + } +}