fix: also shutdown the stream connection in case the timeout exception is

agrawal-siddharth · agrawal-siddharth · commit d62a1d3cd28e · 2024-03-09T06:39:16.000-08:00
triggered.
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java
@@ -353,7 +353,7 @@ public void run() {
           } finally {
             lock.unlock();
           }
-          cleanupInflightRequests();
+          cleanup(/* waitForDone= */ false);
         });
     this.appendThread.start();
   }
@@ -812,7 +812,10 @@ private void appendLoop() {
         this.streamConnection.send(originalRequestBuilder.build());
       }
     }
+    cleanup(/* waitForDone= */ true);
+  }
 
+  private void cleanup(boolean waitForDone) {
     log.info(
         "Cleanup starts. Stream: "
             + streamName
@@ -828,7 +831,9 @@ private void appendLoop() {
     // We can close the stream connection and handle the remaining inflight requests.
     if (streamConnection != null) {
       this.streamConnection.close();
-      waitForDoneCallback(3, TimeUnit.MINUTES);
+      if (waitForDone) {
+        waitForDoneCallback(3, TimeUnit.MINUTES);
+      }
     }
 
     // At this point, there cannot be more callback. It is safe to clean up all inflight requests.
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerTest.java
@@ -650,9 +650,9 @@ public void testThrowExceptionWhileWithinAppendLoop_MaxWaitTimeExceed() throws E
             null,
             client.getSettings(),
             retrySettings);
-    testBigQueryWrite.setResponseSleep(org.threeten.bp.Duration.ofSeconds(3));
+    testBigQueryWrite.setResponseSleep(org.threeten.bp.Duration.ofSeconds(2));
 
-    long appendCount = 10;
+    long appendCount = 2;
     for (int i = 0; i < appendCount; i++) {
       testBigQueryWrite.addResponse(createAppendResponse(i));
     }
@@ -691,6 +691,10 @@ public void testThrowExceptionWhileWithinAppendLoop_MaxWaitTimeExceed() throws E
                         100)
                     .get());
     assertThat(ex.getCause()).hasMessageThat().contains("Request has waited in inflight queue");
+
+    // Verify we can shutdown normally.
+    connectionWorker.close();
+    assertTrue(connectionWorker.isUserClosed());
   }
 
   @Test

Original file line number	Diff line number	Diff line change
`@@ -353,7 +353,7 @@ public void run() {`
`353`	`353`	`} finally {`
`354`	`354`	`lock.unlock();`
`355`	`355`	`}`
`356`		`- cleanupInflightRequests();`
	`356`	`+ cleanup(/* waitForDone= */ false);`
`357`	`357`	`});`
`358`	`358`	`this.appendThread.start();`
`359`	`359`	`}`
`@@ -812,7 +812,10 @@ private void appendLoop() {`
`812`	`812`	`this.streamConnection.send(originalRequestBuilder.build());`
`813`	`813`	`}`
`814`	`814`	`}`
	`815`	`+ cleanup(/* waitForDone= */ true);`
	`816`	`+ }`
`815`	`817`
	`818`	`+ private void cleanup(boolean waitForDone) {`
`816`	`819`	`log.info(`
`817`	`820`	`"Cleanup starts. Stream: "`
`818`	`821`	`+ streamName`
`@@ -828,7 +831,9 @@ private void appendLoop() {`
`828`	`831`	`// We can close the stream connection and handle the remaining inflight requests.`
`829`	`832`	`if (streamConnection != null) {`
`830`	`833`	`this.streamConnection.close();`
`831`		`- waitForDoneCallback(3, TimeUnit.MINUTES);`
	`834`	`+ if (waitForDone) {`
	`835`	`+ waitForDoneCallback(3, TimeUnit.MINUTES);`
	`836`	`+ }`
`832`	`837`	`}`
`833`	`838`
`834`	`839`	`// At this point, there cannot be more callback. It is safe to clean up all inflight requests.`