Interceptors to preserve & restore trace headers (#268)

pkg/options/options.go

@@ -17,6 +17,7 @@ import (
 	"sync"
 	"time"
 
+	"chainguard.dev/go-grpc-kit/pkg/trace"
 	"github.com/chainguard-dev/slogctx"
 	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
 	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
@@ -128,6 +129,7 @@ func GRPCOptions(delegate url.URL) (string, []grpc.DialOption) {
 		enableClientTimeHistogram()
 		return net.JoinHostPort(delegate.Hostname(), port), []grpc.DialOption{
 			grpc.WithStatsHandler(otelgrpc.NewClientHandler()),
+			grpc.WithStatsHandler(trace.PreserveTraceParentHandler),
 			grpc.WithChainUnaryInterceptor(grpc_prometheus.UnaryClientInterceptor, grpc_retry.UnaryClientInterceptor(retryOpts...)),
 			grpc.WithChainStreamInterceptor(grpc_prometheus.StreamClientInterceptor, grpc_retry.StreamClientInterceptor(retryOpts...)),
 			grpc.WithTransportCredentials(insecure.NewCredentials()),
@@ -145,6 +147,7 @@ func GRPCOptions(delegate url.URL) (string, []grpc.DialOption) {
 		enableClientTimeHistogram()
 		return net.JoinHostPort(delegate.Hostname(), port), []grpc.DialOption{
 			grpc.WithStatsHandler(otelgrpc.NewClientHandler()),
+			grpc.WithStatsHandler(trace.PreserveTraceParentHandler),
 			grpc.WithChainUnaryInterceptor(grpc_prometheus.UnaryClientInterceptor, grpc_retry.UnaryClientInterceptor(retryOpts...)),
 			grpc.WithChainStreamInterceptor(grpc_prometheus.StreamClientInterceptor, grpc_retry.StreamClientInterceptor(retryOpts...)),
 			grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{

pkg/trace/preserve.go

@@ -0,0 +1,57 @@
+/*
+Copyright 2024 Chainguard, Inc.
+SPDX-License-Identifier: Apache-2.0
+*/
+
+package trace
+
+import (
+	"context"
+
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/stats"
+)
+
+const (
+	OriginalTraceParentHeader string = "original-traceparent"
+	TraceParentHeader         string = "traceparent"
+)
+
+var (
+	// PreserveTraceParentHandler is a client stats.Handler that preserves the original
+	// traceparent header in the outgoing context, with a different header name.
+	//
+	// This is useful when the next hop in the request chain (like Cloud Run) may lose span
+	// information, and become an unreliable span. In those cases, we just use the original
+	// traceparent header to associate child spans directly with the outgoing span here.
+	PreserveTraceParentHandler stats.Handler = &preserveTraceParentHandler{}
+)
+
+type preserveTraceParentHandler struct{}
+
+// TagRPC implements stats.Handler interface.
+func (*preserveTraceParentHandler) TagRPC(ctx context.Context, _ *stats.RPCTagInfo) context.Context {
+	md, ok := metadata.FromOutgoingContext(ctx)
+	if !ok {
+		md = metadata.MD{}
+	}
+	if tp := md.Get(TraceParentHeader); len(tp) > 0 {
+		md.Set(OriginalTraceParentHeader, tp...)
+	}
+	return metadata.NewOutgoingContext(ctx, md)
+}
+
+// HandleRPC implments stats.Handler interface.
+func (*preserveTraceParentHandler) HandleRPC(context.Context, stats.RPCStats) {
+	// Do nothing
+}
+
+// TagConn implements stats.Handler interface.
+func (*preserveTraceParentHandler) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
+	return ctx
+}
+
+// HandleConn processes the Conn stats.
+func (*preserveTraceParentHandler) HandleConn(context.Context, stats.ConnStats) {
+	// Do nothing
+}

pkg/trace/restore.go

@@ -0,0 +1,49 @@
+/*
+Copyright 2024 Chainguard, Inc.
+SPDX-License-Identifier: Apache-2.0
+*/
+
+package trace
+
+import (
+	"context"
+
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/stats"
+)
+
+var (
+	// RestoreTraceParentHandler is a server stats.Handler that restores the traceparent
+	// stored by the PreserveTraceParentHandler.
+	RestoreTraceParentHandler stats.Handler = &restoreTraceParentHandler{}
+)
+
+type restoreTraceParentHandler struct{}
+
+// TagRPC implements stats.Handler.
+func (r *restoreTraceParentHandler) TagRPC(ctx context.Context, _ *stats.RPCTagInfo) context.Context {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		md = metadata.MD{}
+	}
+	if tp := md.Get(OriginalTraceParentHeader); len(tp) > 0 {
+		md.Set(TraceParentHeader, tp...)
+	}
+	return metadata.NewIncomingContext(ctx, md)
+}
+
+// HandleRPC implements stats.Handler.
+func (r *restoreTraceParentHandler) HandleRPC(_ context.Context, _ stats.RPCStats) {
+	// Do nothing.
+}
+
+// TagConn implements stats.Handler.
+func (r *restoreTraceParentHandler) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
+	// Do nothing.
+	return ctx
+}
+
+// HandleConn implements stats.Handler.
+func (r *restoreTraceParentHandler) HandleConn(_ context.Context, _ stats.ConnStats) {
+	// Do nothing.
+}