Visitar URL original
deflake location-aware tests · googleapis/google-cloud-java@c47d2a9 · GitHub
Skip to content

Commit c47d2a9

Browse files
rahul2393rahul2393
committed
deflake location-aware tests
1 parent 95e6ca5 commit c47d2a9

1 file changed

Lines changed: 47 additions & 22 deletions

File tree

‎java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/LocationAwareSharedBackendReplicaHarnessTest.java‎

Lines changed: 47 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ public class LocationAwareSharedBackendReplicaHarnessTest {
6767
private static final String PROJECT = "fake-project";
6868
private static final String INSTANCE = "fake-instance";
6969
private static final String DATABASE = "fake-database";
70+
private static final AtomicInteger DATABASE_COUNTER = new AtomicInteger();
7071
private static final String TABLE = "T";
7172
private static final String REPLICA_LOCATION = "us-east1";
7273
private static final Statement SEED_QUERY = Statement.of("SELECT 1");
@@ -85,6 +86,14 @@ public class LocationAwareSharedBackendReplicaHarnessTest {
8586
.build();
8687
private static SharedBackendReplicaHarness harness;
8788

89+
// Assigned a fresh value per test so each test uses a distinct database scope. The location-aware
90+
// routing layer keeps a process-wide EndpointLatencyRegistry keyed by (databaseScope,
91+
// operationUid, ...). Each test builds a new Spanner whose per-channel operationUid counter
92+
// restarts at 1, so without a unique scope a prior test's recorded error/latency penalties (e.g.
93+
// an aborted commit on the leader) would leak into a later test and skew its routing, making the
94+
// routing assertions flaky.
95+
private String database;
96+
8897
@BeforeClass
8998
public static void enableLocationAwareRouting() throws Exception {
9099
SpannerOptions.useEnvironment(
@@ -100,6 +109,7 @@ public boolean isEnableLocationApi() {
100109
@Before
101110
public void resetHarness() {
102111
harness.reset();
112+
database = DATABASE + "-" + DATABASE_COUNTER.incrementAndGet();
103113
}
104114

105115
@AfterClass
@@ -118,7 +128,7 @@ public static void restoreEnvironment() throws Exception {
118128
public void singleUseReadReroutesOnResourceExhaustedForBypassTraffic() throws Exception {
119129
try (Spanner spanner = createSpanner(harness)) {
120130
configureBackend(harness, singleRowReadResultSet("b"));
121-
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
131+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
122132

123133
seedLocationMetadata(client);
124134
waitForReplicaRoutedRead(client, harness);
@@ -135,7 +145,7 @@ public void singleUseReadReroutesOnResourceExhaustedForBypassTraffic() throws Ex
135145
KeySet.singleKey(Key.of("b")),
136146
Arrays.asList("k"),
137147
Options.directedRead(DIRECTED_READ_OPTIONS))) {
138-
assertTrue(resultSet.next());
148+
assertTrue(drain(resultSet));
139149
}
140150

141151
String diagnostics = routingDiagnostics(harness);
@@ -154,7 +164,7 @@ public void singleUseReadReroutesOnResourceExhaustedForBypassTraffic() throws Ex
154164
public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() throws Exception {
155165
try (Spanner spanner = createSpanner(harness)) {
156166
configureBackend(harness, singleRowReadResultSet("b"));
157-
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
167+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
158168

159169
seedLocationMetadata(client);
160170
waitForReplicaRoutedRead(client, harness);
@@ -172,7 +182,7 @@ public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() thr
172182
KeySet.singleKey(Key.of("b")),
173183
Arrays.asList("k"),
174184
Options.directedRead(DIRECTED_READ_OPTIONS))) {
175-
assertTrue(firstRead.next());
185+
assertTrue(drain(firstRead));
176186
}
177187

178188
try (ResultSet secondRead =
@@ -183,7 +193,7 @@ public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() thr
183193
KeySet.singleKey(Key.of("b")),
184194
Arrays.asList("k"),
185195
Options.directedRead(DIRECTED_READ_OPTIONS))) {
186-
assertTrue(secondRead.next());
196+
assertTrue(drain(secondRead));
187197
}
188198

189199
String diagnostics = routingDiagnostics(harness);
@@ -210,7 +220,7 @@ public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() thr
210220
public void singleUseReadReroutesOnUnavailableForBypassTraffic() throws Exception {
211221
try (Spanner spanner = createSpanner(harness)) {
212222
configureBackend(harness, singleRowReadResultSet("b"));
213-
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
223+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
214224

215225
seedLocationMetadata(client);
216226
waitForReplicaRoutedRead(client, harness);
@@ -227,7 +237,7 @@ public void singleUseReadReroutesOnUnavailableForBypassTraffic() throws Exceptio
227237
KeySet.singleKey(Key.of("b")),
228238
Arrays.asList("k"),
229239
Options.directedRead(DIRECTED_READ_OPTIONS))) {
230-
assertTrue(resultSet.next());
240+
assertTrue(drain(resultSet));
231241
}
232242

233243
String diagnostics = routingDiagnostics(harness);
@@ -247,7 +257,7 @@ public void singleUseReadCooldownSkipsUnavailableReplicaOnNextRequestForBypassTr
247257
throws Exception {
248258
try (Spanner spanner = createSpanner(harness)) {
249259
configureBackend(harness, singleRowReadResultSet("b"));
250-
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
260+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
251261

252262
seedLocationMetadata(client);
253263
waitForReplicaRoutedRead(client, harness);
@@ -264,7 +274,7 @@ public void singleUseReadCooldownSkipsUnavailableReplicaOnNextRequestForBypassTr
264274
KeySet.singleKey(Key.of("b")),
265275
Arrays.asList("k"),
266276
Options.directedRead(DIRECTED_READ_OPTIONS))) {
267-
assertTrue(firstRead.next());
277+
assertTrue(drain(firstRead));
268278
}
269279

270280
try (ResultSet secondRead =
@@ -275,7 +285,7 @@ public void singleUseReadCooldownSkipsUnavailableReplicaOnNextRequestForBypassTr
275285
KeySet.singleKey(Key.of("b")),
276286
Arrays.asList("k"),
277287
Options.directedRead(DIRECTED_READ_OPTIONS))) {
278-
assertTrue(secondRead.next());
288+
assertTrue(drain(secondRead));
279289
}
280290

281291
String diagnostics = routingDiagnostics(harness);
@@ -303,7 +313,7 @@ public void singleUseReadMidStreamRecvFailureWithoutRetryInfoRetriesForBypassTra
303313
throws Exception {
304314
try (Spanner spanner = createSpanner(harness)) {
305315
configureBackend(harness, multiRowReadResultSet("b", "c", "d"));
306-
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
316+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
307317

308318
seedLocationMetadata(client);
309319
waitForReplicaRoutedRead(client, harness);
@@ -359,7 +369,7 @@ public void readWriteTransactionInlineBeginOnDefaultKeepsReadOnDefaultForBypassT
359369
throws Exception {
360370
try (Spanner spanner = createSpanner(harness)) {
361371
configureBackend(harness, singleRowReadResultSet("b"), /* leaderReplicaIndex= */ 1);
362-
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
372+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
363373

364374
seedLocationMetadata(client);
365375
waitForReplicaRoutedStrongRead(client, harness, /* expectedReplicaIndex= */ 1);
@@ -370,11 +380,11 @@ public void readWriteTransactionInlineBeginOnDefaultKeepsReadOnDefaultForBypassT
370380
.run(
371381
transaction -> {
372382
try (ResultSet resultSet = transaction.executeQuery(SEED_QUERY)) {
373-
assertTrue(resultSet.next());
383+
assertTrue(drain(resultSet));
374384
}
375385
try (ResultSet resultSet =
376386
transaction.read(TABLE, KeySet.singleKey(Key.of("b")), Arrays.asList("k"))) {
377-
assertTrue(resultSet.next());
387+
assertTrue(drain(resultSet));
378388
}
379389
return null;
380390
});
@@ -422,7 +432,7 @@ public void readWriteTransactionAbortedCommitUsesReadAffinityReplicaForBypassTra
422432
throws Exception {
423433
try (Spanner spanner = createSpanner(harness)) {
424434
configureBackend(harness, singleRowReadResultSet("b"), /* leaderReplicaIndex= */ 1);
425-
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
435+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
426436

427437
seedLocationMetadata(client);
428438
waitForReplicaRoutedStrongRead(client, harness, /* expectedReplicaIndex= */ 1);
@@ -437,7 +447,7 @@ public void readWriteTransactionAbortedCommitUsesReadAffinityReplicaForBypassTra
437447
int attempt = attempts.incrementAndGet();
438448
try (ResultSet resultSet =
439449
transaction.read(TABLE, KeySet.singleKey(Key.of("b")), Arrays.asList("k"))) {
440-
assertTrue(resultSet.next());
450+
assertTrue(drain(resultSet));
441451
}
442452

443453
if (attempt == 1) {
@@ -549,6 +559,20 @@ private static void seedLocationMetadata(DatabaseClient client) {
549559
}
550560
}
551561

562+
/**
563+
* Fully consumes the result set so the underlying gRPC streaming read closes promptly. Leaving a
564+
* stream open keeps the routed endpoint's active-request count above zero, which inflates its
565+
* selection cost and can bounce the next request to a different replica, making routing
566+
* assertions flaky under load. Returns whether at least one row was seen.
567+
*/
568+
private static boolean drain(ResultSet resultSet) {
569+
boolean sawRow = false;
570+
while (resultSet.next()) {
571+
sawRow = true;
572+
}
573+
return sawRow;
574+
}
575+
552576
private static int waitForReplicaRoutedRead(
553577
DatabaseClient client, SharedBackendReplicaHarness harness) throws InterruptedException {
554578
long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(10);
@@ -561,7 +585,7 @@ private static int waitForReplicaRoutedRead(
561585
KeySet.singleKey(Key.of("b")),
562586
Arrays.asList("k"),
563587
Options.directedRead(DIRECTED_READ_OPTIONS))) {
564-
if (resultSet.next()) {
588+
if (drain(resultSet)) {
565589
for (int replicaIndex = 0; replicaIndex < harness.replicas.size(); replicaIndex++) {
566590
if (!harness
567591
.replicas
@@ -584,17 +608,18 @@ private static void waitForReplicaRoutedStrongRead(
584608
long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(10);
585609
while (System.nanoTime() < deadlineNanos) {
586610
harness.clearRequests();
611+
boolean sawRow;
587612
try (ResultSet resultSet =
588613
client.singleUse().read(TABLE, KeySet.singleKey(Key.of("b")), Arrays.asList("k"))) {
589-
if (resultSet.next()) {
590-
if (!harness
614+
sawRow = drain(resultSet);
615+
}
616+
if (sawRow
617+
&& !harness
591618
.replicas
592619
.get(expectedReplicaIndex)
593620
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
594621
.isEmpty()) {
595-
return;
596-
}
597-
}
622+
return;
598623
}
599624
Thread.sleep(50L);
600625
}

0 commit comments

Comments
 (0)