stats/opentelemetry: record retry attempts from clientStream

stats/opentelemetry/client_tracing.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"log"
 	"strings"
+	"sync/atomic"
 
 	otelcodes "go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/trace"
@@ -47,6 +48,8 @@ func (h *clientTracingHandler) initializeTraces() {
 }
 
 func (h *clientTracingHandler) unaryInterceptor(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+	ci := &callInfo{numRetries: 0}
+	ctx = setRetryCount(ctx, ci)
 	ctx, _ = getOrCreateCallInfo(ctx, cc, method, opts...)
 
 	var span trace.Span
@@ -57,6 +60,8 @@ func (h *clientTracingHandler) unaryInterceptor(ctx context.Context, method stri
 }
 
 func (h *clientTracingHandler) streamInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+	ci := &callInfo{numRetries: 0}
+	ctx = setRetryCount(ctx, ci)
 	ctx, _ = getOrCreateCallInfo(ctx, cc, method, opts...)
 
 	var span trace.Span
@@ -118,9 +123,15 @@ func (h *clientTracingHandler) TagConn(ctx context.Context, _ *stats.ConnTagInfo
 // HandleConn exists to satisfy stats.Handler for tracing.
 func (h *clientTracingHandler) HandleConn(context.Context, stats.ConnStats) {}
 
+type retryCountKey struct{}
+
 // TagRPC implements per RPC attempt context management for traces.
 func (h *clientTracingHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
 	ctx, ai := getOrCreateRPCAttemptInfo(ctx)
+	if ci, ok := getRetryCount(ctx); ok {
+		ai.previousRPCAttempts = uint32(atomic.LoadInt32(&ci.numRetries))
+	}
+	ai.ctx = ctx
 	ctx, ai = h.traceTagRPC(ctx, ai, info.NameResolutionDelay)
 	return setRPCInfo(ctx, &rpcInfo{ai: ai})
 }

stats/opentelemetry/e2e_test.go

@@ -22,6 +22,7 @@ import (
 	"io"
 	"slices"
 	"strconv"
+	"strings"
 	"testing"
 	"time"
 
@@ -1687,6 +1688,7 @@ func (s) TestTraceSpan_WithRetriesAndNameResolutionDelay(t *testing.T) {
 				t.Fatal(err)
 			}
 			verifyTrace(t, spans, wantSpanInfo)
+			verifyPreviousRPCAttempts(t, spans)
 		})
 	}
 }
@@ -1708,6 +1710,32 @@ func verifyTrace(t *testing.T, spans tracetest.SpanStubs, want traceSpanInfo) {
 	}
 }
 
+func verifyPreviousRPCAttempts(t *testing.T, spans tracetest.SpanStubs) {
+	t.Helper()
+	const maxAttempts = 3
+	foundAttempts := make(map[int]bool)
+	observedSpans := make(map[int][]string)
+
+	for _, span := range spans {
+		if !strings.HasPrefix(span.Name, "Attempt.") {
+			continue
+		}
+		for _, attr := range span.Attributes {
+			if attr.Key == "previous-rpc-attempts" {
+				val := int(attr.Value.AsInt64())
+				foundAttempts[val] = true
+				observedSpans[val] = append(observedSpans[val], span.Name)
+			}
+		}
+	}
+
+	for i := range maxAttempts {
+		if !foundAttempts[i] {
+			t.Errorf("Missing span for retry attempt #%d (expected previous-rpc-attempts = %d)", i+1, i)
+		}
+	}
+}
+
 // TestStreamingRPC_TraceSequenceNumbers verifies that sequence numbers
 // are incremented correctly for multiple messages sent and received
 // during a streaming RPC.

stats/opentelemetry/opentelemetry.go

@@ -179,6 +179,8 @@ type callInfo struct {
 	// nameResolutionEventAdded is set when the resolver delay trace event
 	// is added. Prevents duplicate events, since it is reported per-attempt.
 	nameResolutionEventAdded atomic.Bool
+	// numRetries holds the count of non-transparent retry attempts.
+	numRetries int32
 }
 
 type callInfoKey struct{}
@@ -213,6 +215,16 @@ func getRPCInfo(ctx context.Context) *rpcInfo {
 	return ri
 }
 
+func setRetryCount(ctx context.Context, ci *callInfo) context.Context {
+	return context.WithValue(ctx, retryCountKey{}, ci)
+}
+
+// getRetryCount retrieves the retry count tracking struct from the context.
+func getRetryCount(ctx context.Context) (*callInfo, bool) {
+	ci, ok := ctx.Value(retryCountKey{}).(*callInfo)
+	return ci, ok
+}
+
 func removeLeadingSlash(mn string) string {
 	return strings.TrimLeft(mn, "/")
 }
@@ -242,6 +254,7 @@ type attemptInfo struct {
 	countSentMsg        uint32
 	countRecvMsg        uint32
 	previousRPCAttempts uint32
+	ctx                 context.Context
 }
 
 type clientMetrics struct {

stats/opentelemetry/server_tracing.go

@@ -41,6 +41,7 @@ func (h *serverTracingHandler) initializeTraces() {
 // TagRPC implements per RPC attempt context management for traces.
 func (h *serverTracingHandler) TagRPC(ctx context.Context, _ *stats.RPCTagInfo) context.Context {
 	ctx, ai := getOrCreateRPCAttemptInfo(ctx)
+	ai.ctx = ctx
 	ctx, ai = h.traceTagRPC(ctx, ai)
 	return setRPCInfo(ctx, &rpcInfo{ai: ai})
 }

stats/opentelemetry/trace.go

@@ -50,8 +50,13 @@ func populateSpan(rs stats.RPCStats, ai *attemptInfo) {
 			attribute.Int64("previous-rpc-attempts", int64(ai.previousRPCAttempts)),
 			attribute.Bool("transparent-retry", rs.IsTransparentRetryAttempt),
 		)
-		// increment previous rpc attempts applicable for next attempt
-		atomic.AddUint32(&ai.previousRPCAttempts, 1)
+		// Increment retry count for the next attempt if not a transparent
+		// retry.
+		if !rs.IsTransparentRetryAttempt {
+			if ci, ok := getRetryCount(ai.ctx); ok {
+				atomic.AddInt32(&ci.numRetries, 1)
+			}
+		}
 	case *stats.PickerUpdated:
 		span.AddEvent("Delayed LB pick complete")
 	case *stats.InPayload: