From 439666fc04bf06216627e03ac55f91080aaa8286 Mon Sep 17 00:00:00 2001
From: Qiang Zhao <mattisonchao@apache.org>
Date: Sat, 12 Oct 2024 06:37:11 +0800
Subject: [PATCH] fix(session): avoid cache exceptional session creation future
 (#189)

* fix(session): do not cache failed future

* fix typo

* fix the spotless
---
 .../oxia/client/session/SessionManager.java   |  9 +++++++-
 .../client/session/SessionManagerTest.java    | 23 +++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/client/src/main/java/io/streamnative/oxia/client/session/SessionManager.java b/client/src/main/java/io/streamnative/oxia/client/session/SessionManager.java
index a9f56139..644d8637 100644
--- a/client/src/main/java/io/streamnative/oxia/client/session/SessionManager.java
+++ b/client/src/main/java/io/streamnative/oxia/client/session/SessionManager.java
@@ -58,7 +58,14 @@ public CompletableFuture<Session> getSession(long shardId) {
                     new IllegalStateException("session manager has been closed"));
         }
 
-        return sessionsByShardId.computeIfAbsent(shardId, s -> factory.create(shardId));
+        return sessionsByShardId.compute(
+                shardId,
+                (key, existing) -> {
+                    if (existing != null && !existing.isCompletedExceptionally()) {
+                        return existing;
+                    }
+                    return factory.create(shardId);
+                });
     }
 
     @Override
diff --git a/client/src/test/java/io/streamnative/oxia/client/session/SessionManagerTest.java b/client/src/test/java/io/streamnative/oxia/client/session/SessionManagerTest.java
index ce549426..25db3d2d 100644
--- a/client/src/test/java/io/streamnative/oxia/client/session/SessionManagerTest.java
+++ b/client/src/test/java/io/streamnative/oxia/client/session/SessionManagerTest.java
@@ -79,6 +79,29 @@ void existingSession() {
         verifyNoMoreInteractions(factory, session);
     }
 
+    @Test
+    void existingSessionWithFailure() {
+        var shardId = 1L;
+        // first failed
+        when(factory.create(shardId))
+                .thenReturn(CompletableFuture.failedFuture(new IllegalStateException("failed")));
+        var session1 = manager.getSession(shardId);
+        assertThat(session1).isCompletedExceptionally();
+        verify(factory, times(1)).create(shardId);
+
+        // second should be success
+        when(factory.create(shardId)).thenReturn(CompletableFuture.completedFuture(session));
+        var session2 = manager.getSession(shardId);
+        assertThat(session2).isCompletedWithValue(session);
+        verify(factory, times(2)).create(shardId);
+
+        // third should be cache
+        var session3 = manager.getSession(shardId);
+        assertThat(session3).isSameAs(session2);
+        verify(factory, times(2)).create(shardId);
+        verifyNoMoreInteractions(factory, session);
+    }
+
     @Test
     void close() throws Exception {
         var shardId = 5L;