feat(bigtable): prime initial channel pool connections in parallel (#20010)

## Summary
- `NewBigtableChannelPool` previously primed each connection serially,
so a pool of N connections waited `N * (dial + Prime + retry-backoff)`
before returning. With Prime() on the order of tens of ms per
connection, that dominates client startup time and scales linearly with
pool size.
- This PR fans the prime calls out across an `errgroup` with
`SetLimit(min(N, 10))`. The cap is intentional: very large pools should
not spawn one dial+Prime goroutine per connection; small pools naturally
use `connPoolSize` workers.
- Failure semantics are preserved: the first prime error cancels the
rest via the derived context, and the existing cleanup loop closes any
populated entries before returning the error.
- The Direct Access "firstConn" path still seats the primed probe
connection at index 0 and skips a worker for that slot.

## Test plan
- [x] `go build ./...`
- [x] `go vet ./...`
- [x] `go test -race -timeout 240s ./bigtable/internal/transport/...`
- [x] New `TestNewBigtableChannelPoolParallelPriming` covers pool sizes
below, at, and above the 10-worker cap. Asserts peak concurrent
in-flight `PingAndWarm` matches expected worker count, and wall time
beats the sequential bound (`N * primeDelay`).
- [x] `fakeService` now tracks peak in-flight `PingAndWarm` so the test
can assert concurrency without timing-only signals.
- [x] `TestNewBigtableChannelPool/DialFailure` updated to use
`atomic.Int64` for its counter (the dial function is now called
concurrently).

Author: sushanb

bigtable/go.mod

@@ -23,6 +23,7 @@ require (
 	go.opentelemetry.io/otel/sdk v1.44.0
 	go.opentelemetry.io/otel/sdk/metric v1.44.0
 	golang.org/x/oauth2 v0.36.0
+	golang.org/x/sync v0.20.0
 	google.golang.org/api v0.283.0
 	google.golang.org/genproto v0.0.0-20260519071638-aa98bba5eb94
 	google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa
@@ -55,7 +56,6 @@ require (
 	go.opentelemetry.io/otel/trace v1.44.0 // indirect
 	golang.org/x/crypto v0.51.0 // indirect
 	golang.org/x/net v0.55.0 // indirect
-	golang.org/x/sync v0.20.0 // indirect
 	golang.org/x/sys v0.45.0 // indirect
 	golang.org/x/text v0.37.0 // indirect
 	golang.org/x/time v0.15.0 // indirect

bigtable/internal/transport/connpool.go

@@ -16,7 +16,6 @@ package internal
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"log"
 	"math"
@@ -36,6 +35,7 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
 	"golang.org/x/oauth2/google"
+	"golang.org/x/sync/errgroup"
 	gtransport "google.golang.org/api/transport/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/alts"
@@ -61,6 +61,11 @@ const (
 	artificialLoadIfError        = 10
 	artificialLoadPenalizedTimer = 5 * time.Second
 	requestParamsHeader          = "x-goog-request-params"
+	// maxPrimeWorkers caps the goroutines used to prime initial pool
+	// connections in parallel. Pools smaller than this naturally fan out to
+	// connPoolSize workers; larger pools cap here so we don't spawn one
+	// dial+Prime goroutine per connection.
+	maxPrimeWorkers = 10
 )
 
 // ipProtocol represents the type of IP protocol used.
@@ -502,45 +507,23 @@ func NewBigtableChannelPool(ctx context.Context, connPoolSize int, strategy btop
 		pool.selectFunc = pool.selectRoundRobin
 	}
 
-	var exitSignal error
 	btopt.Debugf(pool.logger, "bigtable_connpool: Creating conn pool with %d connections", connPoolSize)
 	// TODO: Replace this logic with addConnections(...).
 	initialConns := make([]*connEntry, connPoolSize)
-	for i := 0; i < connPoolSize; i++ {
-		select {
-		case <-pool.poolCtx.Done():
-			exitSignal = errors.New("bigtable_connpool: pool context canceled")
-		default:
-		}
-
-		if exitSignal != nil {
-			break
-		}
-
-		var entry *connEntry
-		var err error
-
-		if i == 0 && firstConn != nil {
-			entry = &connEntry{conn: firstConn}
-		} else {
-			entry, err = pool.factory.newEntry(ctx)
-		}
-
-		if err != nil {
-			exitSignal = err
-			break
-		}
-		initialConns[i] = entry
+	primeStart := 0
+	if firstConn != nil {
+		initialConns[0] = &connEntry{conn: firstConn}
+		primeStart = 1
 	}
-	if exitSignal != nil {
-		btopt.Debugf(pool.logger, "bigtable_connpool: error during initial connection creation: %v\n", exitSignal)
-		// Close populated conns
+
+	if err := pool.primeInitialConns(pool.poolCtx, initialConns, primeStart); err != nil {
+		btopt.Debugf(pool.logger, "bigtable_connpool: error during initial connection creation: %v\n", err)
 		for _, entry := range initialConns {
 			if entry != nil && entry.conn != nil {
 				entry.conn.Close()
 			}
 		}
-		return nil, exitSignal
+		return nil, err
 	}
 
 	pool.conns.Store(&initialConns)
@@ -569,6 +552,43 @@ func NewBigtableChannelPool(ctx context.Context, connPoolSize int, strategy btop
 	return pool, nil
 }
 
+// primeInitialConns dials and primes the connections at indices [primeStart, len(out))
+// in parallel, capped at maxPrimeWorkers. Successful entries are written into out at
+// their target index; if any prime fails, the first error is returned and the caller is
+// responsible for closing any populated entries.
+//
+// ctx scopes the prime operations: the first worker error (or ctx cancellation) tears
+// down the rest via errgroup's derived context.
+func (p *BigtableChannelPool) primeInitialConns(ctx context.Context, out []*connEntry, primeStart int) error {
+	jobs := len(out) - primeStart
+	if jobs <= 0 {
+		return nil
+	}
+	if err := ctx.Err(); err != nil {
+		return fmt.Errorf("bigtable_connpool: pool context canceled: %w", err)
+	}
+
+	workers := jobs
+	if workers > maxPrimeWorkers {
+		workers = maxPrimeWorkers
+	}
+
+	g, gctx := errgroup.WithContext(ctx)
+	g.SetLimit(workers)
+	for i := primeStart; i < len(out); i++ {
+		idx := i
+		g.Go(func() error {
+			entry, err := p.factory.newEntry(gctx)
+			if err != nil {
+				return err
+			}
+			out[idx] = entry
+			return nil
+		})
+	}
+	return g.Wait()
+}
+
 // checkIfDirectAccessCompatible attempts to create a single connection using the directAccessDialer,
 // primes it, and checks if direct access was successful
 func (p *BigtableChannelPool) checkIfDirectAccessCompatible() (*BigtableConn, bool) {

bigtable/internal/transport/connpool_helper_test.go

@@ -36,6 +36,8 @@ type fakeService struct {
 	btpb.UnimplementedBigtableServer
 	mu                      sync.Mutex
 	pingCount               int
+	pingInflight            int // Currently in-flight PingAndWarm RPCs
+	pingPeakInflight        int // High-water mark for concurrent PingAndWarm RPCs
 	callCount               int
 	streamSema              chan struct{} // To control stream lifetime
 	delay                   time.Duration // To simulate work
@@ -122,6 +124,8 @@ func (f *fakeService) reset() {
 	defer f.mu.Unlock()
 	f.callCount = 0
 	f.pingCount = 0
+	f.pingInflight = 0
+	f.pingPeakInflight = 0
 	f.serverErr = nil
 	f.pingErrs = nil
 	f.delay = 0
@@ -139,6 +143,10 @@ func (s *fakeService) PingAndWarm(ctx context.Context, req *btpb.PingAndWarmRequ
 	s.mu.Lock()
 	callNum := s.pingCount
 	s.pingCount++
+	s.pingInflight++
+	if s.pingInflight > s.pingPeakInflight {
+		s.pingPeakInflight = s.pingInflight
+	}
 
 	var err error
 	if len(s.pingErrs) > 0 {
@@ -156,6 +164,13 @@ func (s *fakeService) PingAndWarm(ctx context.Context, req *btpb.PingAndWarmRequ
 		s.lastPingAndWarmMetadata, _ = metadata.FromIncomingContext(ctx)
 	}
 	s.mu.Unlock()
+
+	defer func() {
+		s.mu.Lock()
+		s.pingInflight--
+		s.mu.Unlock()
+	}()
+
 	if delay > 0 {
 		select {
 		case <-time.After(delay):
@@ -170,6 +185,12 @@ func (s *fakeService) PingAndWarm(ctx context.Context, req *btpb.PingAndWarmRequ
 	return &btpb.PingAndWarmResponse{}, nil
 }
 
+func (s *fakeService) getPingPeakInflight() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.pingPeakInflight
+}
+
 func (s *fakeService) getCallCount() int {
 	s.mu.Lock()
 	defer s.mu.Unlock()

bigtable/internal/transport/connpool_test.go

@@ -671,10 +671,9 @@ func TestNewBigtableChannelPool(t *testing.T) {
 
 	t.Run("DialFailure", func(t *testing.T) {
 		poolSize := 3
-		dialCount := 0
+		var dialCount atomic.Int64
 		dialFunc := func() (*BigtableConn, error) {
-			dialCount++
-			if dialCount > 1 {
+			if dialCount.Add(1) > 1 {
 				return nil, errors.New("simulated dial error")
 			}
 			fake := &fakeService{}
@@ -689,6 +688,59 @@ func TestNewBigtableChannelPool(t *testing.T) {
 	})
 }
 
+// TestNewBigtableChannelPoolParallelPriming verifies that initial pool priming
+// fans out across multiple workers, capped at maxPrimeWorkers. The fake server
+// holds each PingAndWarm open for primeDelay so we can read the peak concurrent
+// in-flight count; with sequential priming the peak would be 1.
+func TestNewBigtableChannelPoolParallelPriming(t *testing.T) {
+	ctx := context.Background()
+	primeDelay := 200 * time.Millisecond
+
+	cases := []struct {
+		name     string
+		poolSize int
+		wantPeak int
+	}{
+		{name: "BelowWorkerCap", poolSize: 4, wantPeak: 4},
+		{name: "AtWorkerCap", poolSize: maxPrimeWorkers, wantPeak: maxPrimeWorkers},
+		{name: "AboveWorkerCap", poolSize: maxPrimeWorkers * 2, wantPeak: maxPrimeWorkers},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			fake := &fakeService{}
+			fake.setDelay(primeDelay)
+			addr := setupTestServer(t, fake)
+			dialFunc := func() (*BigtableConn, error) { return dialBigtableserver(addr) }
+
+			start := time.Now()
+			pool, err := NewBigtableChannelPool(ctx, tc.poolSize, btopt.RoundRobin, dialFunc, time.Now(), poolOpts()...)
+			elapsed := time.Since(start)
+			if err != nil {
+				t.Fatalf("NewBigtableChannelPool failed: %v", err)
+			}
+			defer pool.Close()
+
+			if got := fake.getPingPeakInflight(); got != tc.wantPeak {
+				t.Errorf("peak in-flight PingAndWarm = %d, want %d", got, tc.wantPeak)
+			}
+
+			// Sanity bound on wall time: at most ceil(poolSize / workers) prime
+			// batches plus generous slack for dial/handshake overhead. A
+			// sequential implementation would take ~poolSize * primeDelay.
+			workers := tc.wantPeak
+			batches := (tc.poolSize + workers - 1) / workers
+			wantUnderSequential := time.Duration(tc.poolSize) * primeDelay
+			wantAtMost := time.Duration(batches)*primeDelay + 2*time.Second
+			if elapsed >= wantUnderSequential {
+				t.Errorf("elapsed %v >= sequential bound %v — priming did not parallelize", elapsed, wantUnderSequential)
+			}
+			if elapsed > wantAtMost {
+				t.Errorf("elapsed %v > expected upper bound %v", elapsed, wantAtMost)
+			}
+		})
+	}
+}
+
 func TestConnEntryCalculateConnLoadWithPenalty(t *testing.T) {
 	entry := &connEntry{}