From 5ee4f6bf76b9eba522c8f0315f04551d3bc328ae Mon Sep 17 00:00:00 2001 From: Naveed Date: Wed, 29 Jul 2026 17:45:28 +0530 Subject: [PATCH 1/2] client: don't reuse a pick's authority override on a later attempt --- stream.go | 11 +- test/authority_override_retry_test.go | 162 ++++++++++++++++++++++++++ 2 files changed, 170 insertions(+), 3 deletions(-) create mode 100644 test/authority_override_retry_test.go diff --git a/stream.go b/stream.go index ebfcbee3d5c6..8c33c5c663d7 100644 --- a/stream.go +++ b/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 { diff --git a/test/authority_override_retry_test.go b/test/authority_override_retry_test.go new file mode 100644 index 000000000000..8ca01302fd7f --- /dev/null +++ b/test/authority_override_retry_test.go @@ -0,0 +1,162 @@ +/* + * + * Copyright 2026 gRPC authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package test + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + + "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/stubserver" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/resolver" + "google.golang.org/grpc/resolver/manual" + "google.golang.org/grpc/status" + + testgrpc "google.golang.org/grpc/interop/grpc_testing" + testpb "google.golang.org/grpc/interop/grpc_testing" +) + +// authorityOverridePicker returns an authority override in the pick result +// metadata for the first pick only, and no metadata for subsequent picks. This +// mirrors an xDS cluster (gRFC A81) where only some of the endpoints carry a +// hostname, so only some picks rewrite the authority. +type authorityOverridePicker struct { + sc balancer.SubConn + authority string + picks atomic.Int32 +} + +func (p *authorityOverridePicker) Pick(balancer.PickInfo) (balancer.PickResult, error) { + res := balancer.PickResult{SubConn: p.sc} + if p.picks.Add(1) == 1 { + 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" + ) + + 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}, + }) + }, + }) + if err != nil { + return err + } + sc.Connect() + return nil + }, + } + stub.Register(balancerName, bf) + + var mu sync.Mutex + var authorities []string + ss := &stubserver.StubServer{ + EmptyCallF: func(ctx context.Context, _ *testpb.Empty) (*testpb.Empty, error) { + md, _ := metadata.FromIncomingContext(ctx) + mu.Lock() + authorities = append(authorities, md.Get(":authority")...) + attempt := len(authorities) + mu.Unlock() + // Fail the first attempt with a retryable code so that the RPC is + // retried, and let the second attempt succeed. + if attempt == 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) + } + + mu.Lock() + defer mu.Unlock() + 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) + } +} From ec538d8439f6538752ec498e85dffdf8f0d9eede Mon Sep 17 00:00:00 2001 From: Naveed Date: Wed, 12 Aug 2026 15:02:31 +0530 Subject: [PATCH 2/2] move test into retry_test.go, deflake picker, drop mutex --- test/authority_override_retry_test.go | 162 -------------------------- test/retry_test.go | 127 ++++++++++++++++++++ 2 files changed, 127 insertions(+), 162 deletions(-) delete mode 100644 test/authority_override_retry_test.go diff --git a/test/authority_override_retry_test.go b/test/authority_override_retry_test.go deleted file mode 100644 index 8ca01302fd7f..000000000000 --- a/test/authority_override_retry_test.go +++ /dev/null @@ -1,162 +0,0 @@ -/* - * - * Copyright 2026 gRPC authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - * - */ - -package test - -import ( - "context" - "fmt" - "sync" - "sync/atomic" - "testing" - - "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/stubserver" - "google.golang.org/grpc/metadata" - "google.golang.org/grpc/resolver" - "google.golang.org/grpc/resolver/manual" - "google.golang.org/grpc/status" - - testgrpc "google.golang.org/grpc/interop/grpc_testing" - testpb "google.golang.org/grpc/interop/grpc_testing" -) - -// authorityOverridePicker returns an authority override in the pick result -// metadata for the first pick only, and no metadata for subsequent picks. This -// mirrors an xDS cluster (gRFC A81) where only some of the endpoints carry a -// hostname, so only some picks rewrite the authority. -type authorityOverridePicker struct { - sc balancer.SubConn - authority string - picks atomic.Int32 -} - -func (p *authorityOverridePicker) Pick(balancer.PickInfo) (balancer.PickResult, error) { - res := balancer.PickResult{SubConn: p.sc} - if p.picks.Add(1) == 1 { - 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" - ) - - 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}, - }) - }, - }) - if err != nil { - return err - } - sc.Connect() - return nil - }, - } - stub.Register(balancerName, bf) - - var mu sync.Mutex - var authorities []string - ss := &stubserver.StubServer{ - EmptyCallF: func(ctx context.Context, _ *testpb.Empty) (*testpb.Empty, error) { - md, _ := metadata.FromIncomingContext(ctx) - mu.Lock() - authorities = append(authorities, md.Get(":authority")...) - attempt := len(authorities) - mu.Unlock() - // Fail the first attempt with a retryable code so that the RPC is - // retried, and let the second attempt succeed. - if attempt == 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) - } - - mu.Lock() - defer mu.Unlock() - 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) - } -} diff --git a/test/retry_test.go b/test/retry_test.go index 87c3ec28213f..d97ef368520c 100644 --- a/test/retry_test.go +++ b/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) + } +}