fix(bigquery/storage/managedwriter): don't hold client mutex across GetWriteStream in resolvePool

[joy-twilio]

Fixes #20173.

Problem

Client.resolvePool held the client-wide c.mu across the GetWriteStream RPC used to discover a stream's location. Since that RPC touches no shared state (only the immutable rawClient), holding the lock across it serialized all concurrent stream opens on a client behind one ~one-RTT RPC.

For workloads that open streams frequently (e.g. a stream per write on the default stream), this made open latency grow linearly with concurrency — ~1s at 20 concurrent opens in a cross-region test, while a single uncontended open is ~115ms.

Fix

Move the GetWriteStream call outside the lock; take c.mu only for the pools map lookup/insert. The lookup and createPool remain atomic under the lock, so the create-once-per-location invariant is unchanged: racing opens for a new location still create exactly one pool (whichever caller takes the lock first creates it; the rest find it in the lookup).

Test

TestResolvePool_ConcurrentOpensAreNotSerialized fires N concurrent resolvePool calls against a fake BigQueryWrite server whose GetWriteStream blocks for a fixed delay, and asserts:

• the batch completes in ~1×delay (parallel) rather than ~N×delay (serialized), and
• exactly one pool is created for the location.

Verified the test fails on the pre-fix code (20 workers → ~2.0s for a 100ms delay) and passes after the fix (~0.1s). go test -race and the existing package unit tests pass.

[gemini-code-assist]

Code Review

This pull request optimizes the resolvePool method in the BigQuery managed writer client by moving the getWriteStream network RPC outside of the client-wide mutex lock, preventing concurrent stream opens from being serialized. A regression test is also introduced to verify parallel execution and pool caching. The reviewer identified a critical race condition where a concurrent c.Close() could set c.pools to nil while getWriteStream is running, leading to a panic when the lock is subsequently acquired. Adding a check to ensure c.pools is not nil after acquiring the lock is recommended to handle this gracefully.

[gemini-code-assist]

Since c.getWriteStream is now called outside of the client mutex c.mu, a concurrent call to c.Close() can run and complete while c.getWriteStream is in progress. When c.Close() completes, it sets c.pools to nil. If resolvePool then acquires c.mu and finds no cached pool, it will proceed to create a new pool and attempt to write to c.pools[loc], which will cause a panic because c.pools is nil.

To prevent this panic and handle the closed client state gracefully, we should check if c.pools is nil immediately after acquiring the lock.

Suggested change

	c.mu.Lock()
	defer c.mu.Unlock()
	if pool, ok := c.pools[loc]; ok {
	return pool, nil
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.pools == nil {
	return nil, fmt.Errorf("client is closed")
	}
	if pool, ok := c.pools[loc]; ok {
	return pool, nil
	}

[joy-twilio]

The comment states that Close() sets c.pools to nil, it doesn't. Close() iterates the map and closes each pool, but never reassigns or nils it (client.go:104-125). In fact c.pools is assigned exactly once, at construction (NewClient, client.go:99), and is never set to nil anywhere in the file. So after Close() the map is still a valid, non-nil map the write at c.pools[loc] = pool cannot panic (writing to a non-nil map is safe, even if its pool values are already closed)