client: don't reuse a pick's authority override on a later attempt

stream.go

@@ -618,7 +618,12 @@ func (a *csAttempt) getTransport() error {
 
 func (a *csAttempt) newStream() error {
 	cs := a.cs
-	cs.callHdr.PreviousAttempts = cs.numRetries
+	// The header is copied because the fields set below, notably the authority
+	// override taken from the pick result, describe the endpoint picked for
+	// this attempt only. Mutating the clientStream's header would carry them
+	// into a later attempt.
+	callHdr := *cs.callHdr
+	callHdr.PreviousAttempts = cs.numRetries
 
 	// Merge metadata stored in PickResult, if any, with existing call metadata.
 	// It is safe to overwrite the csAttempt's context here, since all state
@@ -642,11 +647,11 @@ func (a *csAttempt) newStream() error {
 		// apply it, as specified in gRFC A81.
 		if cs.callInfo.authority == "" {
 			if authMD := a.pickResult.Metadata.Get(":authority"); len(authMD) > 0 {
-				cs.callHdr.Authority = authMD[0]
+				callHdr.Authority = authMD[0]
 			}
 		}
 	}
-	s, err := a.transport.NewStream(a.ctx, cs.callHdr, a.statsHandler)
+	s, err := a.transport.NewStream(a.ctx, &callHdr, a.statsHandler)
 	if err != nil {
 		nse, ok := err.(*transport.NewStreamError)
 		if !ok {

test/retry_test.go

@@ -27,15 +27,20 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal/balancer/stub"
 	"google.golang.org/grpc/internal/grpcsync"
 	"google.golang.org/grpc/internal/stubserver"
 	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/manual"
 	"google.golang.org/grpc/stats"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
@@ -928,3 +933,125 @@ func (s) TestNoRetry(t *testing.T) {
 		})
 	}
 }
+
+// authorityOverridePicker returns an authority override in the pick result
+// metadata for every pick made before the server has seen the first attempt
+// of the RPC, and no metadata for later picks. This mirrors an xDS cluster
+// (gRFC A81) where only some of the endpoints carry a hostname, so only some
+// picks rewrite the authority. The decision is keyed off the server-side
+// attempt count, rather than a per-picker pick count, so that it stays stable
+// if the channel re-picks for the first attempt or installs a new picker
+// instance on a connectivity state change.
+type authorityOverridePicker struct {
+	sc             balancer.SubConn
+	authority      string
+	serverAttempts *atomic.Int32
+}
+
+func (p *authorityOverridePicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
+	res := balancer.PickResult{SubConn: p.sc}
+	if p.serverAttempts.Load() == 0 {
+		res.Metadata = metadata.Pairs(":authority", p.authority)
+	}
+	return res, nil
+}
+
+// TestAuthorityOverrideNotReusedAcrossAttempts verifies that an authority
+// override supplied by the LB picker applies only to the attempt it was picked
+// for. A retry attempt whose pick carries no override must fall back to the
+// channel's authority instead of reusing the previous attempt's override.
+func (s) TestAuthorityOverrideNotReusedAcrossAttempts(t *testing.T) {
+	const (
+		balancerName      = "authority-override-retry-balancer"
+		overrideAuthority = "picked-endpoint.example.com"
+		wantAuthority     = "test.server"
+	)
+
+	// Incremented by the server handler as attempts arrive, and read by the
+	// picker to decide whether to return the authority override.
+	serverAttempts := &atomic.Int32{}
+
+	bf := stub.BalancerFuncs{
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			addrs := ccs.ResolverState.Addresses
+			if len(addrs) == 0 {
+				return nil
+			}
+			var sc balancer.SubConn
+			sc, err := bd.ClientConn.NewSubConn(addrs[:1], balancer.NewSubConnOptions{
+				StateListener: func(state balancer.SubConnState) {
+					bd.ClientConn.UpdateState(balancer.State{
+						ConnectivityState: state.ConnectivityState,
+						Picker:            &authorityOverridePicker{sc: sc, authority: overrideAuthority, serverAttempts: serverAttempts},
+					})
+				},
+			})
+			if err != nil {
+				return err
+			}
+			sc.Connect()
+			return nil
+		},
+	}
+	stub.Register(balancerName, bf)
+
+	// Attempts run sequentially (the retry is only scheduled after the first
+	// attempt's handler has returned), so authorities needs no extra
+	// synchronization.
+	var authorities []string
+	ss := &stubserver.StubServer{
+		EmptyCallF: func(ctx context.Context, _ *testpb.Empty) (*testpb.Empty, error) {
+			md, _ := metadata.FromIncomingContext(ctx)
+			authorities = append(authorities, md.Get(":authority")...)
+			// Fail the first attempt with a retryable code so that the RPC is
+			// retried, and let the second attempt succeed.
+			if serverAttempts.Add(1) == 1 {
+				return nil, status.Error(codes.Unavailable, "forcing a retry")
+			}
+			return &testpb.Empty{}, nil
+		},
+	}
+	if err := ss.StartServer(); err != nil {
+		t.Fatalf("Failed to start server: %v", err)
+	}
+	defer ss.Stop()
+
+	r := manual.NewBuilderWithScheme("whatever")
+	r.InitialState(resolver.State{Addresses: []resolver.Address{{Addr: ss.Address}}})
+
+	sc := fmt.Sprintf(`{
+		"loadBalancingConfig": [{%q: {}}],
+		"methodConfig": [{
+			"name": [{"service": "grpc.testing.TestService"}],
+			"retryPolicy": {
+				"maxAttempts": 2,
+				"initialBackoff": "0.01s",
+				"maxBackoff": "0.01s",
+				"backoffMultiplier": 1.0,
+				"retryableStatusCodes": ["UNAVAILABLE"]
+			}
+		}]
+	}`, balancerName)
+
+	cc, err := grpc.NewClient(r.Scheme()+":///"+wantAuthority, grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithResolvers(r), grpc.WithDefaultServiceConfig(sc))
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if _, err := testgrpc.NewTestServiceClient(cc).EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Fatalf("EmptyCall() failed: %v", err)
+	}
+
+	if len(authorities) != 2 {
+		t.Fatalf("Server saw %d attempts (%q), want 2", len(authorities), authorities)
+	}
+	if authorities[0] != overrideAuthority {
+		t.Errorf("First attempt used authority %q, want %q", authorities[0], overrideAuthority)
+	}
+	if authorities[1] != wantAuthority {
+		t.Errorf("Retry attempt used authority %q, want %q", authorities[1], wantAuthority)
+	}
+}