fix(bigquery/storage/managedwriter): address possible deadlocks (#8507)

The context refactoring in #8275 introduced two possible sources of deadlocks in the ManagedStream AppendRows call path.  This PR addresses those, and augments deadlock testing to cover this case.

Fixes: https://github.com/googleapis/google-cloud-go/issues/8505

bigquery/storage/managedwriter/managed_stream.go

@@ -197,10 +197,11 @@ func (ms *ManagedStream) Finalize(ctx context.Context, opts ...gax.CallOption) (
 func (ms *ManagedStream) appendWithRetry(pw *pendingWrite, opts ...gax.CallOption) error {
 	for {
 		ms.mu.Lock()
-		if ms.err != nil {
-			return ms.err
-		}
+		err := ms.err
 		ms.mu.Unlock()
+		if err != nil {
+			return err
+		}
 		conn, err := ms.pool.selectConn(pw)
 		if err != nil {
 			pw.markDone(nil, err)
@@ -291,10 +292,11 @@ func (ms *ManagedStream) buildRequest(data [][]byte) *storagepb.AppendRowsReques
 func (ms *ManagedStream) AppendRows(ctx context.Context, data [][]byte, opts ...AppendOption) (*AppendResult, error) {
 	// before we do anything, ensure the writer isn't closed.
 	ms.mu.Lock()
-	if ms.err != nil {
-		return nil, ms.err
-	}
+	err := ms.err
 	ms.mu.Unlock()
+	if err != nil {
+		return nil, err
+	}
 	// Ensure we build the request and pending write with a consistent schema version.
 	curSchemaVersion := ms.curDescVersion
 	req := ms.buildRequest(data)

bigquery/storage/managedwriter/managed_stream_test.go

@@ -421,6 +421,17 @@ func TestManagedStream_AppendDeadlocks(t *testing.T) {
 		// Issue two closes, to ensure we're not deadlocking there either.
 		ms.Close()
 		ms.Close()
+
+		// Issue two more appends, ensure we're not deadlocked as the writer is closed.
+		gotErr = ms.appendWithRetry(pw)
+		if !errors.Is(gotErr, io.EOF) {
+			t.Errorf("expected io.EOF, got %v", gotErr)
+		}
+		gotErr = ms.appendWithRetry(pw)
+		if !errors.Is(gotErr, io.EOF) {
+			t.Errorf("expected io.EOF, got %v", gotErr)
+		}
+
 	}
 
 }