From 1380ee5984f35a1ee259a76556533b3b2d1a6481 Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 6 Jul 2026 19:26:19 +0000 Subject: [PATCH 1/4] fix(bigquerystorage): fix flakiness in ConnectionWorkerPoolTest without sleeping The test `testMultiStreamClosed_multiplexingEnabled` was flaky because it relied on requests not finishing during the scaling phase. The mock server had a fixed response delay, which could be shorter than the time it took to complete the request submission loop on slow machines, causing requests to finish early and preventing the pool from scaling up. Instead of relying on timing and sleeps, this change introduces a `CountDownLatch` to synchronize the mock server responses. A new helper `BlockingResponseSupplier` blocks the mock server responses until the test thread has finished sending all requests. This guarantees that all requests remain in-flight during the scaling phase, making the test 100% stable and faster. Fixes https://github.com/googleapis/google-cloud-java/issues/13664 TAG=agy CONV=76d8911b-bc84-4f0a-8a39-dceadbb18cf7 --- .../storage/v1/ConnectionWorkerPoolTest.java | 37 ++++++++++++++++--- 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java index 01de76298955..ad84daa7f5e6 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java @@ -40,6 +40,7 @@ import java.util.HashSet; import java.util.List; import java.util.UUID; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -229,15 +230,14 @@ void testMultiStreamClosed_multiplexingEnabled() throws Exception { createConnectionWorkerPool( /* maxRequests= */ 3, /* maxBytes= */ 1000, java.time.Duration.ofSeconds(5)); - // Sets the sleep time to simulate requests stuck in connection. - testBigQueryWrite.setResponseSleep(Duration.ofMillis(50L)); StreamWriter writeStream1 = getTestStreamWriter(TEST_STREAM_1); StreamWriter writeStream2 = getTestStreamWriter(TEST_STREAM_2); // Try append 20 requests, at the end we should have 2 requests per connection. long appendCount = 20; + CountDownLatch latch = new CountDownLatch(1); for (long i = 0; i < appendCount; i++) { - testBigQueryWrite.addResponse(createAppendResponse(i)); + testBigQueryWrite.addResponse(new BlockingResponseSupplier(createAppendResponse(i), latch)); } List> futures = new ArrayList<>(); @@ -253,13 +253,17 @@ void testMultiStreamClosed_multiplexingEnabled() throws Exception { writeStream, connectionWorkerPool, new String[] {String.valueOf(i)}, i)); } - for (ApiFuture future : futures) { - future.get(); - } // At the end we should scale up to 10 connections. assertThat(connectionWorkerPool.getCreateConnectionCount()).isEqualTo(10); assertThat(connectionWorkerPool.getTotalConnectionCount()).isEqualTo(10); + // Release the latch to allow requests to complete. + latch.countDown(); + + for (ApiFuture future : futures) { + future.get(); + } + // Start testing calling close on each stream. // When we close the first stream, only the connection that only serve stream 1 will be closed. // for which c1 and c3 are closed. @@ -400,6 +404,27 @@ public FakeBigQueryWriteImpl.Response get() { } } + private class BlockingResponseSupplier implements Supplier { + private final AppendRowsResponse response; + private final CountDownLatch latch; + + BlockingResponseSupplier(AppendRowsResponse response, CountDownLatch latch) { + this.response = response; + this.latch = latch; + } + + @Override + public FakeBigQueryWriteImpl.Response get() { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + return new FakeBigQueryWriteImpl.Response(response); + } + } + @Test void testAppendWithRetry() throws Exception { ConnectionWorkerPool connectionWorkerPool = From 70c5e42fd2be26b4ab6a1001c23b353abfce6c7d Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 6 Jul 2026 19:31:10 +0000 Subject: [PATCH 2/4] refactor(bigquerystorage): use lambda instead of class in ConnectionWorkerPoolTest Refactor the flaky test fix to use a lambda expression for the mock server response supplier, removing the need for the helper class BlockingResponseSupplier. TAG=agy CONV=76d8911b-bc84-4f0a-8a39-dceadbb18cf7 --- .../storage/v1/ConnectionWorkerPoolTest.java | 33 +++++++------------ 1 file changed, 11 insertions(+), 22 deletions(-) diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java index ad84daa7f5e6..ee29c1c56731 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java @@ -237,7 +237,17 @@ void testMultiStreamClosed_multiplexingEnabled() throws Exception { long appendCount = 20; CountDownLatch latch = new CountDownLatch(1); for (long i = 0; i < appendCount; i++) { - testBigQueryWrite.addResponse(new BlockingResponseSupplier(createAppendResponse(i), latch)); + long offset = i; + testBigQueryWrite.addResponse( + () -> { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + return new FakeBigQueryWriteImpl.Response(createAppendResponse(offset)); + }); } List> futures = new ArrayList<>(); @@ -404,27 +414,6 @@ public FakeBigQueryWriteImpl.Response get() { } } - private class BlockingResponseSupplier implements Supplier { - private final AppendRowsResponse response; - private final CountDownLatch latch; - - BlockingResponseSupplier(AppendRowsResponse response, CountDownLatch latch) { - this.response = response; - this.latch = latch; - } - - @Override - public FakeBigQueryWriteImpl.Response get() { - try { - latch.await(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException(e); - } - return new FakeBigQueryWriteImpl.Response(response); - } - } - @Test void testAppendWithRetry() throws Exception { ConnectionWorkerPool connectionWorkerPool = From 8a6e462d05b453bcf71788b4ca6eafa4e4c1409f Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 6 Jul 2026 19:33:18 +0000 Subject: [PATCH 3/4] test(bigquerystorage): add comment explaining CountDownLatch in ConnectionWorkerPoolTest Add a comment explaining the purpose of using CountDownLatch in testMultiStreamClosed_multiplexingEnabled. TAG=agy CONV=76d8911b-bc84-4f0a-8a39-dceadbb18cf7 --- .../cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java | 3 +++ 1 file changed, 3 insertions(+) diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java index ee29c1c56731..5293d07c9bdc 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java @@ -235,6 +235,9 @@ void testMultiStreamClosed_multiplexingEnabled() throws Exception { // Try append 20 requests, at the end we should have 2 requests per connection. long appendCount = 20; + // Use a CountDownLatch to block mock server responses. This ensures all 20 requests + // remain in-flight simultaneously, forcing the ConnectionWorkerPool to scale up to + // 10 connections. Without this blocking, requests might finish early and prevent scaling. CountDownLatch latch = new CountDownLatch(1); for (long i = 0; i < appendCount; i++) { long offset = i; From c34336afc7fb86ccc09ff65e440fb42f0115ff3c Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Mon, 6 Jul 2026 19:42:09 +0000 Subject: [PATCH 4/4] test(bigquerystorage): wrap assertions in try-finally to guarantee latch release Wrap connection count assertions in a try-finally block to ensure that latch.countDown() is always called, preventing mock server threads from blocking indefinitely if the assertions fail. TAG=agy CONV=76d8911b-bc84-4f0a-8a39-dceadbb18cf7 --- .../storage/v1/ConnectionWorkerPoolTest.java | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java index 5293d07c9bdc..30a78e0475a7 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java @@ -267,11 +267,13 @@ void testMultiStreamClosed_multiplexingEnabled() throws Exception { } // At the end we should scale up to 10 connections. - assertThat(connectionWorkerPool.getCreateConnectionCount()).isEqualTo(10); - assertThat(connectionWorkerPool.getTotalConnectionCount()).isEqualTo(10); - - // Release the latch to allow requests to complete. - latch.countDown(); + try { + assertThat(connectionWorkerPool.getCreateConnectionCount()).isEqualTo(10); + assertThat(connectionWorkerPool.getTotalConnectionCount()).isEqualTo(10); + } finally { + // Release the latch to allow requests to complete. + latch.countDown(); + } for (ApiFuture future : futures) { future.get();