fix(session): avoid cache exceptional session creation future (#189)

* fix(session): do not cache failed future * fix typo * fix the spotless
streamnative · Oct 11, 2024 · 439666f · 439666f
1 parent 36577ec
commit 439666f
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 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
@@ -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
@@ -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;