fix(storage): decouple Recv loop in gRPCOneshotBidiWriteBufferSender … (#14166)

…to prevent deadlocks

Noticed a deadlock when OneShot sender was bombarded with a lot of data
during a PCU benchmarking run

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: cpriti-os <202586561+cpriti-os@users.noreply.github.com>

Author: cpriti-os

storage/grpc_writer.go

@@ -240,7 +240,7 @@ func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 	return w, nil
 }
 
-// gRPCWriter is a wrapper around the the gRPC client-stream API that manages
+// gRPCWriter is a wrapper around the gRPC client-stream API that manages
 // sending chunks of data provided by the user over the stream.
 type gRPCWriter struct {
 	preRunCtx context.Context
@@ -992,22 +992,6 @@ func (w *gRPCWriter) newGRPCOneshotBidiWriteBufferSender() *gRPCOneshotBidiWrite
 
 func (s *gRPCOneshotBidiWriteBufferSender) err() error { return s.streamErr }
 
-// drainInboundStream calls stream.Recv() repeatedly until an error is returned.
-// It returns the last Resource received on the stream, or nil if no Resource
-// was returned. drainInboundStream always returns a non-nil error. io.EOF
-// indicates all messages were successfully read.
-func drainInboundStream(stream storagepb.Storage_BidiWriteObjectClient) (object *storagepb.Object, err error) {
-	for err == nil {
-		var resp *storagepb.BidiWriteObjectResponse
-		resp, err = stream.Recv()
-		// GetResource() returns nil on a nil response
-		if resp.GetResource() != nil {
-			object = resp.GetResource()
-		}
-	}
-	return object, err
-}
-
 func (s *gRPCOneshotBidiWriteBufferSender) connect(ctx context.Context, cs gRPCBufSenderChans, opts ...gax.CallOption) {
 	s.streamErr = nil
 	ctx = gRPCWriteRequestParams{bucket: s.bucket}.apply(ctx)
@@ -1019,59 +1003,93 @@ func (s *gRPCOneshotBidiWriteBufferSender) connect(ctx context.Context, cs gRPCB
 	}
 
 	go func() {
-		firstSend := true
-		for r := range cs.requests {
-			if r.requestAck {
-				cs.requestAcks <- struct{}{}
-				continue
-			}
+		var sendErr, recvErr error
+		sendDone := make(chan struct{})
+		recvDone := make(chan struct{})
 
-			var bufChecksum *uint32
-			if !s.disableAutoChecksum {
-				bufChecksum = proto.Uint32(crc32.Checksum(r.buf, crc32cTable))
-			}
-			objectChecksums := getObjectChecksums(&getObjectChecksumsParams{
-				sendCRC32C:          s.sendCRC32C,
-				objectAttrs:         s.objectAttrs,
-				fullObjectChecksum:  s.fullObjectChecksum,
-				disableAutoChecksum: s.disableAutoChecksum,
-				finishWrite:         r.finishWrite,
-			})
-			req := bidiWriteObjectRequest(r, bufChecksum, objectChecksums)
-
-			if firstSend {
-				proto.Merge(req, s.firstMessage)
-				firstSend = false
-			}
+		go func() {
+			sendErr = func() error {
+				firstSend := true
+				for {
+					select {
+					case <-recvDone:
+						// Because `requests` is not connected to the gRPC machinery, we
+						// have to check for asynchronous termination on the receive side.
+						return nil
+					case r, ok := <-cs.requests:
+						if !ok {
+							stream.CloseSend()
+							return nil
+						}
+						if r.requestAck {
+							cs.requestAcks <- struct{}{}
+							continue
+						}
+
+						var bufChecksum *uint32
+						if !s.disableAutoChecksum {
+							bufChecksum = proto.Uint32(crc32.Checksum(r.buf, crc32cTable))
+						}
+						objectChecksums := getObjectChecksums(&getObjectChecksumsParams{
+							sendCRC32C:          s.sendCRC32C,
+							objectAttrs:         s.objectAttrs,
+							fullObjectChecksum:  s.fullObjectChecksum,
+							disableAutoChecksum: s.disableAutoChecksum,
+							finishWrite:         r.finishWrite,
+						})
+						req := bidiWriteObjectRequest(r, bufChecksum, objectChecksums)
+
+						if firstSend {
+							proto.Merge(req, s.firstMessage)
+							firstSend = false
+						}
 
-			if err := stream.Send(req); err != nil {
-				_, s.streamErr = drainInboundStream(stream)
-				if err != io.EOF {
-					s.streamErr = err
+						if err := stream.Send(req); err != nil {
+							return err
+						}
+
+						if r.finishWrite {
+							stream.CloseSend()
+							return nil
+						}
+
+						// Oneshot uploads assume all flushes succeed.
+						if r.flush {
+							select {
+							case cs.completions <- gRPCBidiWriteCompletion{flushOffset: r.offset + int64(len(r.buf))}:
+							case <-stream.Context().Done():
+								return stream.Context().Err()
+							}
+						}
+					}
 				}
-				close(cs.completions)
-				return
-			}
+			}()
+			close(sendDone)
+		}()
 
-			if r.finishWrite {
-				stream.CloseSend()
-				// Oneshot uploads only read from the response stream on completion or
-				// failure
-				obj, err := drainInboundStream(stream)
-				if obj == nil || err != io.EOF {
-					s.streamErr = err
-				} else {
-					cs.completions <- gRPCBidiWriteCompletion{flushOffset: obj.GetSize(), resource: obj}
+		go func() {
+			recvErr = func() error {
+				for {
+					resp, err := stream.Recv()
+					if err != nil {
+						return err
+					}
+					if c := completion(resp); c != nil {
+						select {
+						case cs.completions <- *c:
+						case <-stream.Context().Done():
+							return stream.Context().Err()
+						}
+					}
 				}
-				close(cs.completions)
-				return
-			}
+			}()
+			close(recvDone)
+		}()
 
-			// Oneshot uploads assume all flushes succeed
-			if r.flush {
-				cs.completions <- gRPCBidiWriteCompletion{flushOffset: r.offset + int64(len(r.buf))}
-			}
-		}
+		<-sendDone
+		<-recvDone
+		s.streamErr = pickStreamError(recvErr, sendErr)
+		close(cs.completions)
 	}()
 }
 
@@ -1203,7 +1221,11 @@ func (s *gRPCResumableBidiWriteBufferSender) connect(ctx context.Context, cs gRP
 						return err
 					}
 					if c := completion(resp); c != nil {
-						cs.completions <- *c
+						select {
+						case cs.completions <- *c:
+						case <-stream.Context().Done():
+							return stream.Context().Err()
+						}
 					}
 				}
 			}()
@@ -1212,15 +1234,7 @@ func (s *gRPCResumableBidiWriteBufferSender) connect(ctx context.Context, cs gRP
 
 		<-sendDone
 		<-recvDone
-		// Prefer recvErr since that's where RPC errors are delivered
-		if recvErr != nil {
-			s.streamErr = recvErr
-		} else if sendErr != nil {
-			s.streamErr = sendErr
-		}
-		if s.streamErr == io.EOF {
-			s.streamErr = nil
-		}
+		s.streamErr = pickStreamError(recvErr, sendErr)
 		close(cs.completions)
 	}()
 }
@@ -1329,7 +1343,11 @@ func (s *gRPCAppendBidiWriteBufferSender) handleStream(stream storagepb.Storage_
 				s.maybeUpdateFirstMessage(resp)
 
 				if c := completion(resp); c != nil {
-					cs.completions <- *c
+					select {
+					case cs.completions <- *c:
+					case <-stream.Context().Done():
+						return stream.Context().Err()
+					}
 				}
 			}
 		}()
@@ -1338,15 +1356,7 @@ func (s *gRPCAppendBidiWriteBufferSender) handleStream(stream storagepb.Storage_
 
 	<-sendDone
 	<-recvDone
-	// Prefer recvErr since that's where RPC errors are delivered
-	if recvErr != nil {
-		s.streamErr = recvErr
-	} else if sendErr != nil {
-		s.streamErr = sendErr
-	}
-	if s.streamErr == io.EOF {
-		s.streamErr = nil
-	}
+	s.streamErr = pickStreamError(recvErr, sendErr)
 	close(cs.completions)
 }
 
@@ -1613,3 +1623,12 @@ func withBidiWriteObjectRedirectionErrorRetries(s *settings) (newr *retryConfig)
 	}
 	return newr
 }
+
+// pickStreamError determines the final error to be reported by prioritizing recvErr.
+// An io.EOF from a receiver is not considered an error.
+func pickStreamError(recvErr, sendErr error) error {
+	if recvErr != nil && recvErr != io.EOF {
+		return recvErr
+	}
+	return sendErr
+}

storage/grpc_writer_test.go

@@ -16,8 +16,11 @@ package storage
 
 import (
 	"context"
+	"errors"
+	"io"
 	"sync"
 	"testing"
+	"time"
 
 	"cloud.google.com/go/storage/internal/apiv2/storagepb"
 	gax "github.com/googleapis/gax-go/v2"
@@ -294,3 +297,118 @@ func filterDataRequests(reqs []gRPCBidiWriteRequest) []gRPCBidiWriteRequest {
 	}
 	return dataReqs
 }
+
+// Test the logic correctly handles the combination of io.EOF
+// from Recv (recvErr) and a generic error from Send (sendErr).
+func TestGRPCWriterErrorHandling(t *testing.T) {
+	// As this is deeply embedded in the unexported types, we verify the logic
+	// by simulating the exact error assignment sequence.
+	tests := []struct {
+		name      string
+		recvErr   error
+		sendErr   error
+		wantError error
+	}{
+		{
+			name:      "recvErr is io.EOF, sendErr is nil",
+			recvErr:   io.EOF,
+			sendErr:   nil,
+			wantError: nil,
+		},
+		{
+			name:      "recvErr is io.EOF, sendErr is an error",
+			recvErr:   io.EOF,
+			sendErr:   errors.New("send error"),
+			wantError: errors.New("send error"), // Send error takes precedence.
+		},
+		{
+			name:      "recvErr is an error, sendErr is nil",
+			recvErr:   errors.New("recv error"),
+			sendErr:   nil,
+			wantError: errors.New("recv error"), // Recv error takes precedence.
+		},
+		{
+			name:      "recvErr is an error, sendErr is an error",
+			recvErr:   errors.New("recv error"),
+			sendErr:   errors.New("send error"),
+			wantError: errors.New("recv error"), // Recv error takes precedence.
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var streamErr error
+
+			streamErr = pickStreamError(tt.recvErr, tt.sendErr)
+
+			if tt.wantError == nil {
+				if streamErr != nil {
+					t.Errorf("got error %v, want nil", streamErr)
+				}
+			} else {
+				if streamErr == nil || streamErr.Error() != tt.wantError.Error() {
+					t.Errorf("got error %v, want %v", streamErr, tt.wantError)
+				}
+			}
+		})
+	}
+}
+
+// TestGRPCWriter_Deadlock simulates a deadlock scenario if Recv and Send channels
+// were not isolated in gRPCOneshotBidiWriteBufferSender.
+func TestGRPCWriter_Deadlock(t *testing.T) {
+	// A timeout means a deadlock likely occurred.
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+
+	sendDone := make(chan struct{})
+	recvDone := make(chan struct{})
+
+	requests := make(chan gRPCBidiWriteRequest)
+	completions := make(chan gRPCBidiWriteCompletion)
+
+	var sendErr error
+
+	go func() {
+		sendErr = func() error {
+			for {
+				select {
+				case <-recvDone:
+					return nil
+				case r, ok := <-requests:
+					if !ok {
+						return nil
+					}
+					if r.requestAck {
+						continue
+					}
+					// mimic send logic
+					if r.finishWrite {
+						return nil
+					}
+				}
+			}
+		}()
+		close(sendDone)
+	}()
+
+	go func() {
+		// Mimic recv loop that immediately exits.
+		// If recvDone isn't checked by the sender loop, sending
+		// requests could block forever if the consumer closes early.
+		close(recvDone)
+	}()
+
+	// sendDone should be closed immediately.
+	select {
+	case <-sendDone:
+		// Success, no deadlock.
+	case <-ctx.Done():
+		t.Fatal("deadlock detected: send loop did not exit after recvDone was closed")
+	}
+
+	if sendErr != nil {
+		t.Errorf("expected no error, got %v", sendErr)
+	}
+	close(completions)
+}