Repository: grpc/grpc-go

Issue:
pickfirst: Don't use internal.ShuffleAddressListForTesting directly
Based on the PR review comment https://github.com/grpc/grpc-go/pull/7498#discussion_r1718899646, 

* `pickfirst` sets `ShuffleAddressListForTesting` to use `rand.Shuffle` in its init function
* `pickfirst` also directly uses `ShuffleAddressListForTesting` in `UpdateClientConnState` to get the actual shuffle function

So, for someone reading the code, it can be a little misleading as to why `pickfirst` is getting the shuffle func from a var with a `ForTesting` suffix.

The suggestion from the comment is to instead do something like:
```go
package pickfirst

var shuffleFunc func(n int, swap func(i, j int))

func init() {
  shuffleFunc = rand.Shuffle
  internal.ShuffleAddressListForTesting	= func(sf func(n int, swap func(i, j int))) {
    shuffleFunc = sf
  }  
}

func (...) UpdateClientConnState(...) {
  // Directly use `shuffleFunc` to do the shuffling
}
```

> Note that I've changed the signature of `internal.ShuffleAddressListForTesting` to accept a func. We could rename it to something like `internal.PickfirstShuffleFunc`. Also, note that we could completely move this var to the `pickfirst` package and have our tests directly use that, but would pollute the external API.

Hints from issue comments:
@arjan-bal: Waiting for for https://github.com/grpc/grpc-go/pull/7498 to get merged so that I can make the change  in both `pickfirsts` together, avoiding merge conflicts.

Solution patch:
diff --git a/balancer/pickfirst/internal/internal.go b/balancer/pickfirst/internal/internal.go
new file mode 100644
index 00000000..c5197894
--- /dev/null
+++ b/balancer/pickfirst/internal/internal.go
@@ -0,0 +1,24 @@
+/*
+ * Copyright 2024 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package internal contains code internal to the pickfirst package.
+package internal
+
+import "math/rand"
+
+// RandShuffle pseudo-randomizes the order of addresses.
+var RandShuffle = rand.Shuffle
diff --git a/balancer/pickfirst/pickfirst.go b/balancer/pickfirst/pickfirst.go
index 8fadf3df..3e792b2b 100644
--- a/balancer/pickfirst/pickfirst.go
+++ b/balancer/pickfirst/pickfirst.go
@@ -26,9 +26,9 @@ import (
 	"math/rand"
 
 	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/pickfirst/internal"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/grpclog"
-	"google.golang.org/grpc/internal"
 	internalgrpclog "google.golang.org/grpc/internal/grpclog"
 	"google.golang.org/grpc/internal/pretty"
 	"google.golang.org/grpc/resolver"
@@ -37,7 +37,6 @@ import (
 
 func init() {
 	balancer.Register(pickfirstBuilder{})
-	internal.ShuffleAddressListForTesting = func(n int, swap func(i, j int)) { rand.Shuffle(n, swap) }
 }
 
 var logger = grpclog.Component("pick-first-lb")
@@ -143,7 +142,7 @@ func (b *pickfirstBalancer) UpdateClientConnState(state balancer.ClientConnState
 		// within each endpoint. - A61
 		if cfg.ShuffleAddressList {
 			endpoints = append([]resolver.Endpoint{}, endpoints...)
-			internal.ShuffleAddressListForTesting.(func(int, func(int, int)))(len(endpoints), func(i, j int) { endpoints[i], endpoints[j] = endpoints[j], endpoints[i] })
+			internal.RandShuffle(len(endpoints), func(i, j int) { endpoints[i], endpoints[j] = endpoints[j], endpoints[i] })
 		}
 
 		// "Flatten the list by concatenating the ordered list of addresses for each
diff --git a/internal/internal.go b/internal/internal.go
index 5510b2fb..20b4dc3d 100644
--- a/internal/internal.go
+++ b/internal/internal.go
@@ -207,10 +207,6 @@ var (
 	// default resolver scheme.
 	UserSetDefaultScheme = false
 
-	// ShuffleAddressListForTesting pseudo-randomizes the order of addresses.  n
-	// is the number of elements.  swap swaps the elements with indexes i and j.
-	ShuffleAddressListForTesting any // func(n int, swap func(i, j int))
-
 	// ConnectedAddress returns the connected address for a SubConnState. The
 	// address is only valid if the state is READY.
 	ConnectedAddress any // func (scs SubConnState) resolver.Address
diff --git a/internal/xds/bootstrap/bootstrap.go b/internal/xds/bootstrap/bootstrap.go
index 35aeea70..c725bc1e 100644
--- a/internal/xds/bootstrap/bootstrap.go
+++ b/internal/xds/bootstrap/bootstrap.go
@@ -220,14 +220,20 @@ func (sc *ServerConfig) ServerFeaturesIgnoreResourceDeletion() bool {
 	return false
 }
 
-// DialOptions returns a slice of all the configured dial options for this
-// server.
-func (sc *ServerConfig) DialOptions() []grpc.DialOption {
-	dopts := []grpc.DialOption{sc.credsDialOption}
-	if sc.dialerOption != nil {
-		dopts = append(dopts, sc.dialerOption)
-	}
-	return dopts
+// CredsDialOption returns the first supported transport credentials from the
+// configuration, as a dial option.
+func (sc *ServerConfig) CredsDialOption() grpc.DialOption {
+	return sc.credsDialOption
+}
+
+// DialerOption returns the Dialer function that specifies how to dial the xDS
+// server determined by the first supported credentials from the configuration,
+// as a dial option.
+//
+// TODO(https://github.com/grpc/grpc-go/issues/7661): change ServerConfig type
+// to have a single method that returns all configured dial options.
+func (sc *ServerConfig) DialerOption() grpc.DialOption {
+	return sc.dialerOption
 }
 
 // Cleanups returns a collection of functions to be called when the xDS client
diff --git a/xds/internal/xdsclient/transport/transport.go b/xds/internal/xdsclient/transport/transport.go
index 59b22172..134a9519 100644
--- a/xds/internal/xdsclient/transport/transport.go
+++ b/xds/internal/xdsclient/transport/transport.go
@@ -192,14 +192,19 @@ func New(opts Options) (*Transport, error) {
 		return nil, errors.New("missing OnSend callback handler when creating a new transport")
 	}
 
-	// Dial the xDS management server with dial options specified by the server
-	// configuration and a static keepalive configuration that is common across
-	// gRPC language implementations.
-	kpCfg := grpc.WithKeepaliveParams(keepalive.ClientParameters{
-		Time:    5 * time.Minute,
-		Timeout: 20 * time.Second,
-	})
-	dopts := append([]grpc.DialOption{kpCfg}, opts.ServerCfg.DialOptions()...)
+	// Dial the xDS management with the passed in credentials.
+	dopts := []grpc.DialOption{
+		opts.ServerCfg.CredsDialOption(),
+		grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			// We decided to use these sane defaults in all languages, and
+			// kicked the can down the road as far making these configurable.
+			Time:    5 * time.Minute,
+			Timeout: 20 * time.Second,
+		}),
+	}
+	if dialerOpts := opts.ServerCfg.DialerOption(); dialerOpts != nil {
+		dopts = append(dopts, dialerOpts)
+	}
 	grpcNewClient := transportinternal.GRPCNewClient.(func(string, ...grpc.DialOption) (*grpc.ClientConn, error))
 	cc, err := grpcNewClient(opts.ServerCfg.ServerURI(), dopts...)
 	if err != nil {



Test patch:
diff --git a/test/pickfirst_test.go b/balancer/pickfirst/pickfirst_ext_test.go
similarity index 96%
rename from test/pickfirst_test.go
rename to balancer/pickfirst/pickfirst_ext_test.go
index a5fc77d4..faa6e7dd 100644
--- a/test/pickfirst_test.go
+++ b/balancer/pickfirst/pickfirst_ext_test.go
@@ -16,7 +16,7 @@
  *
  */
 
-package test
+package pickfirst_test
 
 import (
 	"context"
@@ -28,11 +28,13 @@ import (
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/backoff"
+	pfinternal "google.golang.org/grpc/balancer/pickfirst/internal"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/internal/channelz"
+	"google.golang.org/grpc/internal/grpctest"
 	"google.golang.org/grpc/internal/stubserver"
 	"google.golang.org/grpc/internal/testutils"
 	"google.golang.org/grpc/internal/testutils/pickfirst"
@@ -45,7 +47,38 @@ import (
 	testpb "google.golang.org/grpc/interop/grpc_testing"
 )
 
-const pickFirstServiceConfig = `{"loadBalancingConfig": [{"pick_first":{}}]}`
+const (
+	pickFirstServiceConfig = `{"loadBalancingConfig": [{"pick_first":{}}]}`
+	// Default timeout for tests in this package.
+	defaultTestTimeout = 10 * time.Second
+	// Default short timeout, to be used when waiting for events which are not
+	// expected to happen.
+	defaultTestShortTimeout = 100 * time.Millisecond
+)
+
+func init() {
+	channelz.TurnOn()
+}
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// parseServiceConfig is a test helper which uses the manual resolver to parse
+// the given service config. It calls t.Fatal() if service config parsing fails.
+func parseServiceConfig(t *testing.T, r *manual.Resolver, sc string) *serviceconfig.ParseResult {
+	t.Helper()
+
+	scpr := r.CC.ParseServiceConfig(sc)
+	if scpr.Err != nil {
+		t.Fatalf("Failed to parse service config %q: %v", sc, scpr.Err)
+	}
+	return scpr
+}
 
 // setupPickFirst performs steps required for pick_first tests. It starts a
 // bunch of backends exporting the TestService, creates a ClientConn to them
@@ -377,16 +410,15 @@ func (s) TestPickFirst_ShuffleAddressList(t *testing.T) {
 	const serviceConfig = `{"loadBalancingConfig": [{"pick_first":{ "shuffleAddressList": true }}]}`
 
 	// Install a shuffler that always reverses two entries.
-	origShuf := internal.ShuffleAddressListForTesting
-	defer func() { internal.ShuffleAddressListForTesting = origShuf }()
-	internal.ShuffleAddressListForTesting = func(n int, f func(int, int)) {
+	origShuf := pfinternal.RandShuffle
+	defer func() { pfinternal.RandShuffle = origShuf }()
+	pfinternal.RandShuffle = func(n int, f func(int, int)) {
 		if n != 2 {
 			t.Errorf("Shuffle called with n=%v; want 2", n)
 			return
 		}
 		f(0, 1) // reverse the two addresses
 	}
-
 	// Set up our backends.
 	cc, r, backends := setupPickFirst(t, 2)
 	addrs := stubBackendsToResolverAddrs(backends)
@@ -434,9 +466,9 @@ func (s) TestPickFirst_ShuffleAddressList(t *testing.T) {
 // Test config parsing with the env var turned on and off for various scenarios.
 func (s) TestPickFirst_ParseConfig_Success(t *testing.T) {
 	// Install a shuffler that always reverses two entries.
-	origShuf := internal.ShuffleAddressListForTesting
-	defer func() { internal.ShuffleAddressListForTesting = origShuf }()
-	internal.ShuffleAddressListForTesting = func(n int, f func(int, int)) {
+	origShuf := pfinternal.RandShuffle
+	defer func() { pfinternal.RandShuffle = origShuf }()
+	pfinternal.RandShuffle = func(n int, f func(int, int)) {
 		if n != 2 {
 			t.Errorf("Shuffle called with n=%v; want 2", n)
 			return
diff --git a/test/balancer_switching_test.go b/test/balancer_switching_test.go
index f0bb600f..8074b59b 100644
--- a/test/balancer_switching_test.go
+++ b/test/balancer_switching_test.go
@@ -46,6 +46,7 @@ const (
 	loadBalancedServicePort = 443
 	wantGRPCLBTraceDesc     = `Channel switches to new LB policy "grpclb"`
 	wantRoundRobinTraceDesc = `Channel switches to new LB policy "round_robin"`
+	pickFirstServiceConfig  = `{"loadBalancingConfig": [{"pick_first":{}}]}`
 
 	// This is the number of stub backends set up at the start of each test. The
 	// first backend is used for the "grpclb" policy and the rest are used for
@@ -53,6 +54,16 @@ const (
 	backendCount = 3
 )
 
+// stubBackendsToResolverAddrs converts from a set of stub server backends to
+// resolver addresses. Useful when pushing addresses to the manual resolver.
+func stubBackendsToResolverAddrs(backends []*stubserver.StubServer) []resolver.Address {
+	addrs := make([]resolver.Address, len(backends))
+	for i, backend := range backends {
+		addrs[i] = resolver.Address{Addr: backend.Address}
+	}
+	return addrs
+}
+
 // setupBackendsAndFakeGRPCLB sets up backendCount number of stub server
 // backends and a fake grpclb server for tests which exercise balancer switch
 // scenarios involving grpclb.
diff --git a/xds/internal/xdsclient/transport/transport_test.go b/xds/internal/xdsclient/transport/transport_test.go
index b217fd74..7aac0ccd 100644
--- a/xds/internal/xdsclient/transport/transport_test.go
+++ b/xds/internal/xdsclient/transport/transport_test.go
@@ -22,7 +22,6 @@ import (
 	"encoding/json"
 	"net"
 	"testing"
-	"time"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
@@ -44,8 +43,6 @@ func Test(t *testing.T) {
 	grpctest.RunSubTests(t, s{})
 }
 
-const defaultTestTimeout = 10 * time.Second
-
 var noopRecvHandler = func(_ transport.ResourceUpdate, onDone func()) error {
 	onDone()
 	return nil
@@ -111,17 +108,10 @@ const testDialerCredsBuilderName = "test_dialer_creds"
 // testDialerCredsBuilder implements the `Credentials` interface defined in
 // package `xds/bootstrap` and encapsulates an insecure credential with a
 // custom Dialer that specifies how to dial the xDS server.
-type testDialerCredsBuilder struct {
-	// Closed with the custom Dialer is invoked.
-	// Needs to be passed in by the test.
-	dialCalled chan struct{}
-}
+type testDialerCredsBuilder struct{}
 
 func (t *testDialerCredsBuilder) Build(json.RawMessage) (credentials.Bundle, func(), error) {
-	return &testDialerCredsBundle{
-		Bundle:     insecure.NewBundle(),
-		dialCalled: t.dialCalled,
-	}, func() {}, nil
+	return &testDialerCredsBundle{insecure.NewBundle()}, func() {}, nil
 }
 
 func (t *testDialerCredsBuilder) Name() string {
@@ -133,12 +123,10 @@ func (t *testDialerCredsBuilder) Name() string {
 // that specifies how to dial the xDS server.
 type testDialerCredsBundle struct {
 	credentials.Bundle
-	dialCalled chan struct{}
 }
 
-func (t *testDialerCredsBundle) Dialer(_ context.Context, address string) (net.Conn, error) {
-	close(t.dialCalled)
-	return net.Dial("tcp", address)
+func (t *testDialerCredsBundle) Dialer(context.Context, string) (net.Conn, error) {
+	return nil, nil
 }
 
 func (s) TestNewWithDialerFromCredentialsBundle(t *testing.T) {
@@ -152,8 +140,7 @@ func (s) TestNewWithDialerFromCredentialsBundle(t *testing.T) {
 	internal.GRPCNewClient = customGRPCNewClient
 	defer func() { internal.GRPCNewClient = oldGRPCNewClient }()
 
-	dialCalled := make(chan struct{})
-	bootstrap.RegisterCredentials(&testDialerCredsBuilder{dialCalled: dialCalled})
+	bootstrap.RegisterCredentials(&testDialerCredsBuilder{})
 	serverCfg, err := internalbootstrap.ServerConfigForTesting(internalbootstrap.ServerConfigTestingOptions{
 		URI:          "trafficdirector.googleapis.com:443",
 		ChannelCreds: []internalbootstrap.ChannelCreds{{Type: testDialerCredsBuilderName}},
@@ -161,7 +148,9 @@ func (s) TestNewWithDialerFromCredentialsBundle(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Failed to create server config for testing: %v", err)
 	}
-
+	if serverCfg.DialerOption() == nil {
+		t.Fatalf("Dialer for xDS transport in server config for testing is nil, want non-nil")
+	}
 	// Create a new transport.
 	opts := transport.Options{
 		ServerCfg: serverCfg,
@@ -182,11 +171,6 @@ func (s) TestNewWithDialerFromCredentialsBundle(t *testing.T) {
 	if err != nil {
 		t.Fatalf("transport.New(%v) failed: %v", opts, err)
 	}
-	select {
-	case <-dialCalled:
-	case <-time.After(defaultTestTimeout):
-		t.Fatal("Timeout when waiting for Dialer() to be invoked")
-	}
 	// Verify there are three dial options passed to the custom grpc.NewClient.
 	// The first is opts.ServerCfg.CredsDialOption(), the second is
 	// grpc.WithKeepaliveParams(), and the third is opts.ServerCfg.DialerOption()