Skip to content

Commit 9ebdec6

Browse files
committed
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
1 parent 0b5ac92 commit 9ebdec6

File tree

2 files changed

+58
-7
lines changed

2 files changed

+58
-7
lines changed

core/src/main/java/com/google/adk/sessions/InMemorySessionService.java

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -192,13 +192,16 @@ public Completable deleteSession(String appName, String userId, String sessionId
192192
Objects.requireNonNull(userId, "userId cannot be null");
193193
Objects.requireNonNull(sessionId, "sessionId cannot be null");
194194

195-
ConcurrentMap<String, ConcurrentMap<String, Session>> appSessionsMap = sessions.get(appName);
196-
if (appSessionsMap != null) {
197-
ConcurrentMap<String, Session> userSessionsMap = appSessionsMap.get(userId);
198-
if (userSessionsMap != null) {
199-
userSessionsMap.remove(sessionId);
200-
}
201-
}
195+
sessions.computeIfPresent(appName, (app, appSessionsMap) -> {
196+
appSessionsMap.computeIfPresent(userId, (user, userSessionsMap) -> {
197+
userSessionsMap.remove(sessionId);
198+
// If userSessionsMap is now empty, return null to automatically remove the userId key
199+
return userSessionsMap.isEmpty() ? null : userSessionsMap;
200+
});
201+
// If appSessionsMap is now empty, return null to automatically remove the appName key
202+
return appSessionsMap.isEmpty() ? null : appSessionsMap;
203+
});
204+
202205
return Completable.complete();
203206
}
204207

core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import com.google.adk.events.Event;
2121
import com.google.adk.events.EventActions;
2222
import io.reactivex.rxjava3.core.Single;
23+
import java.lang.reflect.Field;
2324
import java.util.HashMap;
2425
import java.util.Optional;
2526
import java.util.concurrent.ConcurrentHashMap;
@@ -247,4 +248,51 @@ public void sequentialAgents_shareTempState() {
247248
assertThat(retrievedSession.state()).doesNotContainKey("temp:agent1_output");
248249
assertThat(retrievedSession.state()).containsEntry("temp:agent2_output", "processed_data");
249250
}
251+
252+
@Test
253+
public void deleteSession_cleansUpEmptyParentMaps() throws Exception {
254+
InMemorySessionService sessionService = new InMemorySessionService();
255+
256+
Session session = sessionService.createSession("app-name", "user-id").blockingGet();
257+
258+
sessionService.deleteSession(session.appName(), session.userId(), session.id()).blockingAwait();
259+
260+
// Use reflection to access the private 'sessions' field
261+
Field field = InMemorySessionService.class.getDeclaredField("sessions");
262+
field.setAccessible(true);
263+
ConcurrentMap<?, ?> sessions = (ConcurrentMap<?, ?>) field.get(sessionService);
264+
265+
// After deleting the only session for "user-id" under "app-name",
266+
// both the userId map and the appName map should have been removed
267+
assertThat(sessions).isEmpty();
268+
}
269+
270+
@Test
271+
public void deleteSession_doesNotRemoveUserMapWhenOtherSessionsExist() throws Exception {
272+
InMemorySessionService sessionService = new InMemorySessionService();
273+
274+
Session session1 = sessionService.createSession("app-name", "user-id").blockingGet();
275+
Session session2 = sessionService.createSession("app-name", "user-id").blockingGet();
276+
277+
// Delete only one of the two sessions
278+
sessionService.deleteSession(session1.appName(), session1.userId(), session1.id()).blockingAwait();
279+
280+
// session2 should still be retrievable
281+
assertThat(
282+
sessionService
283+
.getSession(session2.appName(), session2.userId(), session2.id(), Optional.empty())
284+
.blockingGet())
285+
.isNotNull();
286+
287+
// The userId entry should still exist (not pruned) because session2 remains
288+
Field field = InMemorySessionService.class.getDeclaredField("sessions");
289+
field.setAccessible(true);
290+
@SuppressWarnings("unchecked")
291+
ConcurrentMap<String, ConcurrentMap<String, ConcurrentMap<String, ?>>> sessions =
292+
(ConcurrentMap<String, ConcurrentMap<String, ConcurrentMap<String, ?>>>) field.get(sessionService);
293+
294+
assertThat(sessions.get("app-name")).isNotNull();
295+
assertThat(sessions.get("app-name").get("user-id")).isNotNull();
296+
assertThat(sessions.get("app-name").get("user-id")).hasSize(1);
297+
}
250298
}

0 commit comments

Comments
 (0)