From 9ebdec69bc5cc78b14a3e0e2df17e9a86405932f Mon Sep 17 00:00:00 2001 From: YuqiGuo105 Date: Thu, 26 Feb 2026 23:06:38 -0700 Subject: [PATCH] fix: memory leak in InMemorySessionService.deleteSession() - clean up empty parent maps When deleteSession() removes the last session for a userId, the now-empty userId map remained in the parent appName map. Similarly, if all users under an appName were deleted, the empty appName map remained in the top-level sessions map. This caused linear memory growth with the number of unique users/apps, eventually leading to OutOfMemoryError. Replace the manual get+remove pattern with computeIfPresent() calls that atomically remove parent map entries when they become empty. When the remapping function returns null, ConcurrentHashMap automatically removes the key - this is thread-safe without additional locking. Before: sessions = { appName: { userId: {} } } <-- empty maps leak After: sessions = {} <-- fully cleaned up Add two new tests: - deleteSession_cleansUpEmptyParentMaps: verifies via reflection that the internal sessions map is completely empty after deleting the only session - deleteSession_doesNotRemoveUserMapWhenOtherSessionsExist: verifies that a userId entry is NOT pruned when the user still has remaining sessions Fixes #687 --- .../adk/sessions/InMemorySessionService.java | 17 ++++--- .../sessions/InMemorySessionServiceTest.java | 48 +++++++++++++++++++ 2 files changed, 58 insertions(+), 7 deletions(-) diff --git a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java index b2a584b11..345811633 100644 --- a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java +++ b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java @@ -192,13 +192,16 @@ public Completable deleteSession(String appName, String userId, String sessionId Objects.requireNonNull(userId, "userId cannot be null"); Objects.requireNonNull(sessionId, "sessionId cannot be null"); - ConcurrentMap> appSessionsMap = sessions.get(appName); - if (appSessionsMap != null) { - ConcurrentMap userSessionsMap = appSessionsMap.get(userId); - if (userSessionsMap != null) { - userSessionsMap.remove(sessionId); - } - } + sessions.computeIfPresent(appName, (app, appSessionsMap) -> { + appSessionsMap.computeIfPresent(userId, (user, userSessionsMap) -> { + userSessionsMap.remove(sessionId); + // If userSessionsMap is now empty, return null to automatically remove the userId key + return userSessionsMap.isEmpty() ? null : userSessionsMap; + }); + // If appSessionsMap is now empty, return null to automatically remove the appName key + return appSessionsMap.isEmpty() ? null : appSessionsMap; + }); + return Completable.complete(); } diff --git a/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java b/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java index 41e156ffd..53380d072 100644 --- a/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java +++ b/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java @@ -20,6 +20,7 @@ import com.google.adk.events.Event; import com.google.adk.events.EventActions; import io.reactivex.rxjava3.core.Single; +import java.lang.reflect.Field; import java.util.HashMap; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; @@ -247,4 +248,51 @@ public void sequentialAgents_shareTempState() { assertThat(retrievedSession.state()).doesNotContainKey("temp:agent1_output"); assertThat(retrievedSession.state()).containsEntry("temp:agent2_output", "processed_data"); } + + @Test + public void deleteSession_cleansUpEmptyParentMaps() throws Exception { + InMemorySessionService sessionService = new InMemorySessionService(); + + Session session = sessionService.createSession("app-name", "user-id").blockingGet(); + + sessionService.deleteSession(session.appName(), session.userId(), session.id()).blockingAwait(); + + // Use reflection to access the private 'sessions' field + Field field = InMemorySessionService.class.getDeclaredField("sessions"); + field.setAccessible(true); + ConcurrentMap sessions = (ConcurrentMap) field.get(sessionService); + + // After deleting the only session for "user-id" under "app-name", + // both the userId map and the appName map should have been removed + assertThat(sessions).isEmpty(); + } + + @Test + public void deleteSession_doesNotRemoveUserMapWhenOtherSessionsExist() throws Exception { + InMemorySessionService sessionService = new InMemorySessionService(); + + Session session1 = sessionService.createSession("app-name", "user-id").blockingGet(); + Session session2 = sessionService.createSession("app-name", "user-id").blockingGet(); + + // Delete only one of the two sessions + sessionService.deleteSession(session1.appName(), session1.userId(), session1.id()).blockingAwait(); + + // session2 should still be retrievable + assertThat( + sessionService + .getSession(session2.appName(), session2.userId(), session2.id(), Optional.empty()) + .blockingGet()) + .isNotNull(); + + // The userId entry should still exist (not pruned) because session2 remains + Field field = InMemorySessionService.class.getDeclaredField("sessions"); + field.setAccessible(true); + @SuppressWarnings("unchecked") + ConcurrentMap>> sessions = + (ConcurrentMap>>) field.get(sessionService); + + assertThat(sessions.get("app-name")).isNotNull(); + assertThat(sessions.get("app-name").get("user-id")).isNotNull(); + assertThat(sessions.get("app-name").get("user-id")).hasSize(1); + } }