feat(bigquery/storage/managedwriter): graceful connection drains (#11463)

shollyman · web-flow · commit b29912faab73 · 2025-01-17T10:37:06.000-08:00
* feat(bigquery/storage/managedwriter): graceful connection drains

This PR enables a more graceful shutdown when we're doing client
initiated reconnects, typically in response to things like schema
changes.

With this PR, rather than immediately cancelling the previous connection
context, we instead use a goroutine to defer this by a fixed interval,
which gives our recv goroutine more time to process any incoming
notifications from the service.

This can yield a net reduction in duplicate rows for some scenarios,
particularly when EnableWriteRetries is enabled and we see sequential
disconnects (e.g. a multiplex connection observing a number of schema
updates).

It does, however, yield possible short term increases in resource usage,
as we're now potentially retaining goroutines/channels/connections for a
longer interval during these reconnect events.  This PR updates testing
that looks for leaking resources with this new understanding
accordingly.
diff --git a/bigquery/storage/managedwriter/connection.go b/bigquery/storage/managedwriter/connection.go
@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"io"
 	"sync"
+	"time"
 
 	"cloud.google.com/go/bigquery/storage/apiv1/storagepb"
 	"github.com/googleapis/gax-go/v2"
@@ -36,6 +37,8 @@ const (
 
 var (
 	errNoRouterForPool = errors.New("no router for connection pool")
+	// TODO(https://github.com/googleapis/google-cloud-go/issues/11460): revisit if this should be user configurable
+	gracefulReconnectDuration = 20 * time.Second
 )
 
 // connectionPool represents a pooled set of connections.
@@ -489,7 +492,12 @@ func (co *connection) getStream(arc *storagepb.BigQueryWrite_AppendRowsClient, f
 		close(co.pending)
 	}
 	if co.cancel != nil {
-		co.cancel()
+		// Delay cancellation to give queued writes a chance to drain normally.
+		oldCancel := co.cancel
+		go func() {
+			time.Sleep(gracefulReconnectDuration)
+			oldCancel()
+		}()
 		co.ctx, co.cancel = context.WithCancel(co.pool.ctx)
 	}
 
diff --git a/bigquery/storage/managedwriter/managed_stream_test.go b/bigquery/storage/managedwriter/managed_stream_test.go
@@ -519,11 +519,13 @@ func TestManagedStream_LeakingGoroutinesReconnect(t *testing.T) {
 		[]byte("foo"),
 	}
 
+	appendCount := 15
+	// Give a small budget for additional goroutines to account for jitter.
 	threshold := runtime.NumGoroutine() + 5
 
 	// Send a bunch of appends that will trigger reconnects and monitor that
 	// goroutine growth stays within bounded threshold.
-	for i := 0; i < 30; i++ {
+	for i := 0; i < appendCount; i++ {
 		writeCtx := context.Background()
 		r, err := ms.AppendRows(writeCtx, fakeData)
 		if err != nil {
@@ -539,12 +541,24 @@ func TestManagedStream_LeakingGoroutinesReconnect(t *testing.T) {
 		if testArc.openCount != i+2 {
 			t.Errorf("should trigger a reconnect, but found openCount %d", testArc.openCount)
 		}
-		if i%10 == 0 {
+		// Bump our threshold each append.  We add 2 goroutines each reconnect due to a new processor
+		// and a graceful cancellation, but this should recover over time.
+		threshold = threshold + 2
+
+		if i%5 == 0 {
 			if current := runtime.NumGoroutine(); current > threshold {
 				t.Errorf("potential goroutine leak, append %d: current %d, threshold %d", i, current, threshold)
 			}
 		}
 	}
+	// Verify goroutine count drops after graceful shutdowns should have concluded.
+	time.Sleep(gracefulReconnectDuration)
+	threshold = threshold - (2 * appendCount)
+
+	if current := runtime.NumGoroutine(); current > threshold {
+		t.Errorf("potential goroutine leak after recovery: current %d, threshold %d", current, threshold)
+	}
+
 }
 
 func TestManagedWriter_CancellationDuringRetry(t *testing.T) {
