fix: remove unrecoverable connection from connection pool during multiplexing

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java

@@ -378,6 +378,11 @@ public String getWriterId() {
  return writerId;
  }
 
+ boolean isConnectionInUnrecoverableState() {
+ // If final status is set, there's no
+ return connectionFinalStatus != null;
+ }
+
  /** Close the stream writer. Shut down all resources. */
  @Override
  public void close() {

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPool.java

@@ -234,9 +234,17 @@ ApiFuture<AppendRowsResponse> append(StreamWriter streamWriter, ProtoRows rows,
  streamWriter,
  (key, existingStream) -> {
  // Stick to the existing stream if it's not overwhelmed.
- if (existingStream != null && !existingStream.getLoad().isOverwhelmed()) {
+ if (existingStream != null
+ && !existingStream.getLoad().isOverwhelmed()
+ && !existingStream.isConnectionInUnrecoverableState()) {
  return existingStream;
  }
+ if (existingStream != null && existingStream.isConnectionInUnrecoverableState()) {
+ existingStream = null;
+ }
+ // Before search for the next connection to attach, clear the finalized connections
+ // first so that they will not be selected.
+ clearFinalizedConnectionWorker();
  // Try to create or find another existing stream to reuse.
  ConnectionWorker createdOrExistingConnection = null;
  try {
@@ -299,7 +307,6 @@ private ConnectionWorker createOrReuseConnectionWorker(
  }
  return createConnectionWorker(streamWriter.getStreamName(), streamWriter.getProtoSchema());
  } else {
-
  // Stick to the original connection if all the connections are overwhelmed.
  if (existingConnectionWorker != null) {
  return existingConnectionWorker;
@@ -310,6 +317,18 @@ private ConnectionWorker createOrReuseConnectionWorker(
  }
  }
 
+ private void clearFinalizedConnectionWorker() {
+ Set<ConnectionWorker> connectionWorkerSet = new HashSet<>();
+ for (ConnectionWorker existingWorker : connectionWorkerPool) {
+ if (existingWorker.isConnectionInUnrecoverableState()) {
+ connectionWorkerSet.add(existingWorker);
+ }
+ }
+ for (ConnectionWorker workerToRemove : connectionWorkerSet) {
+ connectionWorkerPool.remove(workerToRemove);
+ }
+ }
+
  /** Select out the best connection worker among the given connection workers. */
  static ConnectionWorker pickBestLoadConnection(
  Comparator<Load> comparator, List<ConnectionWorker> connectionWorkerList) {

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java

@@ -31,6 +31,7 @@
 import java.io.IOException;
 import java.time.Duration;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
@@ -450,6 +451,15 @@ static void cleanUp() {
  connectionPoolMap.clear();
  }
 
+ @VisibleForTesting
+ ConnectionWorkerPool getTestOnlyConnectionWorkerPool() {
+ ConnectionWorkerPool connectionWorkerPool = null;
+ for (Entry<ConnectionPoolKey, ConnectionWorkerPool> entry : connectionPoolMap.entrySet()) {
+ connectionWorkerPool = entry.getValue();
+ }
+ return connectionWorkerPool;
+ }
+
  /** A builder of {@link StreamWriter}s. */
  public static final class Builder {
  private static final long DEFAULT_MAX_INFLIGHT_REQUESTS = 1000L;

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java

@@ -69,6 +69,7 @@ public class StreamWriterTest {
  private static final Logger log = Logger.getLogger(StreamWriterTest.class.getName());
  private static final String TEST_STREAM_1 = "projects/p/datasets/d1/tables/t1/streams/_default";
  private static final String TEST_STREAM_2 = "projects/p/datasets/d2/tables/t2/streams/_default";
+ private static final String TEST_STREAM_3 = "projects/p/datasets/d3/tables/t3/streams/_default";
  private static final String TEST_STREAM_SHORTEN = "projects/p/datasets/d2/tables/t2/_default";
  private static final String EXPLICIT_STEAM = "projects/p/datasets/d1/tables/t1/streams/s1";
  private static final String TEST_TRACE_ID = "DATAFLOW:job_id";
@@ -1090,6 +1091,115 @@ public void testExtractDatasetName() throws Exception {
  Assert.assertTrue(ex.getMessage().contains("The passed in stream name does not match"));
  }
 
+ @Test
+ public void testRetryInUnrecoverableStatus_MultiplexingCase() throws Exception {
+ ConnectionWorkerPool.setOptions(
+ Settings.builder().setMinConnectionsPerRegion(1).setMaxConnectionsPerRegion(4).build());
+ ConnectionWorkerPool.enableTestingLogic();
+
+ // Setup: create three stream writers, two of them are writing to the same stream.
+ // Those four stream writers should be assigned to the same connection.
+ // 1. Submit three requests at first to trigger connection retry limitation.
+ // 2. At this point the connection should be entering a unrecoverable state.
+ // 3. Further submit requests to those stream writers would trigger connection reassignment.
+ StreamWriter writer1 = getMultiplexingStreamWriter(TEST_STREAM_1);
+ StreamWriter writer2 = getMultiplexingStreamWriter(TEST_STREAM_2);
+ StreamWriter writer3 = getMultiplexingStreamWriter(TEST_STREAM_3);
+ StreamWriter writer4 = getMultiplexingStreamWriter(TEST_STREAM_3);
+
+ testBigQueryWrite.setCloseForeverAfter(2);
+ testBigQueryWrite.setTimesToClose(1);
+ testBigQueryWrite.addResponse(createAppendResponse(0));
+ testBigQueryWrite.addResponse(createAppendResponse(1));
+
+ // Connection will be failed after triggering the third append.
+ ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer1, new String[] {"A"}, 0);
+ ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer2, new String[] {"B"}, 1);
+ ApiFuture<AppendRowsResponse> appendFuture3 = sendTestMessage(writer3, new String[] {"C"}, 2);
+ TimeUnit.SECONDS.sleep(1);
+ assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue());
+ assertEquals(1, appendFuture2.get().getAppendResult().getOffset().getValue());
+ assertThrows(
+ ExecutionException.class,
+ () -> {
+ assertEquals(2, appendFuture3.get().getAppendResult().getOffset().getValue());
+ });
+ assertEquals(writer1.getTestOnlyConnectionWorkerPool().getTotalConnectionCount(), 1);
+ assertEquals(writer1.getTestOnlyConnectionWorkerPool().getCreateConnectionCount(), 1);
+
+ // Insert another request to the writer attached to closed connection would create another
+ // connection.
+
+ testBigQueryWrite.setCloseForeverAfter(0);
+ testBigQueryWrite.addResponse(createAppendResponse(4));
+ testBigQueryWrite.addResponse(createAppendResponse(5));
+ testBigQueryWrite.addResponse(createAppendResponse(6));
+ ApiFuture<AppendRowsResponse> appendFuture4 = sendTestMessage(writer4, new String[] {"A"}, 2);
+ ApiFuture<AppendRowsResponse> appendFuture5 = sendTestMessage(writer1, new String[] {"A"}, 3);
+ ApiFuture<AppendRowsResponse> appendFuture6 = sendTestMessage(writer2, new String[] {"B"}, 4);
+ assertEquals(4, appendFuture4.get().getAppendResult().getOffset().getValue());
+ assertEquals(5, appendFuture5.get().getAppendResult().getOffset().getValue());
+ assertEquals(6, appendFuture6.get().getAppendResult().getOffset().getValue());
+ assertEquals(writer1.getTestOnlyConnectionWorkerPool().getTotalConnectionCount(), 1);
+ assertEquals(writer1.getTestOnlyConnectionWorkerPool().getCreateConnectionCount(), 2);
+
+ writer1.close();
+ writer2.close();
+ writer3.close();
+ writer4.close();
+ assertEquals(writer1.getTestOnlyConnectionWorkerPool().getTotalConnectionCount(), 0);
+ }
+
+ @Test
+ public void testCloseWhileInUnrecoverableState() throws Exception {
+ ConnectionWorkerPool.setOptions(
+ Settings.builder().setMinConnectionsPerRegion(1).setMaxConnectionsPerRegion(4).build());
+ ConnectionWorkerPool.enableTestingLogic();
+
+ // Setup: create three stream writers
+ // 1. Submit three requests at first to trigger connection retry limitation.
+ // 2. Submit request to writer3 to trigger reassignment
+ // 3. Close the previous two writers would be succesful
+ StreamWriter writer1 = getMultiplexingStreamWriter(TEST_STREAM_1);
+ StreamWriter writer2 = getMultiplexingStreamWriter(TEST_STREAM_2);
+ StreamWriter writer3 = getMultiplexingStreamWriter(TEST_STREAM_3);
+
+ testBigQueryWrite.setCloseForeverAfter(2);
+ testBigQueryWrite.setTimesToClose(1);
+ testBigQueryWrite.addResponse(createAppendResponse(0));
+ testBigQueryWrite.addResponse(createAppendResponse(1));
+
+ // Connection will be failed after triggering the third append.
+ ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer1, new String[] {"A"}, 0);
+ ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer2, new String[] {"B"}, 1);
+ ApiFuture<AppendRowsResponse> appendFuture3 = sendTestMessage(writer3, new String[] {"C"}, 2);
+ TimeUnit.SECONDS.sleep(1);
+ assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue());
+ assertEquals(1, appendFuture2.get().getAppendResult().getOffset().getValue());
+ assertThrows(
+ ExecutionException.class,
+ () -> {
+ assertEquals(2, appendFuture3.get().getAppendResult().getOffset().getValue());
+ });
+ assertEquals(writer1.getTestOnlyConnectionWorkerPool().getTotalConnectionCount(), 1);
+ assertEquals(writer1.getTestOnlyConnectionWorkerPool().getCreateConnectionCount(), 1);
+
+ writer1.close();
+ writer2.close();
+ // We will still be left with one request
+ assertEquals(writer1.getTestOnlyConnectionWorkerPool().getCreateConnectionCount(), 1);
+ }
+
+ public StreamWriter getMultiplexingStreamWriter(String streamName) throws IOException {
+ return StreamWriter.newBuilder(streamName, client)
+ .setWriterSchema(createProtoSchema())
+ .setEnableConnectionPool(true)
+ .setMaxInflightRequests(10)
+ .setLocation("US")
+ .setMaxRetryDuration(java.time.Duration.ofMillis(100))
+ .build();
+ }
+
  // Timeout to ensure close() doesn't wait for done callback timeout.
  @Test(timeout = 10000)
  public void testCloseDisconnectedStream() throws Exception {