@@ -20,6 +20,7 @@ package pickfirst_test
2020
2121import (
2222 "context"
23+ "encoding/json"
2324 "errors"
2425 "fmt"
2526 "strings"
@@ -28,11 +29,14 @@ import (
2829
2930 "google.golang.org/grpc"
3031 "google.golang.org/grpc/backoff"
32+ "google.golang.org/grpc/balancer"
33+ pfbalancer "google.golang.org/grpc/balancer/pickfirst"
3134 pfinternal "google.golang.org/grpc/balancer/pickfirst/internal"
3235 "google.golang.org/grpc/codes"
3336 "google.golang.org/grpc/connectivity"
3437 "google.golang.org/grpc/credentials/insecure"
3538 "google.golang.org/grpc/internal"
39+ "google.golang.org/grpc/internal/balancer/stub"
3640 "google.golang.org/grpc/internal/channelz"
3741 "google.golang.org/grpc/internal/grpctest"
3842 "google.golang.org/grpc/internal/stubserver"
@@ -463,6 +467,85 @@ func (s) TestPickFirst_ShuffleAddressList(t *testing.T) {
463467 }
464468}
465469
470+ // Tests the PF LB policy with shuffling enabled. It explicitly unsets the
471+ // Endpoints field in the resolver update to test the shuffling of the
472+ // Addresses.
473+ func (s ) TestPickFirst_ShuffleAddressListNoEndpoints (t * testing.T ) {
474+ // Install a shuffler that always reverses two entries.
475+ origShuf := pfinternal .RandShuffle
476+ defer func () { pfinternal .RandShuffle = origShuf }()
477+ pfinternal .RandShuffle = func (n int , f func (int , int )) {
478+ if n != 2 {
479+ t .Errorf ("Shuffle called with n=%v; want 2" , n )
480+ return
481+ }
482+ f (0 , 1 ) // reverse the two addresses
483+ }
484+
485+ pfBuilder := balancer .Get (pfbalancer .Name )
486+ shuffleConfig , err := pfBuilder .(balancer.ConfigParser ).ParseConfig (json .RawMessage (`{ "shuffleAddressList": true }` ))
487+ if err != nil {
488+ t .Fatal (err )
489+ }
490+ noShuffleConfig , err := pfBuilder .(balancer.ConfigParser ).ParseConfig (json .RawMessage (`{ "shuffleAddressList": false }` ))
491+ if err != nil {
492+ t .Fatal (err )
493+ }
494+ var activeCfg serviceconfig.LoadBalancingConfig
495+
496+ bf := stub.BalancerFuncs {
497+ Init : func (bd * stub.BalancerData ) {
498+ bd .ChildBalancer = pfBuilder .Build (bd .ClientConn , bd .BuildOptions )
499+ },
500+ Close : func (bd * stub.BalancerData ) {
501+ bd .ChildBalancer .Close ()
502+ },
503+ UpdateClientConnState : func (bd * stub.BalancerData , ccs balancer.ClientConnState ) error {
504+ ccs .BalancerConfig = activeCfg
505+ ccs .ResolverState .Endpoints = nil
506+ return bd .ChildBalancer .UpdateClientConnState (ccs )
507+ },
508+ }
509+
510+ stub .Register (t .Name (), bf )
511+ svcCfg := fmt .Sprintf (`{ "loadBalancingConfig": [{%q: {}}] }` , t .Name ())
512+ // Set up our backends.
513+ cc , r , backends := setupPickFirst (t , 2 , grpc .WithDefaultServiceConfig (svcCfg ))
514+ addrs := stubBackendsToResolverAddrs (backends )
515+
516+ ctx , cancel := context .WithTimeout (context .Background (), defaultTestTimeout )
517+ defer cancel ()
518+
519+ // Push an update with both addresses and shuffling disabled. We should
520+ // connect to backend 0.
521+ activeCfg = noShuffleConfig
522+ resolverState := resolver.State {Addresses : addrs }
523+ r .UpdateState (resolverState )
524+ if err := pickfirst .CheckRPCsToBackend (ctx , cc , addrs [0 ]); err != nil {
525+ t .Fatal (err )
526+ }
527+
528+ // Send a config with shuffling enabled. This will reverse the addresses,
529+ // but the channel should still be connected to backend 0.
530+ activeCfg = shuffleConfig
531+ r .UpdateState (resolverState )
532+ if err := pickfirst .CheckRPCsToBackend (ctx , cc , addrs [0 ]); err != nil {
533+ t .Fatal (err )
534+ }
535+
536+ // Send a resolver update with no addresses. This should push the channel
537+ // into TransientFailure.
538+ r .UpdateState (resolver.State {})
539+ testutils .AwaitState (ctx , t , cc , connectivity .TransientFailure )
540+
541+ // Send the same config as last time with shuffling enabled. Since we are
542+ // not connected to backend 0, we should connect to backend 1.
543+ r .UpdateState (resolverState )
544+ if err := pickfirst .CheckRPCsToBackend (ctx , cc , addrs [1 ]); err != nil {
545+ t .Fatal (err )
546+ }
547+ }
548+
466549// Test config parsing with the env var turned on and off for various scenarios.
467550func (s ) TestPickFirst_ParseConfig_Success (t * testing.T ) {
468551 // Install a shuffler that always reverses two entries.
0 commit comments