chore: make BlobDescriptorStream re-open logic smarter

BenWhitehead · BenWhitehead · commit a4d919ed36d3 · 2025-03-14T16:50:07.000-04:00
If #send(BidiReadObjectRequest) is called and the stream isn't open, automatically add the necessary BidiReadObjectSpec into the request.
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/BlobDescriptorState.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/BlobDescriptorState.java
@@ -28,6 +28,7 @@
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.ReentrantLock;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -39,6 +40,7 @@ final class BlobDescriptorState {
   private final AtomicReference<@MonotonicNonNull Object> metadata;
   private final AtomicLong readIdSeq;
   private final Map<Long, BlobDescriptorStreamRead> outstandingReads;
+  private final ReentrantLock lock;
 
   BlobDescriptorState(BidiReadObjectRequest openRequest) {
     this.openRequest = openRequest;
@@ -47,9 +49,16 @@ final class BlobDescriptorState {
     this.metadata = new AtomicReference<>();
     this.readIdSeq = new AtomicLong(1);
     this.outstandingReads = new HashMap<>();
+    this.lock = new ReentrantLock();
   }
 
   BidiReadObjectRequest getOpenRequest() {
+    Object obj = metadata.get();
+    if (obj != null && obj.getGeneration() != openRequest.getReadObjectSpec().getGeneration()) {
+      BidiReadObjectRequest.Builder b = openRequest.toBuilder();
+      b.getReadObjectSpecBuilder().setGeneration(obj.getGeneration());
+      return b.build();
+    }
     return openRequest;
   }
 
@@ -76,20 +85,29 @@ long newReadId() {
 
   @Nullable
   BlobDescriptorStreamRead getOutstandingRead(long key) {
-    synchronized (this) {
+    lock.lock();
+    try {
       return outstandingReads.get(key);
+    } finally {
+      lock.unlock();
     }
   }
 
   void putOutstandingRead(long key, BlobDescriptorStreamRead value) {
-    synchronized (this) {
+    lock.lock();
+    try {
       outstandingReads.put(key, value);
+    } finally {
+      lock.unlock();
     }
   }
 
   void removeOutstandingRead(long key) {
-    synchronized (this) {
+    lock.lock();
+    try {
       outstandingReads.remove(key);
+    } finally {
+      lock.unlock();
     }
   }
 
@@ -103,22 +121,28 @@ String getRoutingToken() {
   }
 
   BlobDescriptorStreamRead assignNewReadId(long oldReadId) {
-    synchronized (this) {
+    lock.lock();
+    try {
       BlobDescriptorStreamRead remove = outstandingReads.remove(oldReadId);
-      checkState(remove != null, "unable to locate old");
+      checkState(remove != null, "unable to locate old read");
       long newReadId = newReadId();
       BlobDescriptorStreamRead withNewReadId = remove.withNewReadId(newReadId);
       outstandingReads.put(newReadId, withNewReadId);
       return withNewReadId;
+    } finally {
+      lock.unlock();
     }
   }
 
   List<ReadRange> getOutstandingReads() {
-    synchronized (this) {
+    lock.lock();
+    try {
       return outstandingReads.values().stream()
           .filter(BlobDescriptorStreamRead::readyToSend)
           .map(BlobDescriptorStreamRead::makeReadRange)
           .collect(ImmutableList.toImmutableList());
+    } finally {
+      lock.unlock();
     }
   }
 }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/BlobDescriptorStream.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/BlobDescriptorStream.java
@@ -127,7 +127,9 @@ public void close() throws IOException {
   public void send(BidiReadObjectRequest request) {
     checkOpen();
     if (requestStream == null) {
-      restart();
+      BidiReadObjectRequest merged =
+          state.getOpenRequest().toBuilder().clearReadRanges().mergeFrom(request).build();
+      getRequestStream().send(merged);
     } else {
       getRequestStream().send(request);
     }
@@ -369,7 +371,7 @@ private OnSuccess restartReadFromCurrentOffset(long id) {
         BlobDescriptorStreamRead readWithNewId = state.assignNewReadId(id);
         BidiReadObjectRequest requestWithNewReadId =
             BidiReadObjectRequest.newBuilder().addReadRanges(readWithNewId.makeReadRange()).build();
-        requestStream.send(requestWithNewReadId);
+        BlobDescriptorStream.this.send(requestWithNewReadId);
       };
     }
 
@@ -378,11 +380,11 @@ public void onComplete() {}
   }
 
   private class MonitoringResponseObserver implements ResponseObserver<BidiReadObjectResponse> {
-    private final BlobDescriptorStream.BidiReadObjectResponseObserver delegate;
+    private final ResponseObserver<BidiReadObjectResponse> delegate;
     private final SettableApiFuture<Void> openSignal;
     private final SettableApiFuture<Void> closeSignal;
 
-    private MonitoringResponseObserver(BidiReadObjectResponseObserver delegate) {
+    private MonitoringResponseObserver(ResponseObserver<BidiReadObjectResponse> delegate) {
       this.delegate = delegate;
       this.openSignal = SettableApiFuture.create();
       this.closeSignal = SettableApiFuture.create();
