fix(bigtable): guard NewStream OnFinish against grpc-go double-fire (#20295)

## Summary

grpc-go v1.82.x can invoke a registered `grpc.OnFinish` callback more
than once on some stream-creation failure paths (closed `ClientConn` is
the observed one: the retry unwinder fires it once from inside
`withRetry`, then `newClientStream`'s deferred stream-teardown fires it
again). Compounding that, the old `BigtableChannelPool.NewStream` also
manually decremented `streamingLoad` on the returned-error path assuming
OnFinish had NOT fired. Combined: `streamingLoad` went 2–3 decrements
deep per failed NewStream. Debug pages surfaced this as **\"Streaming in
flight: -9\"** on an idle classic pool.

## Fix

Trust `OnFinish` as the single source of truth for load accounting +
error attribution. Gate the callback body with
`atomic.Bool.CompareAndSwap(false, true)` so it runs exactly once
regardless of how many times grpc-go fires it. Belt-and-suspenders:
after `entry.conn.NewStream` returns an error, call the same accounting
closure directly — the CAS makes it a no-op if grpc-go already fired,
and it covers the zero-fire path if grpc-go ever returns an error
without arming OnFinish.

- Live path: one atomic Load per NewStream on the happy path; cache-map
lookup only on the recovery branch.
- Non-goal: no changes to per-attempt tracer / metric emission logic.

## Test plan

- [x] New regression test
\`TestPoolNewStream/ImmediateFailureLoadStaysNonNegative\` — closes the
underlying \`ClientConn\`, calls \`NewStream\` 5×, asserts
\`streamingLoad >= 0\` after each and \`== 0\` at end. Fails on pre-fix
code with \`streamingLoad\` going to \`-5\`.
- [x] New coverage test
\`TestPoolNewStream/OnFinishFiresAtLeastOncePerFailedNewStream\` — wraps
a user OnFinish counter around each NewStream on a closed conn, asserts
fires ∈ {1, 2, 3} — pins the \"at least once\" invariant the whole fix
leans on, and cries out early if grpc-go ever regresses to zero-fire.
- [x] Existing \`TestPoolNewStream\` subtests unchanged and passing
(\`Strategy_least_in_flight\`, \`Strategy_round_robin\`,
\`Strategy_power_of_two_least_in_flight\`, \`EmptyPoolNewStream\`,
\`NewStreamServerError\`).
- [x] \`go test ./internal/transport/ -race -count=1 -short\` clean.

Author: sushanb

bigtable/internal/transport/connpool.go

@@ -769,11 +769,26 @@ func (p *BigtableChannelPool) getBigtableConn() *BigtableConn {
 }
 
 // NewStream selects a connection by the configured load-balancing strategy
-// and opens a stream on it. grpc.OnFinish fires exactly once for any stream
-// that was successfully created (normal completion, context cancellation,
-// transport teardown), so it is the single source of truth for both load
-// accounting and per-stream error attribution — no need to wrap the
-// returned ClientStream.
+// and opens a stream on it. Load accounting and per-stream error attribution
+// happen in a single `finish` closure that runs AT MOST ONCE per NewStream
+// call — the CAS in the closure body enforces that regardless of how many
+// times grpc-go invokes OnFinish (which has been observed to be 0, 1, or
+// more depending on the stream-creation path):
+//
+//   - Nominal path: grpc-go fires OnFinish once at stream teardown; the
+//     CAS wins on the first fire, all bookkeeping runs.
+//   - Double-fire path: grpc-go re-fires OnFinish on some stream-creation
+//     failures (retry unwinder + deferred stream-teardown both call it);
+//     the CAS wins once, subsequent fires are no-ops. Without the guard,
+//     streamingLoad would drift negative — the debug pages first surfaced
+//     this as "Streaming in flight: -9" on an idle classic pool.
+//   - Zero-fire path: if `entry.conn.NewStream` returns an error before
+//     grpc-go arms the OnFinish trigger, the deferred `finish(err)` at
+//     the return site picks it up. Belt-and-suspenders — no current
+//     grpc-go version does this on any path bigtable exercises, but the
+//     accounting invariant is now load-bearing and OnFinish is a
+//     best-effort callback, so we don't want a future interceptor
+//     addition to silently leak +1.
 func (p *BigtableChannelPool) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
 	entry, err := p.selectFunc()
 	if err != nil {
@@ -782,26 +797,36 @@ func (p *BigtableChannelPool) NewStream(ctx context.Context, desc *grpc.StreamDe
 
 	entry.streamingLoad.Add(1)
 
-	onFinish := grpc.OnFinish(func(err error) {
+	// onFinishFired is the exactly-once gate for `finish`. Named for the
+	// callback state ("did the accounting closure run?") — not the stream
+	// state — because both the OnFinish callback and the error-path
+	// fallback below race to be the first-and-only caller. Reads of this
+	// field are only ever "did we already fire?".
+	var onFinishFired atomic.Bool
+	finish := func(err error) {
+		if !onFinishFired.CompareAndSwap(false, true) {
+			return
+		}
 		if err != nil {
 			entry.errorCount.Add(1)
 			entry.applyErrorPenalty(err)
 		}
 		entry.streamingLoad.Add(-1)
-	})
+	}
 	// Prepend onto a fresh slice so we never write into spare capacity of
 	// the caller's opts (which would race with concurrent NewStream calls
 	// that share the same backing array).
-	opts = append([]grpc.CallOption{onFinish}, opts...)
+	opts = append([]grpc.CallOption{grpc.OnFinish(finish)}, opts...)
 
 	stream, err := entry.conn.NewStream(ctx, desc, method, opts...)
 	if err != nil {
-		entry.errorCount.Add(1)
-		entry.applyErrorPenalty(err)
-		entry.streamingLoad.Add(-1) // Decrement immediately on creation failure
+		// Zero-fire fallback: if grpc-go never armed OnFinish before
+		// returning this error, `finish` wins the CAS here. If it DID
+		// fire OnFinish (once or twice), the CAS is already lost and
+		// this call is a no-op — safe to unconditionally invoke.
+		finish(err)
 		return nil, err
 	}
-
 	return stream, nil
 }
 

bigtable/internal/transport/connpool_test.go

@@ -631,6 +631,41 @@ func TestPoolNewStream(t *testing.T) {
 			t.Errorf("Load is %d, want 0 after stream error", pool.getConns()[0].streamingLoad.Load())
 		}
 	})
+
+	// Regression: when entry.conn.NewStream itself returns an error (e.g.
+	// the ClientConn is closed), grpc-go's deferred endOfClientStream can
+	// fire every registered OnFinish. Load accounting must not double-
+	// decrement on that path or streamingLoad goes negative and stays
+	// negative.
+	t.Run("ImmediateFailureLoadStaysNonNegative", func(t *testing.T) {
+		poolSize := 1
+		fake := &fakeService{}
+		addr := setupTestServer(t, fake)
+		dialFunc := func() (*BigtableConn, error) { return dialBigtableserver(addr) }
+		pool, err := NewBigtableChannelPool(ctx, poolSize, btopt.RoundRobin, dialFunc, time.Now(), poolOpts()...)
+		if err != nil {
+			t.Fatalf("Failed to create pool: %v", err)
+		}
+		defer pool.Close()
+
+		entry := pool.getConns()[0]
+		// Close the underlying ClientConn so any subsequent NewStream fails
+		// immediately with codes.Canceled / codes.Unavailable.
+		entry.conn.Close()
+
+		for i := 0; i < 5; i++ {
+			_, err := pool.NewStream(ctx, &grpc.StreamDesc{StreamName: "StreamingCall"}, "/grpc.testing.BenchmarkService/StreamingCall")
+			if err == nil {
+				t.Fatalf("attempt %d: NewStream unexpectedly succeeded on a closed conn", i)
+			}
+			if got := entry.streamingLoad.Load(); got < 0 {
+				t.Fatalf("attempt %d: streamingLoad went negative (%d) — accounting fired more than once", i, got)
+			}
+		}
+		if got := entry.streamingLoad.Load(); got != 0 {
+			t.Errorf("streamingLoad after 5 failed NewStreams = %d, want 0", got)
+		}
+	})
 }
 
 func TestNewBigtableChannelPool(t *testing.T) {