stats/opentelemetry: record retry attempts from clientStream

stats/opentelemetry/client_tracing.go

@@ -48,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
@@ -58,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
@@ -124,14 +128,9 @@ 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)
-	var counter *int32
-	if val := ctx.Value(retryCountKey{}); val != nil {
-		counter = val.(*int32)
-	} else {
-		counter = new(int32)
-		ctx = context.WithValue(ctx, retryCountKey{}, counter)
+	if ci, ok := getRetryCount(ctx); ok {
+		ai.previousRPCAttempts = uint32(atomic.LoadInt32(&ci.numRetries))
 	}
-	ai.previousRPCAttempts = uint32(atomic.LoadInt32(counter))
 	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"
 
@@ -928,7 +929,7 @@ func (s) TestMetricsAndTracesOptionEnabled(t *testing.T) {
 				},
 				{
 					Key:   "previous-rpc-attempts",
-					Value: attribute.IntValue(1),
+					Value: attribute.IntValue(0),
 				},
 				{
 					Key:   "transparent-retry",
@@ -1021,7 +1022,7 @@ func (s) TestMetricsAndTracesOptionEnabled(t *testing.T) {
 				},
 				{
 					Key:   "previous-rpc-attempts",
-					Value: attribute.IntValue(1),
+					Value: attribute.IntValue(0),
 				},
 				{
 					Key:   "transparent-retry",
@@ -1144,7 +1145,7 @@ func (s) TestSpan(t *testing.T) {
 				},
 				{
 					Key:   "previous-rpc-attempts",
-					Value: attribute.IntValue(1),
+					Value: attribute.IntValue(0),
 				},
 				{
 					Key:   "transparent-retry",
@@ -1229,7 +1230,7 @@ func (s) TestSpan(t *testing.T) {
 				},
 				{
 					Key:   "previous-rpc-attempts",
-					Value: attribute.IntValue(1),
+					Value: attribute.IntValue(0),
 				},
 				{
 					Key:   "transparent-retry",
@@ -1354,7 +1355,7 @@ func (s) TestSpan_WithW3CContextPropagator(t *testing.T) {
 				},
 				{
 					Key:   "previous-rpc-attempts",
-					Value: attribute.IntValue(1),
+					Value: attribute.IntValue(0),
 				},
 				{
 					Key:   "transparent-retry",
@@ -1439,7 +1440,7 @@ func (s) TestSpan_WithW3CContextPropagator(t *testing.T) {
 				},
 				{
 					Key:   "previous-rpc-attempts",
-					Value: attribute.IntValue(1),
+					Value: attribute.IntValue(0),
 				},
 				{
 					Key:   "transparent-retry",
@@ -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.
@@ -1777,7 +1805,7 @@ func (s) TestStreamingRPC_TraceSequenceNumbers(t *testing.T) {
 			attributes: []attribute.KeyValue{
 				attribute.Bool("Client", true),
 				attribute.Bool("FailFast", true),
-				attribute.Int("previous-rpc-attempts", 1),
+				attribute.Int("previous-rpc-attempts", 0),
 				attribute.Bool("transparent-retry", false),
 			},
 		},

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, "/")
 }

stats/opentelemetry/trace.go

@@ -41,24 +41,22 @@ func populateSpan(rs stats.RPCStats, ai *attemptInfo) {
 
 	switch rs := rs.(type) {
 	case *stats.Begin:
-		retryCount := ai.previousRPCAttempts
-		if !rs.IsTransparentRetryAttempt {
-			if val := ai.ctx.Value(retryCountKey{}); val != nil {
-				// Atomic increment and get new value
-				retryCount = uint32(atomic.AddInt32(val.(*int32), 1))
-			}
-		}
 		// Note: Go always added Client and FailFast attributes even though they are not
 		// defined by the OpenCensus gRPC spec. Thus, they are unimportant for
 		// correctness.
 		span.SetAttributes(
 			attribute.Bool("Client", rs.Client),
 			attribute.Bool("FailFast", rs.FailFast),
-			attribute.Int64("previous-rpc-attempts", int64(retryCount)),
+			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: