diff --git a/internal/balancergroup/balancergroup_test.go b/internal/balancergroup/balancergroup_test.go index 9dfad1a76b72..5d3029c882a8 100644 --- a/internal/balancergroup/balancergroup_test.go +++ b/internal/balancergroup/balancergroup_test.go @@ -484,6 +484,67 @@ func (s) TestBalancerGroupBuildOptions(t *testing.T) { } } +func (s) TestBalancerGroup_UpdateClientConnState_AfterClose(t *testing.T) { + balancerName := t.Name() + clientConnStateCh := make(chan struct{}, 1) + + stub.Register(balancerName, stub.BalancerFuncs{ + UpdateClientConnState: func(_ *stub.BalancerData, _ balancer.ClientConnState) error { + clientConnStateCh <- struct{}{} + return nil + }, + }) + + bg := New(Options{ + CC: testutils.NewBalancerClientConn(t), + BuildOpts: balancer.BuildOptions{}, + StateAggregator: nil, + Logger: nil, + }) + + bg.Add(testBalancerIDs[0], balancer.Get(balancerName)) + bg.Close() + + if err := bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{}); err != nil { + t.Fatalf("Expected nil error, got %v", err) + } + + select { + case <-clientConnStateCh: + t.Fatalf("UpdateClientConnState was called after BalancerGroup was closed") + case <-time.After(defaultTestShortTimeout): + } +} + +func (s) TestBalancerGroup_ResolverError_AfterClose(t *testing.T) { + balancerName := t.Name() + resolveErrorCh := make(chan struct{}, 1) + + stub.Register(balancerName, stub.BalancerFuncs{ + ResolverError: func(_ *stub.BalancerData, _ error) { + resolveErrorCh <- struct{}{} + }, + }) + + bg := New(Options{ + CC: testutils.NewBalancerClientConn(t), + BuildOpts: balancer.BuildOptions{}, + StateAggregator: nil, + Logger: nil, + }) + + bg.Add(testBalancerIDs[0], balancer.Get(balancerName)) + bg.Close() + + bg.ResolverError(errors.New("test error")) + + select { + case <-resolveErrorCh: + t.Fatalf("ResolverError was called on sub-balancer after BalancerGroup was closed") + case <-time.After(defaultTestShortTimeout): + } +} + func (s) TestBalancerExitIdleOne(t *testing.T) { const balancerName = "stub-balancer-test-balancergroup-exit-idle-one" exitIdleCh := make(chan struct{}, 1) @@ -505,7 +566,7 @@ func (s) TestBalancerExitIdleOne(t *testing.T) { builder := balancer.Get(balancerName) bg.Add(testBalancerIDs[0], builder) - // Call ExitIdle on the child policy. + // Call ExitIdleOne on the child policy. bg.ExitIdleOne(testBalancerIDs[0]) select { case <-time.After(time.Second): @@ -514,6 +575,62 @@ func (s) TestBalancerExitIdleOne(t *testing.T) { } } +func (s) TestBalancerGroup_ExitIdleOne_AfterClose(t *testing.T) { + balancerName := t.Name() + exitIdleCh := make(chan struct{}) + + stub.Register(balancerName, stub.BalancerFuncs{ + ExitIdle: func(_ *stub.BalancerData) { + close(exitIdleCh) + }, + }) + + bg := New(Options{ + CC: testutils.NewBalancerClientConn(t), + BuildOpts: balancer.BuildOptions{}, + StateAggregator: nil, + Logger: nil, + }) + + bg.Add(testBalancerIDs[0], balancer.Get(balancerName)) + bg.Close() + bg.ExitIdleOne(testBalancerIDs[0]) + + select { + case <-time.After(defaultTestShortTimeout): + case <-exitIdleCh: + t.Fatalf("ExitIdleOne called ExitIdle on sub-balancer after BalancerGroup was closed") + } +} + +func (s) TestBalancerGroup_ExitIdleOne_NonExistentID(t *testing.T) { + balancerName := t.Name() + exitIdleCh := make(chan struct{}, 1) + + stub.Register(balancerName, stub.BalancerFuncs{ + ExitIdle: func(_ *stub.BalancerData) { + exitIdleCh <- struct{}{} + }, + }) + + bg := New(Options{ + CC: testutils.NewBalancerClientConn(t), + BuildOpts: balancer.BuildOptions{}, + StateAggregator: nil, + Logger: nil, + }) + defer bg.Close() + + bg.Add(testBalancerIDs[0], balancer.Get(balancerName)) + bg.ExitIdleOne("non-existent-id") + + select { + case <-time.After(defaultTestShortTimeout): + case <-exitIdleCh: + t.Fatalf("ExitIdleOne called ExitIdle on wrong sub-balancer ID") + } +} + // TestBalancerGracefulSwitch tests the graceful switch functionality for a // child of the balancer group. At first, the child is configured as a round // robin load balancer, and thus should behave accordingly. The test then @@ -639,3 +756,92 @@ func (s) TestBalancerGracefulSwitch(t *testing.T) { } } } + +func (s) TestBalancerExitIdle_All(t *testing.T) { + balancer1 := t.Name() + "-1" + balancer2 := t.Name() + "-2" + + testID1, testID2 := testBalancerIDs[0], testBalancerIDs[1] + + exitIdleCh1, exitIdleCh2 := make(chan struct{}, 1), make(chan struct{}, 1) + + stub.Register(balancer1, stub.BalancerFuncs{ + ExitIdle: func(_ *stub.BalancerData) { + exitIdleCh1 <- struct{}{} + }, + }) + + stub.Register(balancer2, stub.BalancerFuncs{ + ExitIdle: func(_ *stub.BalancerData) { + exitIdleCh2 <- struct{}{} + }, + }) + + cc := testutils.NewBalancerClientConn(t) + bg := New(Options{ + CC: cc, + BuildOpts: balancer.BuildOptions{}, + StateAggregator: nil, + Logger: nil, + }) + defer bg.Close() + + bg.Add(testID1, balancer.Get(balancer1)) + bg.Add(testID2, balancer.Get(balancer2)) + + bg.ExitIdle() + + errCh := make(chan error, 2) + + go func() { + select { + case <-exitIdleCh1: + errCh <- nil + case <-time.After(defaultTestTimeout): + errCh <- fmt.Errorf("timeout waiting for ExitIdle on balancer1") + } + }() + + go func() { + select { + case <-exitIdleCh2: + errCh <- nil + case <-time.After(defaultTestTimeout): + errCh <- fmt.Errorf("timeout waiting for ExitIdle on balancer2") + } + }() + + for i := 0; i < 2; i++ { + if err := <-errCh; err != nil { + t.Fatal(err) + } + } +} + +func (s) TestBalancerGroup_ExitIdle_AfterClose(t *testing.T) { + balancerName := t.Name() + exitIdleCh := make(chan struct{}, 1) + + stub.Register(balancerName, stub.BalancerFuncs{ + ExitIdle: func(_ *stub.BalancerData) { + exitIdleCh <- struct{}{} + }, + }) + + bg := New(Options{ + CC: testutils.NewBalancerClientConn(t), + BuildOpts: balancer.BuildOptions{}, + StateAggregator: nil, + Logger: nil, + }) + + bg.Add(testBalancerIDs[0], balancer.Get(balancerName)) + bg.Close() + bg.ExitIdle() + + select { + case <-exitIdleCh: + t.Fatalf("ExitIdle was called on sub-balancer even after BalancerGroup was closed") + case <-time.After(defaultTestShortTimeout): + } +}