Skip to content

Commit 091dddb

Browse files
GaoleMenggcf-owl-bot[bot]
GaoleMeng
and
gcf-owl-bot[bot]
authored
fix: remove unrecoverable connection from connection pool during multiplexing (#1967)
* feat: Split writer into connection worker and wrapper, this is a prerequisite for multiplexing client * feat: add connection worker pool skeleton, used for multiplexing client * feat: add Load api for connection worker for multiplexing client * feat: add multiplexing support to connection worker. We will treat every new stream name as a switch of destinationt * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * feat: port the multiplexing client core algorithm and basic tests also fixed a tiny bug inside fake bigquery write impl for getting thre response from offset * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * feat: wire multiplexing connection pool to stream writer * feat: some fixes for multiplexing client * feat: fix some todos, and reject the mixed behavior of passed in client or not * feat: fix the bug that we may peek into the write_stream field but it's possible the proto schema does not contain this field * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * feat: fix the bug that we may peek into the write_stream field but it's possible the proto schema does not contain this field * feat: add getInflightWaitSeconds implementation * feat: Add schema comparision in connection loop to ensure schema update for the same stream name can be notified * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * feat: add schema update support to multiplexing * fix: fix windows build bug: windows Instant resolution is different with linux * fix: fix another failing tests for windows build * fix: fix another test failure for Windows build * feat: Change new thread for each retry to be a thread pool to avoid create/tear down too much threads if lots of retries happens * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * fix: add back the background executor provider that's accidentally removed * feat: throw error when use connection pool for explicit stream * fix: Add precision truncation to the passed in value from JSON float and double type. * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * modify the bom version * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * fix deadlockissue in ConnectionWorkerPool * fix: fix deadlock issue during close + append for multiplexing * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * fix: fix one potential root cause of deadlock issue for non-multiplexing case * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * Add timeout to inflight queue waiting, and also add some extra log * feat: allow java client lib handle switch table schema for the same stream name * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent 305f71e commit 091dddb

File tree

4 files changed

+146
-2
lines changed

4 files changed

+146
-2
lines changed

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -378,6 +378,11 @@ public String getWriterId() {
378378
return writerId;
379379
}
380380

381+
boolean isConnectionInUnrecoverableState() {
382+
// If final status is set, there's no
383+
return connectionFinalStatus != null;
384+
}
385+
381386
/** Close the stream writer. Shut down all resources. */
382387
@Override
383388
public void close() {

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

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -234,9 +234,17 @@ ApiFuture<AppendRowsResponse> append(StreamWriter streamWriter, ProtoRows rows,
234234
streamWriter,
235235
(key, existingStream) -> {
236236
// Stick to the existing stream if it's not overwhelmed.
237-
if (existingStream != null && !existingStream.getLoad().isOverwhelmed()) {
237+
if (existingStream != null
238+
&& !existingStream.getLoad().isOverwhelmed()
239+
&& !existingStream.isConnectionInUnrecoverableState()) {
238240
return existingStream;
239241
}
242+
if (existingStream != null && existingStream.isConnectionInUnrecoverableState()) {
243+
existingStream = null;
244+
}
245+
// Before search for the next connection to attach, clear the finalized connections
246+
// first so that they will not be selected.
247+
clearFinalizedConnectionWorker();
240248
// Try to create or find another existing stream to reuse.
241249
ConnectionWorker createdOrExistingConnection = null;
242250
try {
@@ -299,7 +307,6 @@ private ConnectionWorker createOrReuseConnectionWorker(
299307
}
300308
return createConnectionWorker(streamWriter.getStreamName(), streamWriter.getProtoSchema());
301309
} else {
302-
303310
// Stick to the original connection if all the connections are overwhelmed.
304311
if (existingConnectionWorker != null) {
305312
return existingConnectionWorker;
@@ -310,6 +317,18 @@ private ConnectionWorker createOrReuseConnectionWorker(
310317
}
311318
}
312319

320+
private void clearFinalizedConnectionWorker() {
321+
Set<ConnectionWorker> connectionWorkerSet = new HashSet<>();
322+
for (ConnectionWorker existingWorker : connectionWorkerPool) {
323+
if (existingWorker.isConnectionInUnrecoverableState()) {
324+
connectionWorkerSet.add(existingWorker);
325+
}
326+
}
327+
for (ConnectionWorker workerToRemove : connectionWorkerSet) {
328+
connectionWorkerPool.remove(workerToRemove);
329+
}
330+
}
331+
313332
/** Select out the best connection worker among the given connection workers. */
314333
static ConnectionWorker pickBestLoadConnection(
315334
Comparator<Load> comparator, List<ConnectionWorker> connectionWorkerList) {

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import java.io.IOException;
3232
import java.time.Duration;
3333
import java.util.Map;
34+
import java.util.Map.Entry;
3435
import java.util.Objects;
3536
import java.util.UUID;
3637
import java.util.concurrent.ConcurrentHashMap;
@@ -450,6 +451,15 @@ static void cleanUp() {
450451
connectionPoolMap.clear();
451452
}
452453

454+
@VisibleForTesting
455+
ConnectionWorkerPool getTestOnlyConnectionWorkerPool() {
456+
ConnectionWorkerPool connectionWorkerPool = null;
457+
for (Entry<ConnectionPoolKey, ConnectionWorkerPool> entry : connectionPoolMap.entrySet()) {
458+
connectionWorkerPool = entry.getValue();
459+
}
460+
return connectionWorkerPool;
461+
}
462+
453463
/** A builder of {@link StreamWriter}s. */
454464
public static final class Builder {
455465
private static final long DEFAULT_MAX_INFLIGHT_REQUESTS = 1000L;

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

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ public class StreamWriterTest {
6969
private static final Logger log = Logger.getLogger(StreamWriterTest.class.getName());
7070
private static final String TEST_STREAM_1 = "projects/p/datasets/d1/tables/t1/streams/_default";
7171
private static final String TEST_STREAM_2 = "projects/p/datasets/d2/tables/t2/streams/_default";
72+
private static final String TEST_STREAM_3 = "projects/p/datasets/d3/tables/t3/streams/_default";
7273
private static final String TEST_STREAM_SHORTEN = "projects/p/datasets/d2/tables/t2/_default";
7374
private static final String EXPLICIT_STEAM = "projects/p/datasets/d1/tables/t1/streams/s1";
7475
private static final String TEST_TRACE_ID = "DATAFLOW:job_id";
@@ -1090,6 +1091,115 @@ public void testExtractDatasetName() throws Exception {
10901091
Assert.assertTrue(ex.getMessage().contains("The passed in stream name does not match"));
10911092
}
10921093

1094+
@Test
1095+
public void testRetryInUnrecoverableStatus_MultiplexingCase() throws Exception {
1096+
ConnectionWorkerPool.setOptions(
1097+
Settings.builder().setMinConnectionsPerRegion(1).setMaxConnectionsPerRegion(4).build());
1098+
ConnectionWorkerPool.enableTestingLogic();
1099+
1100+
// Setup: create three stream writers, two of them are writing to the same stream.
1101+
// Those four stream writers should be assigned to the same connection.
1102+
// 1. Submit three requests at first to trigger connection retry limitation.
1103+
// 2. At this point the connection should be entering a unrecoverable state.
1104+
// 3. Further submit requests to those stream writers would trigger connection reassignment.
1105+
StreamWriter writer1 = getMultiplexingStreamWriter(TEST_STREAM_1);
1106+
StreamWriter writer2 = getMultiplexingStreamWriter(TEST_STREAM_2);
1107+
StreamWriter writer3 = getMultiplexingStreamWriter(TEST_STREAM_3);
1108+
StreamWriter writer4 = getMultiplexingStreamWriter(TEST_STREAM_3);
1109+
1110+
testBigQueryWrite.setCloseForeverAfter(2);
1111+
testBigQueryWrite.setTimesToClose(1);
1112+
testBigQueryWrite.addResponse(createAppendResponse(0));
1113+
testBigQueryWrite.addResponse(createAppendResponse(1));
1114+
1115+
// Connection will be failed after triggering the third append.
1116+
ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer1, new String[] {"A"}, 0);
1117+
ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer2, new String[] {"B"}, 1);
1118+
ApiFuture<AppendRowsResponse> appendFuture3 = sendTestMessage(writer3, new String[] {"C"}, 2);
1119+
TimeUnit.SECONDS.sleep(1);
1120+
assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue());
1121+
assertEquals(1, appendFuture2.get().getAppendResult().getOffset().getValue());
1122+
assertThrows(
1123+
ExecutionException.class,
1124+
() -> {
1125+
assertEquals(2, appendFuture3.get().getAppendResult().getOffset().getValue());
1126+
});
1127+
assertEquals(writer1.getTestOnlyConnectionWorkerPool().getTotalConnectionCount(), 1);
1128+
assertEquals(writer1.getTestOnlyConnectionWorkerPool().getCreateConnectionCount(), 1);
1129+
1130+
// Insert another request to the writer attached to closed connection would create another
1131+
// connection.
1132+
1133+
testBigQueryWrite.setCloseForeverAfter(0);
1134+
testBigQueryWrite.addResponse(createAppendResponse(4));
1135+
testBigQueryWrite.addResponse(createAppendResponse(5));
1136+
testBigQueryWrite.addResponse(createAppendResponse(6));
1137+
ApiFuture<AppendRowsResponse> appendFuture4 = sendTestMessage(writer4, new String[] {"A"}, 2);
1138+
ApiFuture<AppendRowsResponse> appendFuture5 = sendTestMessage(writer1, new String[] {"A"}, 3);
1139+
ApiFuture<AppendRowsResponse> appendFuture6 = sendTestMessage(writer2, new String[] {"B"}, 4);
1140+
assertEquals(4, appendFuture4.get().getAppendResult().getOffset().getValue());
1141+
assertEquals(5, appendFuture5.get().getAppendResult().getOffset().getValue());
1142+
assertEquals(6, appendFuture6.get().getAppendResult().getOffset().getValue());
1143+
assertEquals(writer1.getTestOnlyConnectionWorkerPool().getTotalConnectionCount(), 1);
1144+
assertEquals(writer1.getTestOnlyConnectionWorkerPool().getCreateConnectionCount(), 2);
1145+
1146+
writer1.close();
1147+
writer2.close();
1148+
writer3.close();
1149+
writer4.close();
1150+
assertEquals(writer1.getTestOnlyConnectionWorkerPool().getTotalConnectionCount(), 0);
1151+
}
1152+
1153+
@Test
1154+
public void testCloseWhileInUnrecoverableState() throws Exception {
1155+
ConnectionWorkerPool.setOptions(
1156+
Settings.builder().setMinConnectionsPerRegion(1).setMaxConnectionsPerRegion(4).build());
1157+
ConnectionWorkerPool.enableTestingLogic();
1158+
1159+
// Setup: create three stream writers
1160+
// 1. Submit three requests at first to trigger connection retry limitation.
1161+
// 2. Submit request to writer3 to trigger reassignment
1162+
// 3. Close the previous two writers would be succesful
1163+
StreamWriter writer1 = getMultiplexingStreamWriter(TEST_STREAM_1);
1164+
StreamWriter writer2 = getMultiplexingStreamWriter(TEST_STREAM_2);
1165+
StreamWriter writer3 = getMultiplexingStreamWriter(TEST_STREAM_3);
1166+
1167+
testBigQueryWrite.setCloseForeverAfter(2);
1168+
testBigQueryWrite.setTimesToClose(1);
1169+
testBigQueryWrite.addResponse(createAppendResponse(0));
1170+
testBigQueryWrite.addResponse(createAppendResponse(1));
1171+
1172+
// Connection will be failed after triggering the third append.
1173+
ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer1, new String[] {"A"}, 0);
1174+
ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer2, new String[] {"B"}, 1);
1175+
ApiFuture<AppendRowsResponse> appendFuture3 = sendTestMessage(writer3, new String[] {"C"}, 2);
1176+
TimeUnit.SECONDS.sleep(1);
1177+
assertEquals(0, appendFuture1.get().getAppendResult().getOffset().getValue());
1178+
assertEquals(1, appendFuture2.get().getAppendResult().getOffset().getValue());
1179+
assertThrows(
1180+
ExecutionException.class,
1181+
() -> {
1182+
assertEquals(2, appendFuture3.get().getAppendResult().getOffset().getValue());
1183+
});
1184+
assertEquals(writer1.getTestOnlyConnectionWorkerPool().getTotalConnectionCount(), 1);
1185+
assertEquals(writer1.getTestOnlyConnectionWorkerPool().getCreateConnectionCount(), 1);
1186+
1187+
writer1.close();
1188+
writer2.close();
1189+
// We will still be left with one request
1190+
assertEquals(writer1.getTestOnlyConnectionWorkerPool().getCreateConnectionCount(), 1);
1191+
}
1192+
1193+
public StreamWriter getMultiplexingStreamWriter(String streamName) throws IOException {
1194+
return StreamWriter.newBuilder(streamName, client)
1195+
.setWriterSchema(createProtoSchema())
1196+
.setEnableConnectionPool(true)
1197+
.setMaxInflightRequests(10)
1198+
.setLocation("US")
1199+
.setMaxRetryDuration(java.time.Duration.ofMillis(100))
1200+
.build();
1201+
}
1202+
10931203
// Timeout to ensure close() doesn't wait for done callback timeout.
10941204
@Test(timeout = 10000)
10951205
public void testCloseDisconnectedStream() throws Exception {

0 commit comments

Comments
 (0)