Repository: grpc/grpc-go
File: xds/internal/balancer/clusterresolver/e2e_test/eds_impl_test.go (Go)

Code before:
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			origDualstackEndpointsEnabled := envconfig.XDSDualstackEndpointsEnabled
			defer func() {
				envconfig.XDSDualstackEndpointsEnabled = origDualstackEndpointsEnabled
			}()
			envconfig.XDSDualstackEndpointsEnabled = tc.dualstackEndpointsEnabled

			// Wrap the round robin balancer to intercept resolver updates.
			originalRRBuilder := balancer.Get(roundrobin.Name)
			defer func() {
				balancer.Register(originalRRBuilder)
			}()
			resolverStateMu := sync.Mutex{}
			resolverState := resolver.State{}
			stub.Register(roundrobin.Name, stub.BalancerFuncs{
				Init: func(bd *stub.BalancerData) {
					bd.Data = originalRRBuilder.Build(bd.ClientConn, bd.BuildOptions)
				},
				Close: func(bd *stub.BalancerData) {
					bd.Data.(balancer.Balancer).Close()
				},
				UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
					resolverStateMu.Lock()
					defer resolverStateMu.Unlock()
					resolverState = ccs.ResolverState
					return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
				},
			})

			// Spin up a management server to receive xDS resources from.
			mgmtServer := e2e.StartManagementServer(t, e2e.ManagementServerOptions{})

			// Create bootstrap configuration pointing to the above management server.
			nodeID := uuid.New().String()
			bootstrapContents := e2e.DefaultBootstrapContents(t, nodeID, mgmtServer.Address)

			// Create xDS resources for consumption by the test. We start off with a
			// single backend in a single EDS locality.
			resources := clientEndpointsResource(nodeID, edsServiceName, []e2e.LocalityOptions{{
				Name:   localityName1,
				Weight: 1,
				Backends: []e2e.BackendOptions{{
					Ports: ports,
				}},
			}})
			if err := mgmtServer.Update(ctx, resources); err != nil {
				t.Fatal(err)

Review comment (suggestion):
I think this would be more concise and express the same thing if it used an `atomic.Pointer` instead of a separate mutex.

Code after:
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			origDualstackEndpointsEnabled := envconfig.XDSDualstackEndpointsEnabled
			defer func() {
				envconfig.XDSDualstackEndpointsEnabled = origDualstackEndpointsEnabled
			}()
			envconfig.XDSDualstackEndpointsEnabled = tc.dualstackEndpointsEnabled

			// Wrap the round robin balancer to intercept resolver updates.
			originalRRBuilder := balancer.Get(roundrobin.Name)
			defer func() {
				balancer.Register(originalRRBuilder)
			}()
			resolverState := atomic.Pointer[resolver.State]{}
			resolverState.Store(&resolver.State{})
			stub.Register(roundrobin.Name, stub.BalancerFuncs{
				Init: func(bd *stub.BalancerData) {
					bd.Data = originalRRBuilder.Build(bd.ClientConn, bd.BuildOptions)
				},
				Close: func(bd *stub.BalancerData) {
					bd.Data.(balancer.Balancer).Close()
				},
				UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
					resolverState.Store(&ccs.ResolverState)
					return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
				},
			})

			// Spin up a management server to receive xDS resources from.
			mgmtServer := e2e.StartManagementServer(t, e2e.ManagementServerOptions{})

			// Create bootstrap configuration pointing to the above management server.
			nodeID := uuid.New().String()
			bootstrapContents := e2e.DefaultBootstrapContents(t, nodeID, mgmtServer.Address)

			// Create xDS resources for consumption by the test. We start off with a
			// single backend in a single EDS locality.
			resources := clientEndpointsResource(nodeID, edsServiceName, []e2e.LocalityOptions{{
				Name:   localityName1,
				Weight: 1,
				Backends: []e2e.BackendOptions{{
					Ports: ports,
				}},
			}})
			if err := mgmtServer.Update(ctx, resources); err != nil {
				t.Fatal(err)
			}