Fix performance issues with many offline sessions

Fixes: #13340
keycloak · Nov 11, 2022 · 0b7e317 · 0b7e317
1 parent 02a6956
commit 0b7e317
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 12 deletions.
diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/session/JpaUserSessionPersisterProvider.java b/model/jpa/src/main/java/org/keycloak/models/jpa/session/JpaUserSessionPersisterProvider.java
@@ -352,7 +352,7 @@ public Stream<UserSessionModel> loadUserSessionsStream(RealmModel realm, ClientM
         query.setParameter("offline", offlineStr);
         query.setParameter("realmId", realm.getId());
 
-        return loadUserSessionsWithClientSessions(query, offlineStr);
+        return loadUserSessionsWithClientSessions(query, offlineStr, true);
     }
 
     @Override
@@ -368,7 +368,7 @@ public Stream<UserSessionModel> loadUserSessionsStream(RealmModel realm, UserMod
         query.setParameter("realmId", realm.getId());
         query.setParameter("userId", user.getId());
 
-        return loadUserSessionsWithClientSessions(query, offlineStr);
+        return loadUserSessionsWithClientSessions(query, offlineStr, true);
     }
 
     public Stream<UserSessionModel> loadUserSessionsStream(Integer firstResult, Integer maxResults, boolean offline,
@@ -379,10 +379,19 @@ public Stream<UserSessionModel> loadUserSessionsStream(Integer firstResult, Inte
             .setParameter("offline", offlineStr)
             .setParameter("lastSessionId", lastUserSessionId), firstResult, maxResults);
 
-        return loadUserSessionsWithClientSessions(query, offlineStr);
+        return loadUserSessionsWithClientSessions(query, offlineStr, false);
     }
 
-    private Stream<UserSessionModel> loadUserSessionsWithClientSessions(TypedQuery<PersistentUserSessionEntity> query, String offlineStr) {
+    /**
+     *
+     * @param query
+     * @param offlineStr
+     * @param useExact If {@code true}, then only client sessions from the user sessions
+     *          obtained from the {@code query} are loaded. If {@code false}, then user session
+     *          is taken as range, and all client sessions in between the minimum and maximum session
+     * @return
+     */
+    private Stream<UserSessionModel> loadUserSessionsWithClientSessions(TypedQuery<PersistentUserSessionEntity> query, String offlineStr, boolean useExact) {
 
         List<PersistentUserSessionAdapter> userSessionAdapters = closing(query.getResultStream()
                 .map(this::toAdapter)
@@ -392,16 +401,25 @@ private Stream<UserSessionModel> loadUserSessionsWithClientSessions(TypedQuery<P
         Map<String, PersistentUserSessionAdapter> sessionsById = userSessionAdapters.stream()
                 .collect(Collectors.toMap(UserSessionModel::getId, Function.identity()));
 
+        Set<String> userSessionIds = sessionsById.keySet();
+
         Set<String> removedClientUUIDs = new HashSet<>();
 
         if (!sessionsById.isEmpty()) {
-            String fromUserSessionId = userSessionAdapters.get(0).getId();
-            String toUserSessionId = userSessionAdapters.get(userSessionAdapters.size() - 1).getId();
-
-            TypedQuery<PersistentClientSessionEntity> queryClientSessions = em.createNamedQuery("findClientSessionsOrderedById", PersistentClientSessionEntity.class);
-            queryClientSessions.setParameter("offline", offlineStr);
-            queryClientSessions.setParameter("fromSessionId", fromUserSessionId);
-            queryClientSessions.setParameter("toSessionId", toUserSessionId);
+            TypedQuery<PersistentClientSessionEntity> queryClientSessions;
+            if (useExact) {
+                queryClientSessions = em.createNamedQuery("findClientSessionsOrderedByIdExact", PersistentClientSessionEntity.class);
+                queryClientSessions.setParameter("offline", offlineStr);
+                queryClientSessions.setParameter("userSessionIds", userSessionIds);
+            } else {
+                String fromUserSessionId = userSessionAdapters.get(0).getId();
+                String toUserSessionId = userSessionAdapters.get(userSessionAdapters.size() - 1).getId();
+
+                queryClientSessions = em.createNamedQuery("findClientSessionsOrderedByIdInterval", PersistentClientSessionEntity.class);
+                queryClientSessions.setParameter("offline", offlineStr);
+                queryClientSessions.setParameter("fromSessionId", fromUserSessionId);
+                queryClientSessions.setParameter("toSessionId", toUserSessionId);
+            }
 
             closing(queryClientSessions.getResultStream()).forEach(clientSession -> {
                 PersistentUserSessionAdapter userSession = sessionsById.get(clientSession.getUserSessionId());

diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/session/PersistentClientSessionEntity.java b/model/jpa/src/main/java/org/keycloak/models/jpa/session/PersistentClientSessionEntity.java
@@ -44,7 +44,8 @@
         //                 It is removed from here and added manually in JpaUtils to give a native implementation if needed
         //@NamedQuery(name="deleteExpiredClientSessions", query="delete from PersistentClientSessionEntity sess where sess.userSessionId IN (select u.userSessionId from PersistentUserSessionEntity u where u.realmId = :realmId AND u.offline = :offline AND u.lastSessionRefresh < :lastSessionRefresh)"),
         @NamedQuery(name="findClientSessionsByUserSession", query="select sess from PersistentClientSessionEntity sess where sess.userSessionId=:userSessionId and sess.offline = :offline"),
-        @NamedQuery(name="findClientSessionsOrderedById", query="select sess from PersistentClientSessionEntity sess where sess.offline = :offline and sess.userSessionId >= :fromSessionId and sess.userSessionId <= :toSessionId order by sess.userSessionId"),
+        @NamedQuery(name="findClientSessionsOrderedByIdInterval", query="select sess from PersistentClientSessionEntity sess where sess.offline = :offline and sess.userSessionId >= :fromSessionId and sess.userSessionId <= :toSessionId order by sess.userSessionId"),
+        @NamedQuery(name="findClientSessionsOrderedByIdExact", query="select sess from PersistentClientSessionEntity sess where sess.offline = :offline and sess.userSessionId IN (:userSessionIds)"),
         @NamedQuery(name="findClientSessionsCountByClient", query="select count(sess) from PersistentClientSessionEntity sess where sess.offline = :offline and sess.clientId = :clientId"),
         @NamedQuery(name="findClientSessionsCountByExternalClient", query="select count(sess) from PersistentClientSessionEntity sess where sess.offline = :offline and sess.clientStorageProvider = :clientStorageProvider and sess.externalClientId = :externalClientId")
 })