@@ -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