From 43acf3075fd019f4771101a162c8e029220bb331 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Wed, 7 Aug 2024 14:50:14 +0200
Subject: [PATCH] chore: unflake
 blockAndTimeoutOnPoolExhaustion_withAcquireSessionTimeout

Fixes 3251
---
 .../google/cloud/spanner/SessionPoolTest.java   | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java
index 58b0280cf08..f4430861194 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java
@@ -1294,7 +1294,7 @@ public void blockAndTimeoutOnPoolExhaustion_withAcquireSessionTimeout() throws E
             .setMinSessions(minSessions)
             .setMaxSessions(1)
             .setInitialWaitForSessionTimeoutMillis(20L)
-            .setAcquireSessionTimeout(Duration.ofMillis(20L))
+            .setAcquireSessionTimeout(null)
             .build();
     setupMockSessionCreation();
     pool = createPool();
@@ -1307,18 +1307,18 @@ public void blockAndTimeoutOnPoolExhaustion_withAcquireSessionTimeout() throws E
     Future<Void> fut =
         executor.submit(
             () -> {
-              latch.countDown();
               PooledSessionFuture session = pool.getSession();
+              latch.countDown();
+              session.get();
               session.close();
               return null;
             });
     // Wait until the background thread is actually waiting for a session.
     latch.await();
     // Wait until the request has timed out.
-    int waitCount = 0;
-    while (pool.getNumWaiterTimeouts() == 0L && waitCount < 5000) {
-      Thread.sleep(1L);
-      waitCount++;
+    Stopwatch watch = Stopwatch.createStarted();
+    while (pool.getNumWaiterTimeouts() == 0L && watch.elapsed(TimeUnit.MILLISECONDS) < 1000) {
+      Thread.yield();
     }
     // Return the checked out session to the pool so the async request will get a session and
     // finish.
@@ -1326,10 +1326,11 @@ public void blockAndTimeoutOnPoolExhaustion_withAcquireSessionTimeout() throws E
     // Verify that the async request also succeeds.
     fut.get(10L, TimeUnit.SECONDS);
     executor.shutdown();
+    assertTrue(executor.awaitTermination(10L, TimeUnit.SECONDS));
 
     // Verify that the session was returned to the pool and that we can get it again.
-    Session session = pool.getSession();
-    assertThat(session).isNotNull();
+    PooledSessionFuture session = pool.getSession();
+    assertThat(session.get()).isNotNull();
     session.close();
     assertThat(pool.getNumWaiterTimeouts()).isAtLeast(1L);
   }