@@ -20,6 +20,7 @@ package pickfirst_test
20
20
21
21
import (
22
22
"context"
23
+ "encoding/json"
23
24
"errors"
24
25
"fmt"
25
26
"strings"
@@ -28,11 +29,14 @@ import (
28
29
29
30
"google.golang.org/grpc"
30
31
"google.golang.org/grpc/backoff"
32
+ "google.golang.org/grpc/balancer"
33
+ pfbalancer "google.golang.org/grpc/balancer/pickfirst"
31
34
pfinternal "google.golang.org/grpc/balancer/pickfirst/internal"
32
35
"google.golang.org/grpc/codes"
33
36
"google.golang.org/grpc/connectivity"
34
37
"google.golang.org/grpc/credentials/insecure"
35
38
"google.golang.org/grpc/internal"
39
+ "google.golang.org/grpc/internal/balancer/stub"
36
40
"google.golang.org/grpc/internal/channelz"
37
41
"google.golang.org/grpc/internal/grpctest"
38
42
"google.golang.org/grpc/internal/stubserver"
@@ -463,6 +467,85 @@ func (s) TestPickFirst_ShuffleAddressList(t *testing.T) {
463
467
}
464
468
}
465
469
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
+
466
549
// Test config parsing with the env var turned on and off for various scenarios.
467
550
func (s ) TestPickFirst_ParseConfig_Success (t * testing.T ) {
468
551
// Install a shuffler that always reverses two entries.
0 commit comments