chore: generalise session pool class for multiplexed session. (#2964)

* chore: generalise session pool class for multiplexed session. * chore: add back previous code. * chore: address comments.
googleapis · Mar 25, 2024 · e043de8 · e043de8
1 parent d0ba290
commit e043de8
Show file tree

Hide file tree

Showing 3 changed files with 172 additions and 76 deletions.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java
@@ -257,7 +257,9 @@ private <T> T runWithSessionRetry(Function<Session, T> callable) {
       try {
         return callable.apply(session);
       } catch (SessionNotFoundException e) {
-        session = pool.replaceSession(e, session);
+        session =
+            (PooledSessionFuture)
+                pool.getPooledSessionReplacementHandler().replaceSession(e, session);
       }
     }
   }