Skip to content

Commit dd8c4d1

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 dd8c4d1

2 files changed

Lines changed: 66 additions & 7 deletions

File tree

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: 56 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,59 @@ 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+
// Fixes https://github.com/google/adk-java/issues/687
253+
@Test
254+
public void deleteSession_cleansUpEmptyParentMaps() throws Exception {
255+
InMemorySessionService sessionService = new InMemorySessionService();
256+
Session session = sessionService.createSession("app-name", "user-id").blockingGet();
257+
258+
sessionService.deleteSession(session.appName(), session.userId(), session.id()).blockingAwait();
259+
260+
assertThat(getSessionsMap(sessionService)).isEmpty();
261+
}
262+
263+
@Test
264+
public void deleteSession_doesNotRemoveUserMapWhenOtherSessionsExist() throws Exception {
265+
InMemorySessionService sessionService = new InMemorySessionService();
266+
Session session1 = sessionService.createSession("app-name", "user-id").blockingGet();
267+
Session session2 = sessionService.createSession("app-name", "user-id").blockingGet();
268+
269+
sessionService.deleteSession(session1.appName(), session1.userId(), session1.id()).blockingAwait();
270+
271+
// session2 still retrievable
272+
assertThat(
273+
sessionService
274+
.getSession(session2.appName(), session2.userId(), session2.id(), Optional.empty())
275+
.blockingGet())
276+
.isNotNull();
277+
278+
// userId entry survives because session2 remains
279+
assertThat(getSessionsMap(sessionService).get("app-name").get("user-id")).hasSize(1);
280+
}
281+
282+
// Reproduces issue #687: user1/user2/user3 each delete their only session → no ghost maps left.
283+
@Test
284+
public void deleteSession_multipleUsers_cleansUpAllEmptyMaps() throws Exception {
285+
InMemorySessionService sessionService = new InMemorySessionService();
286+
Session s1 = sessionService.createSession("myApp", "user1").blockingGet();
287+
Session s2 = sessionService.createSession("myApp", "user2").blockingGet();
288+
Session s3 = sessionService.createSession("myApp", "user3").blockingGet();
289+
290+
sessionService.deleteSession(s1.appName(), s1.userId(), s1.id()).blockingAwait();
291+
sessionService.deleteSession(s2.appName(), s2.userId(), s2.id()).blockingAwait();
292+
sessionService.deleteSession(s3.appName(), s3.userId(), s3.id()).blockingAwait();
293+
294+
assertThat(getSessionsMap(sessionService)).isEmpty();
295+
}
296+
297+
// Helper: accesses the private sessions field via reflection.
298+
@SuppressWarnings("unchecked")
299+
private static ConcurrentMap<String, ConcurrentMap<String, ConcurrentMap<String, ?>>>
300+
getSessionsMap(InMemorySessionService service) throws Exception {
301+
Field field = InMemorySessionService.class.getDeclaredField("sessions");
302+
field.setAccessible(true);
303+
return (ConcurrentMap<String, ConcurrentMap<String, ConcurrentMap<String, ?>>>)
304+
field.get(service);
305+
}
250306
}

0 commit comments

Comments
 (0)