From 9ebdec69bc5cc78b14a3e0e2df17e9a86405932f Mon Sep 17 00:00:00 2001 From: YuqiGuo105 Date: Thu, 26 Feb 2026 23:06:38 -0700 Subject: [PATCH 1/3] 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); + } } From d02d25c93dccc359c2ebfa2ecf9651623c28f073 Mon Sep 17 00:00:00 2001 From: yongkiy-google Date: Wed, 25 Mar 2026 15:46:03 +0100 Subject: [PATCH 2/3] Update InMemorySessionService.java --- .../adk/sessions/InMemorySessionService.java | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 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 3519fe0f6..3e96f959e 100644 --- a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java +++ b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java @@ -186,16 +186,21 @@ public Completable deleteSession(String appName, String userId, String sessionId Objects.requireNonNull(userId, "userId cannot be null"); Objects.requireNonNull(sessionId, "sessionId cannot be null"); - 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; + 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; }); - // If appSessionsMap is now empty, return null to automatically remove the appName key - return appSessionsMap.isEmpty() ? null : appSessionsMap; - }); - + return Completable.complete(); } From 316702ef323dc924a54b3168de3d4f7fc3b6d31d Mon Sep 17 00:00:00 2001 From: yongkiy-google Date: Wed, 25 Mar 2026 15:48:02 +0100 Subject: [PATCH 3/3] Update InMemorySessionServiceTest.java --- .../google/adk/sessions/InMemorySessionServiceTest.java | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) 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 aa763b691..c260f6695 100644 --- a/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java +++ b/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java @@ -274,7 +274,9 @@ public void deleteSession_cleansUpEmptyParentMaps() throws Exception { Session session = sessionService.createSession("app-name", "user-id").blockingGet(); - sessionService.deleteSession(session.appName(), session.userId(), session.id()).blockingAwait(); + sessionService + .deleteSession(session1.appName(), session1.userId(), session1.id()) + .blockingAwait(); // Use reflection to access the private 'sessions' field Field field = InMemorySessionService.class.getDeclaredField("sessions"); @@ -308,7 +310,8 @@ public void deleteSession_doesNotRemoveUserMapWhenOtherSessionsExist() throws Ex field.setAccessible(true); @SuppressWarnings("unchecked") ConcurrentMap>> sessions = - (ConcurrentMap>>) field.get(sessionService); + (ConcurrentMap>>) + field.get(sessionService); assertThat(sessions.get("app-name")).isNotNull(); assertThat(sessions.get("app-name").get("user-id")).isNotNull();