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 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;