fix(storage): fix MultiRangeDownloader deadlock (#12548)

tritone · web-flow · commit 2eb23bb01ffe · 2025-07-15T20:52:27.000Z
This fixes some potential deadlocks in MultiRangeDownloader which were caused by not releasing a lock on early return, or holding it when it wasn't needed (when sending on a chan).
diff --git a/storage/client_test.go b/storage/client_test.go
@@ -1564,8 +1564,8 @@ func TestMultiRangeDownloaderEmulated(t *testing.T) {
 		reader.Wait()
 		for _, k := range res {
 			if !bytes.Equal(k.buf.Bytes(), content[k.offset:k.offset+k.limit]) {
-				t.Errorf("Error in read range offset %v, limit %v, got: %v; want: %v",
-					k.offset, k.limit, k.buf.Bytes(), content[k.offset:k.offset+k.limit])
+				t.Errorf("Error in read range offset %v, limit %v, got: %v bytes; want: %v bytes",
+					k.offset, k.limit, len(k.buf.Bytes()), len(content[k.offset:k.offset+k.limit]))
 			}
 			if k.err != nil {
 				t.Errorf("read range %v to %v : %v", k.offset, k.limit, k.err)
diff --git a/storage/grpc_client.go b/storage/grpc_client.go
@@ -1269,6 +1269,7 @@ func (c *grpcStorageClient) NewMultiRangeDownloader(ctx context.Context, params
 				if err == nil {
 					mrd.mu.Lock()
 					if len(mrd.activeRanges) == 0 && mrd.numActiveRanges == 0 {
+						mrd.mu.Unlock()
 						mrd.closeReceiver <- true
 						mrd.closeSender <- true
 						return
@@ -1277,34 +1278,37 @@ func (c *grpcStorageClient) NewMultiRangeDownloader(ctx context.Context, params
 					arr := resp.GetObjectDataRanges()
 					for _, val := range arr {
 						id := val.GetReadRange().GetReadId()
-						mrd.mu.Lock()
-						_, ok := mrd.activeRanges[id]
-						if !ok {
-							// it's ok to ignore responses for read_id not in map as user would have been notified by callback.
-							continue
-						}
-						_, err = mrd.activeRanges[id].writer.Write(val.GetChecksummedData().GetContent())
-						if err != nil {
-							mrd.activeRanges[id].callback(mrd.activeRanges[id].offset, mrd.activeRanges[id].totalBytesWritten, err)
-							mrd.numActiveRanges--
-							delete(mrd.activeRanges, id)
-						} else {
-							mrd.activeRanges[id] = mrdRange{
-								readID:              mrd.activeRanges[id].readID,
-								writer:              mrd.activeRanges[id].writer,
-								offset:              mrd.activeRanges[id].offset,
-								limit:               mrd.activeRanges[id].limit,
-								currentBytesWritten: mrd.activeRanges[id].currentBytesWritten + int64(len(val.GetChecksummedData().GetContent())),
-								totalBytesWritten:   mrd.activeRanges[id].totalBytesWritten + int64(len(val.GetChecksummedData().GetContent())),
-								callback:            mrd.activeRanges[id].callback,
+						func() {
+							mrd.mu.Lock()
+							defer mrd.mu.Unlock()
+							currRange, ok := mrd.activeRanges[id]
+							if !ok {
+								// it's ok to ignore responses for read_id not in map as user would have been notified by callback.
+								return
 							}
-						}
-						if val.GetRangeEnd() {
-							mrd.activeRanges[id].callback(mrd.activeRanges[id].offset, mrd.activeRanges[id].totalBytesWritten, nil)
-							mrd.numActiveRanges--
-							delete(mrd.activeRanges, id)
-						}
-						mrd.mu.Unlock()
+							_, err = currRange.writer.Write(val.GetChecksummedData().GetContent())
+							if err != nil {
+								currRange.callback(currRange.offset, currRange.totalBytesWritten, err)
+								mrd.numActiveRanges--
+								delete(mrd.activeRanges, id)
+							} else {
+								currRange = mrdRange{
+									readID:              currRange.readID,
+									writer:              currRange.writer,
+									offset:              currRange.offset,
+									limit:               currRange.limit,
+									currentBytesWritten: currRange.currentBytesWritten + int64(len(val.GetChecksummedData().GetContent())),
+									totalBytesWritten:   currRange.totalBytesWritten + int64(len(val.GetChecksummedData().GetContent())),
+									callback:            currRange.callback,
+								}
+								mrd.activeRanges[id] = currRange
+							}
+							if val.GetRangeEnd() {
+								currRange.callback(currRange.offset, currRange.totalBytesWritten, nil)
+								mrd.numActiveRanges--
+								delete(mrd.activeRanges, id)
+							}
+						}()
 					}
 				}
 			}
@@ -1459,8 +1463,8 @@ func (mrd *gRPCBidiReader) add(output io.Writer, offset, limit int64, callback f
 		spec := mrdRange{readID: id, writer: output, offset: offset, limit: limit, currentBytesWritten: 0, totalBytesWritten: 0, callback: callback}
 		mrd.mu.Lock()
 		mrd.numActiveRanges++
-		mrd.rangesToRead <- []mrdRange{spec}
 		mrd.mu.Unlock()
+		mrd.rangesToRead <- []mrdRange{spec}
 	} else {
 		callback(offset, 0, errors.New("storage: cannot add range because the stream is closed"))
 	}
diff --git a/storage/integration_test.go b/storage/integration_test.go
@@ -416,6 +416,45 @@ func TestIntegration_MultiRangeDownloader(t *testing.T) {
 	})
 }
 
+// Test many concurrent reads on the same MultiRangeDownloader to try to detect
+// potential deadlocks.
+func TestIntegration_MRDManyReads(t *testing.T) {
+	multiTransportTest(skipAllButBidi(context.Background(), "Bidi Read API 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 := "MultiRangeDownloaderManyReads"
+		// 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)
+		}
+		// Previously 100 ranges here worked in a few seconds, but 1000 caused
+		// a deadlock. After this PR, 1000 also works in under 10s.
+		// 1000 is larger than the size of the buffer for the new ranges
+		// to be added.
+		for i := 0; i < 1000; i++ {
+			reader.Add(io.Discard, 0, 100, func(_ int64, _ int64, err error) {
+				if err != nil {
+					t.Errorf("error in call %v: %v", i, err)
+				}
+			})
+		}
+		reader.Wait()
+		if err = reader.Close(); err != nil {
+			t.Fatalf("Error while closing reader %v", err)
+		}
+	})
+}
+
 // TestIntegration_MRDCallbackReturnsDataLength tests if the callback returns the correct data
 // read length or not.
 func TestIntegration_MRDCallbackReturnsDataLength(t *testing.T) {
