chore: attempt to drain gRPC ReadObject stream iterator (#2700)

BenWhitehead · web-flow · commit a6790cbe5cdb · 2024-09-10T18:46:39.000-04:00
Related #2696
diff --git a/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicUnbufferedReadableByteChannel.java b/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicUnbufferedReadableByteChannel.java
@@ -182,6 +182,7 @@ public boolean hasNext() {
         if (!result.isDone()) {
           result.setException(StorageException.coalesce(e));
         }
+        reset();
         throw e;
       }
     }
@@ -194,15 +195,31 @@ public ReadObjectResponse next() {
         if (!result.isDone()) {
           result.setException(StorageException.coalesce(e));
         }
+        reset();
         throw e;
       }
     }
 
     @Override
     public void close() {
       if (serverStream != null) {
-        // todo: do we need to "drain" anything?
         serverStream.cancel();
+        if (responseIterator != null) {
+          IOException ioException = null;
+          while (responseIterator.hasNext()) {
+            try {
+              ReadObjectResponse next = responseIterator.next();
+              ResponseContentLifecycleHandle handle = rclm.get(next);
+              handle.close();
+            } catch (IOException e) {
+              if (ioException == null) {
+                ioException = e;
+              } else if (ioException != e) {
+                ioException.addSuppressed(e);
+              }
+            }
+          }
+        }
       }
     }
 
@@ -223,5 +240,11 @@ private Iterator<ReadObjectResponse> ensureResponseIteratorOpen() {
         }
       }
     }
+
+    private void reset() {
+      serverStream = null;
+      responseIterator = null;
+      streamInitialized = false;
+    }
   }
 }
diff --git a/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java b/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java
@@ -1098,7 +1098,7 @@ static void closeAllStreams(Iterable<InputStream> inputStreams) throws IOExcepti
         } catch (IOException e) {
           if (ioException == null) {
             ioException = e;
-          } else {
+          } else if (ioException != e) {
             ioException.addSuppressed(e);
           }
         }

Original file line number	Diff line number	Diff line change
`@@ -182,6 +182,7 @@ public boolean hasNext() {`
`182`	`182`	`if (!result.isDone()) {`
`183`	`183`	`result.setException(StorageException.coalesce(e));`
`184`	`184`	`}`
	`185`	`+ reset();`
`185`	`186`	`throw e;`
`186`	`187`	`}`
`187`	`188`	`}`
`@@ -194,15 +195,31 @@ public ReadObjectResponse next() {`
`194`	`195`	`if (!result.isDone()) {`
`195`	`196`	`result.setException(StorageException.coalesce(e));`
`196`	`197`	`}`
	`198`	`+ reset();`
`197`	`199`	`throw e;`
`198`	`200`	`}`
`199`	`201`	`}`
`200`	`202`
`201`	`203`	`@Override`
`202`	`204`	`public void close() {`
`203`	`205`	`if (serverStream != null) {`
`204`		`- // todo: do we need to "drain" anything?`
`205`	`206`	`serverStream.cancel();`
	`207`	`+ if (responseIterator != null) {`
	`208`	`+ IOException ioException = null;`
	`209`	`+ while (responseIterator.hasNext()) {`
	`210`	`+ try {`
	`211`	`+ ReadObjectResponse next = responseIterator.next();`
	`212`	`+ ResponseContentLifecycleHandle handle = rclm.get(next);`
	`213`	`+ handle.close();`
	`214`	`+ } catch (IOException e) {`
	`215`	`+ if (ioException == null) {`
	`216`	`+ ioException = e;`
	`217`	`+ } else if (ioException != e) {`
	`218`	`+ ioException.addSuppressed(e);`
	`219`	`+ }`
	`220`	`+ }`
	`221`	`+ }`
	`222`	`+ }`
`206`	`223`	`}`
`207`	`224`	`}`
`208`	`225`
`@@ -223,5 +240,11 @@ private Iterator<ReadObjectResponse> ensureResponseIteratorOpen() {`
`223`	`240`	`}`
`224`	`241`	`}`
`225`	`242`	`}`
	`243`	`+`
	`244`	`+ private void reset() {`
	`245`	`+ serverStream = null;`
	`246`	`+ responseIterator = null;`
	`247`	`+ streamInitialized = false;`
	`248`	`+ }`
`226`	`249`	`}`
`227`	`250`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1098,7 +1098,7 @@ static void closeAllStreams(Iterable<InputStream> inputStreams) throws IOExcepti`
`1098`	`1098`	`} catch (IOException e) {`
`1099`	`1099`	`if (ioException == null) {`
`1100`	`1100`	`ioException = e;`
`1101`		`- } else {`
	`1101`	`+ } else if (ioException != e) {`
`1102`	`1102`	`ioException.addSuppressed(e);`
`1103`	`1103`	`}`
`1104`	`1104`	`}`