Repository: grpc/grpc-go
File: balancer/pickfirst/pickfirstleaf/pickfirstleaf_ext_test.go (Go)

Code before:
		}
	}
}

// resolverAddrs  returns a list of resolver addresses for the stub server
// backends. Useful when pushing addresses to the manual resolver.
func (b *backendManager) resolverAddrs() []resolver.Address {
	addrs := make([]resolver.Address, len(b.backends))
	for i, backend := range b.backends {
		addrs[i] = resolver.Address{Addr: backend.Address}
	}
	return addrs
}

type ccStateSubscriber struct {
	transitions []connectivity.State
}

func (c *ccStateSubscriber) OnMessage(msg any) {
	c.transitions = append(c.transitions, msg.(connectivity.State))
}

// hangingServer is a server that accept a TCP connection and remains idle until
// asked to close or respond to the connection. They can be used to control
// how long it takes for a subchannel to report a TRANSIENT_FAILURE in tests.
type hangingServer struct {
	addr      resolver.Address
	listener  net.Listener
	closeConn func()
	contacted *grpcsync.Event
	readyChan chan struct{}
}

func newHangingServerGroup(t *testing.T, count int) []*hangingServer {
	servers := []*hangingServer{}
	for i := 0; i < count; i++ {
		lis, err := net.Listen("tcp", "localhost:0")
		if err != nil {
			t.Fatalf("Error while listening. Err: %v", err)
		}
		closeChan := make(chan struct{})
		contacted := grpcsync.NewEvent()
		readyChan := make(chan struct{})

		go func() {
			conn, err := lis.Accept()
			if err != nil {
				t.Error(err)
				return
			}
			defer conn.Close()

Review comment (bug):
We have a bunch of utilities in the `testutils` package that can be used I feel, instead of implementing this new hangingServer.

There is a blocking dialer that lets you block, fail or succeed connection attempts: https://github.com/grpc/grpc-go/blob/ec10e73f02a38fcf1e4f6e3cdee5f73c89bfa4ac/internal/testutils/blocking_context_dialer.go#L33
There is a restartable listener: https://github.com/grpc/grpc-go/blob/master/internal/testutils/restartable_listener.go#L37
There is a wrapped listener that returns a wrapped conn: https://github.com/grpc/grpc-go/blob/master/internal/testutils/wrappers.go

Code after:
	b.Balancer.Close()
}

func (b *stateStoringBalancer) ExitIdle() {
	if ib, ok := b.Balancer.(balancer.ExitIdler); ok {
		ib.ExitIdle()
	}
}

type stateStoringBalancerBuilder struct {
	balancer chan *stateStoringBalancer
}

func (b *stateStoringBalancerBuilder) Name() string {
	return stateStoringBalancerName
}

func (b *stateStoringBalancerBuilder) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
	bal := &stateStoringBalancer{}
	bal.Balancer = balancer.Get(pickfirstleaf.Name).Build(&stateStoringCCWrapper{cc, bal}, opts)
	b.balancer <- bal
	return bal
}

func (b *stateStoringBalancer) subConnStates() []scState {
	b.mu.Lock()
	defer b.mu.Unlock()
	ret := []scState{}
	for _, s := range b.scStates {
		ret = append(ret, *s)
	}
	return ret
}

func (b *stateStoringBalancer) addSCState(state *scState) {
	b.mu.Lock()
	b.scStates = append(b.scStates, state)
	b.mu.Unlock()
}

type stateStoringCCWrapper struct {
	balancer.ClientConn
	b *stateStoringBalancer
}

func (ccw *stateStoringCCWrapper) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
	oldListener := opts.StateListener
	scs := &scState{
		State: connectivity.Idle,
		Addrs: addrs,
	}