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 index 4608dc942a..998549b923 100644 --- 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,8 @@ public void run() { } finally { lock.unlock(); } - cleanupInflightRequests(); + cleanup( + /* avoidBlocking= */ true); // don't perform blocking operations while on user thread }); this.appendThread.start(); } @@ -812,7 +813,10 @@ private void appendLoop() { this.streamConnection.send(originalRequestBuilder.build()); } } + cleanup(/* avoidBlocking= */ false); + } + private void cleanup(boolean avoidBlocking) { log.info( "Cleanup starts. Stream: " + streamName @@ -828,7 +832,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 (!avoidBlocking) { + 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 index 71e4d47673..26214b5cd1 100644 --- 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