Revert "fix(storage): fix stream termination in MRD. (googleapis#11432)" (googleapis#11778)

This reverts commit 3d4e62f.

We still need to do some debugging on an integration test failure for this.

Fixes googleapis#11769

Author: tritone

storage/grpc_client.go

@@ -1185,8 +1185,6 @@ func (c *grpcStorageClient) NewMultiRangeDownloader(ctx context.Context, params
 		done:             false,
 		activeTask:       0,
 		streamRecreation: false,
-		endReceiver:      false,
-		endSender:        false,
 	}
 
 	// streamManager goroutine runs in background where we send message to gcs and process response.
@@ -1197,21 +1195,18 @@ func (c *grpcStorageClient) NewMultiRangeDownloader(ctx context.Context, params
 			case <-rr.ctx.Done():
 				rr.mu.Lock()
 				rr.done = true
-				rr.endSender = true
-				if rr.stream != nil {
-					rr.stream.CloseSend()
-				}
 				rr.mu.Unlock()
 				return
 			case <-rr.managerRetry:
-				// We are not closing stream here as it is already closed and we are retring it.
 				return
 			case <-rr.closeManager:
 				rr.mu.Lock()
-				if rr.stream != nil {
-					rr.stream.CloseSend()
+				if len(rr.mp) != 0 {
+					for key := range rr.mp {
+						rr.mp[key].callback(rr.mp[key].offset, rr.mp[key].totalBytesWritten, fmt.Errorf("stream closed early"))
+						delete(rr.mp, key)
+					}
 				}
-				rr.endSender = true
 				rr.activeTask = 0
 				rr.mu.Unlock()
 				return
@@ -1260,29 +1255,11 @@ func (c *grpcStorageClient) NewMultiRangeDownloader(ctx context.Context, params
 		for {
 			select {
 			case <-rr.ctx.Done():
-				rr.mu.Lock()
-				rr.endReceiver = true
 				rr.done = true
-				if len(rr.mp) != 0 {
-					drainInboundReadStream(rr.stream)
-				}
-				for key := range rr.mp {
-					rr.mp[key].callback(rr.mp[key].offset, rr.mp[key].totalBytesWritten, rr.ctx.Err())
-					delete(rr.mp, key)
-				}
-				rr.activeTask = 0
-				rr.mu.Unlock()
 				return
 			case <-rr.receiverRetry:
-				// We are not draining from stream here as it is already closed and we are retring it.
 				return
 			case <-rr.closeReceiver:
-				rr.mu.Lock()
-				if len(rr.mp) != 0 {
-					drainInboundReadStream(rr.stream)
-				}
-				rr.endReceiver = true
-				rr.mu.Unlock()
 				return
 			default:
 				// This function reads the data sent for a particular range request and has a callback
@@ -1292,9 +1269,7 @@ func (c *grpcStorageClient) NewMultiRangeDownloader(ctx context.Context, params
 					rr.readHandle = resp.GetReadHandle().GetHandle()
 				}
 				if err == io.EOF {
-					rr.mu.Lock()
-					rr.endReceiver = true
-					rr.mu.Unlock()
+					err = nil
 				}
 				if err != nil {
 					// cancel stream and reopen the stream again.
@@ -1366,8 +1341,6 @@ func (c *grpcStorageClient) NewMultiRangeDownloader(ctx context.Context, params
 		err = rr.retryStream(err)
 		if err != nil {
 			rr.mu.Lock()
-			rr.endReceiver = true
-			rr.endSender = true
 			for key := range rr.mp {
 				rr.mp[key].callback(rr.mp[key].offset, rr.mp[key].totalBytesWritten, err)
 				delete(rr.mp, key)
@@ -1377,10 +1350,6 @@ func (c *grpcStorageClient) NewMultiRangeDownloader(ctx context.Context, params
 			rr.mu.Unlock()
 			rr.close()
 		} else {
-			rr.mu.Lock()
-			rr.endReceiver = false
-			rr.endSender = false
-			rr.mu.Unlock()
 			// If stream recreation happened successfully lets again start
 			// both the goroutine making the whole flow asynchronous again.
 			if thread == "receiver" {
@@ -1514,39 +1483,18 @@ func (mr *gRPCBidiReader) wait() {
 
 // Close will notify stream manager goroutine that the reader has been closed, if it's still running.
 func (mr *gRPCBidiReader) close() error {
-	mr.closeManager <- true
-	mr.closeReceiver <- true
-	mr.mu.Lock()
-	for key := range mr.mp {
-		mr.mp[key].callback(mr.mp[key].offset, mr.mp[key].totalBytesWritten, fmt.Errorf("stream closed early"))
-		delete(mr.mp, key)
+	if mr.cancel != nil {
+		mr.cancel()
 	}
+	mr.mu.Lock()
 	mr.done = true
 	mr.activeTask = 0
 	mr.mu.Unlock()
-	mr.mu.Lock()
-	tryClosing := !(mr.endReceiver && mr.endSender)
-	mr.mu.Unlock()
-
-	for tryClosing {
-		mr.mu.Lock()
-		tryClosing = !(mr.endReceiver && mr.endSender)
-		mr.mu.Unlock()
-	}
-	defer mr.cancel()
+	mr.closeReceiver <- true
+	mr.closeManager <- true
 	return nil
 }
 
-// drainInboundReadStream calls stream.Recv() repeatedly until an error is returned.
-// drainInboundReadStream always returns a non-nil error. io.EOF indicates all
-// messages were successfully read.
-func drainInboundReadStream(stream storagepb.Storage_BidiReadObjectClient) (err error) {
-	for err == nil {
-		_, err = stream.Recv()
-	}
-	return err
-}
-
 func (mrr *gRPCBidiReader) getHandle() []byte {
 	return mrr.readHandle
 }
@@ -1977,8 +1925,6 @@ type gRPCBidiReader struct {
 	objectSize       int64               // always use the mutex when accessing this variable
 	retrier          func(error, string)
 	streamRecreation bool // This helps us identify if stream recreation is in progress or not. If stream recreation gets called from two goroutine then this will stop second one.
-	endReceiver      bool
-	endSender        bool
 }
 
 // gRPCReader is used by storage.Reader if the experimental option WithGRPCBidiReads is passed.
@@ -2707,11 +2653,11 @@ func bucketContext(ctx context.Context, bucket string) context.Context {
 	return gax.InsertMetadataIntoOutgoingContext(ctx, hds...)
 }
 
-// drainInboundWriteStream calls stream.Recv() repeatedly until an error is returned.
+// 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. drainInboundWriteStream always returns a non-nil error. io.EOF
+// was returned. drainInboundStream always returns a non-nil error. io.EOF
 // indicates all messages were successfully read.
-func drainInboundWriteStream(stream storagepb.Storage_BidiWriteObjectClient) (object *storagepb.Object, err error) {
+func drainInboundStream(stream storagepb.Storage_BidiWriteObjectClient) (object *storagepb.Object, err error) {
 	for err == nil {
 		var resp *storagepb.BidiWriteObjectResponse
 		resp, err = stream.Recv()
@@ -2791,7 +2737,7 @@ func (s *gRPCOneshotBidiWriteBufferSender) sendBuffer(ctx context.Context, buf [
 
 	sendErr := s.stream.Send(req)
 	if sendErr != nil {
-		obj, err = drainInboundWriteStream(s.stream)
+		obj, err = drainInboundStream(s.stream)
 		s.stream = nil
 		if sendErr != io.EOF {
 			err = sendErr
@@ -2804,7 +2750,7 @@ func (s *gRPCOneshotBidiWriteBufferSender) sendBuffer(ctx context.Context, buf [
 		s.stream.CloseSend()
 		// Oneshot uploads only read from the response stream on completion or
 		// failure
-		obj, err = drainInboundWriteStream(s.stream)
+		obj, err = drainInboundStream(s.stream)
 		s.stream = nil
 		if err == io.EOF {
 			err = nil
@@ -2916,7 +2862,7 @@ func (s *gRPCResumableBidiWriteBufferSender) sendBuffer(ctx context.Context, buf
 
 	sendErr := s.stream.Send(req)
 	if sendErr != nil {
-		obj, err = drainInboundWriteStream(s.stream)
+		obj, err = drainInboundStream(s.stream)
 		s.stream = nil
 		if err == io.EOF {
 			// This is unexpected - we got an error on Send(), but not on Recv().
@@ -2928,7 +2874,7 @@ func (s *gRPCResumableBidiWriteBufferSender) sendBuffer(ctx context.Context, buf
 
 	if finishWrite {
 		s.stream.CloseSend()
-		obj, err = drainInboundWriteStream(s.stream)
+		obj, err = drainInboundStream(s.stream)
 		s.stream = nil
 		if err == io.EOF {
 			err = nil
@@ -3025,5 +2971,6 @@ func checkCanceled(err error) error {
 	if status.Code(err) == codes.Canceled {
 		return context.Canceled
 	}
+
 	return err
 }

storage/integration_test.go

@@ -597,123 +597,6 @@ func TestIntegration_MRDWithNonRetriableError(t *testing.T) {
 	})
 }
 
-// Test that context cancellation correctly stops a multi range download before completion.
-func TestIntegration_MultiRangeDownloaderContextCancel(t *testing.T) {
-	multiTransportTest(skipHTTP("gRPC implementation specific test"), t, func(t *testing.T, ctx context.Context, bucket, _ string, client *Client) {
-		ctx, close := context.WithDeadline(ctx, time.Now().Add(time.Second*30))
-		defer close()
-		content := make([]byte, 5<<20)
-		rand.New(rand.NewSource(0)).Read(content)
-		objName := "mrdnonretry"
-		// Upload test data.
-		obj := client.Bucket(bucket).Object(objName)
-		if err := writeObject(ctx, obj, "text/plain", content); err != nil {
-			t.Fatal(err)
-		}
-		defer func() {
-			if err := obj.Delete(ctx); err != nil {
-				log.Printf("failed to delete test object: %v", err)
-			}
-		}()
-		// Create a multi-range-reader and then cancel the context before completing the reads.
-		readerCtx, cancel := context.WithCancel(ctx)
-		reader, err := obj.NewMultiRangeDownloader(readerCtx)
-		if err != nil {
-			t.Fatalf("NewMultiRangeDownloader: %v", err)
-		}
-		res := make([]multiRangeDownloaderOutput, 3)
-		callback := func(x, y int64, err error) {
-			res[0].offset = x
-			res[0].limit = y
-			res[0].err = err
-		}
-		callback1 := func(x, y int64, err error) {
-			res[1].offset = x
-			res[1].limit = y
-			res[1].err = err
-		}
-		callback2 := func(x, y int64, err error) {
-			res[2].offset = x
-			res[2].limit = y
-			res[2].err = err
-		}
-		// Add one range on the reader, and then cancel the context.
-		reader.Add(&res[0].buf, 0, int64(len(content)), callback)
-		// As context is cancelled remaining ranges would result in context cancelled error or stream is closed errors.
-		cancel()
-		reader.Add(&res[1].buf, -10, 0, callback1)
-		reader.Add(&res[2].buf, 0, 10, callback2)
-		reader.Wait()
-		// we can get stream is closed, can't add range error in case process is over before we add the range.
-		expErr := fmt.Errorf("stream is closed, can't add range")
-		for i, k := range res {
-			// if we get nil error for any callback other than first, that should be an error.
-			if i == 0 && k.err == nil && !bytes.Equal(content, k.buf.Bytes()) {
-				t.Errorf("Error in read range offset %v, limit %v, got: %v; want: %v",
-					k.offset, k.limit, len(k.buf.Bytes()), len(content))
-			}
-			if k.err == nil && k.err.Error() != expErr.Error() && !errors.Is(err, context.Canceled) && !(status.Code(err) == codes.Canceled) {
-				t.Fatalf("read range %v to %v: got error %v, want nil, context.Canceled or stream is closed error", k.offset, k.limit, k.err)
-			}
-		}
-		if err = reader.Close(); err != nil {
-			t.Fatalf("Error while closing reader %v", err)
-		}
-	})
-}
-
-func TestIntegration_MultiRangeDownloaderSuddenClose(t *testing.T) {
-	multiTransportTest(skipHTTP("gRPC implementation specific test"), t, func(t *testing.T, ctx context.Context, bucket string, _ string, client *Client) {
-		content := make([]byte, 5<<20)
-		rand.New(rand.NewSource(0)).Read(content)
-		objName := "MultiRangeDownloader"
-
-		// Upload test data.
-		obj := client.Bucket(bucket).Object(objName)
-		if err := writeObject(ctx, obj, "text/plain", content); err != nil {
-			t.Fatal(err)
-		}
-		defer func() {
-			if err := obj.Delete(ctx); err != nil {
-				log.Printf("failed to delete test object: %v", err)
-			}
-		}()
-		reader, err := obj.NewMultiRangeDownloader(ctx)
-		if err != nil {
-			t.Fatalf("NewMultiRangeDownloader: %v", err)
-		}
-		res := make([]multiRangeDownloaderOutput, 3)
-		callback := func(x, y int64, err error) {
-			res[0].offset = x
-			res[0].limit = y
-			res[0].err = err
-		}
-		callback1 := func(x, y int64, err error) {
-			res[1].offset = x
-			res[1].limit = y
-			res[1].err = err
-		}
-		callback2 := func(x, y int64, err error) {
-			res[2].offset = x
-			res[2].limit = y
-			res[2].err = err
-		}
-		// Add three ranges on the reader, and then do a sudden close.
-		reader.Add(&res[0].buf, 0, int64(len(content)), callback)
-		reader.Close()
-		reader.Add(&res[1].buf, -10, 0, callback1)
-		reader.Add(&res[2].buf, 0, 10, callback2)
-		// we can get stream is closed, can't add range error in case process is over before we add the range.
-		expErr := fmt.Errorf("stream is closed, can't add range")
-		expErr2 := fmt.Errorf("stream closed early")
-		for _, k := range res {
-			if k.err.Error() != expErr.Error() && k.err.Error() != expErr2.Error() {
-				t.Fatalf("read range %v to %v: got error %v, want stream closed error", k.offset, k.limit, k.err)
-			}
-		}
-	})
-}
-
 // Test in a GCE environment expected to be located in one of:
 // - us-west1-a, us-west1-b, us-west-c
 //