KEYCLOAK-11019 Initial support for lazy offline user-session loading

Co-authored-by: Thomas Darimont <thomas.darimont@googlemail.com>
Co-authored-by: Thomas Darimont <thomas.darimont@gmail.com>

model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProvider.java

@@ -30,6 +30,7 @@ import org.keycloak.device.DeviceActivityManager;
 import org.keycloak.models.AuthenticatedClientSessionModel;
 import org.keycloak.models.ClientModel;
 import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.ModelException;
 import org.keycloak.models.OfflineUserSessionModel;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserModel;
@@ -60,9 +61,11 @@ import org.keycloak.models.sessions.infinispan.util.SessionTimeouts;
 
 import java.io.Serializable;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
@@ -105,6 +108,8 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
     protected final RemoteCacheInvoker remoteCacheInvoker;
     protected final InfinispanKeyGenerator keyGenerator;
 
+    protected final boolean loadOfflineSessionsStatsFromDatabase;
+
     public InfinispanUserSessionProvider(KeycloakSession session,
                                          RemoteCacheInvoker remoteCacheInvoker,
                                          CrossDCLastSessionRefreshStore lastSessionRefreshStore,
@@ -114,7 +119,8 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
                                          Cache<String, SessionEntityWrapper<UserSessionEntity>> sessionCache,
                                          Cache<String, SessionEntityWrapper<UserSessionEntity>> offlineSessionCache,
                                          Cache<UUID, SessionEntityWrapper<AuthenticatedClientSessionEntity>> clientSessionCache,
-                                         Cache<UUID, SessionEntityWrapper<AuthenticatedClientSessionEntity>> offlineClientSessionCache) {
+                                         Cache<UUID, SessionEntityWrapper<AuthenticatedClientSessionEntity>> offlineClientSessionCache,
+                                         boolean loadOfflineSessionsStatsFromDatabase) {
         this.session = session;
 
         this.sessionCache = sessionCache;
@@ -134,6 +140,7 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
         this.persisterLastSessionRefreshStore = persisterLastSessionRefreshStore;
         this.remoteCacheInvoker = remoteCacheInvoker;
         this.keyGenerator = keyGenerator;
+        this.loadOfflineSessionsStatsFromDatabase = loadOfflineSessionsStatsFromDatabase;
 
         session.getTransactionManager().enlistAfterCompletion(clusterEventsSenderTx);
         session.getTransactionManager().enlistAfterCompletion(sessionTx);
@@ -239,19 +246,78 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
 
         entity.setStarted(currentTime);
         entity.setLastSessionRefresh(currentTime);
-
-
     }
 
-
     @Override
     public UserSessionModel getUserSession(RealmModel realm, String id) {
         return getUserSession(realm, id, false);
     }
 
     protected UserSessionAdapter getUserSession(RealmModel realm, String id, boolean offline) {
-        UserSessionEntity entity = getUserSessionEntity(realm, id, offline);
-        return wrap(realm, entity, offline);
+
+        UserSessionEntity userSessionEntityFromCache = getUserSessionEntity(realm, id, offline);
+        if (userSessionEntityFromCache != null) {
+            return wrap(realm, userSessionEntityFromCache, offline);
+        }
+
+        if (!offline) {
+            return null;
+        }
+
+        // Try to recover from potentially lost offline-sessions by attempting to fetch and re-import
+        // the offline session information from the PersistenceProvider.
+        UserSessionEntity userSessionEntityFromPersistenceProvider = getUserSessionEntityFromPersistenceProvider(realm, id, offline);
+        if (userSessionEntityFromPersistenceProvider != null) {
+            // we successfully recovered the offline session!
+            return wrap(realm, userSessionEntityFromPersistenceProvider, offline);
+        }
+
+        // no luck, the session is really not there anymore
+        return null;
+    }
+
+    private UserSessionEntity getUserSessionEntityFromPersistenceProvider(RealmModel realm, String sessionId, boolean offline) {
+
+        log.debugf("Offline user-session not found in infinispan, attempting UserSessionPersisterProvider lookup for sessionId=%s", sessionId);
+        UserSessionPersisterProvider persister = session.getProvider(UserSessionPersisterProvider.class);
+        UserSessionModel persistentUserSession = persister.loadUserSession(realm, sessionId, offline);
+
+        if (persistentUserSession == null) {
+            log.debugf("Offline user-session not found in UserSessionPersisterProvider for sessionId=%s", sessionId);
+            return null;
+        }
+
+        return importUserSession(realm, offline, persistentUserSession);
+    }
+
+    private UserSessionEntity getUserSessionEntityFromCacheOrImportIfNecessary(RealmModel realm, boolean offline, UserSessionModel persistentUserSession) {
+
+        UserSessionEntity userSessionEntity = getUserSessionEntity(realm, persistentUserSession.getId(), offline);
+        if (userSessionEntity != null) {
+            // user session present in cache, return existing session
+            return userSessionEntity;
+        }
+
+        return importUserSession(realm, offline, persistentUserSession);
+    }
+
+    private UserSessionEntity importUserSession(RealmModel realm, boolean offline, UserSessionModel persistentUserSession) {
+
+        String sessionId = persistentUserSession.getId();
+
+        log.debugf("Attempting to import user-session for sessionId=%s offline=%s", sessionId, offline);
+        session.sessions().importUserSessions(Collections.singleton(persistentUserSession), offline);
+        log.debugf("user-session imported, trying another lookup for sessionId=%s offline=%s", sessionId, offline);
+
+        UserSessionEntity ispnUserSessionEntity = getUserSessionEntity(realm, sessionId, offline);
+
+        if (ispnUserSessionEntity != null) {
+            log.debugf("user-session found after import for sessionId=%s offline=%s", sessionId, offline);
+            return ispnUserSessionEntity;
+        }
+
+        log.debugf("user-session could not be found after import for sessionId=%s offline=%s", sessionId, offline);
+        return null;
     }
 
     private UserSessionEntity getUserSessionEntity(RealmModel realm, String id, boolean offline) {
@@ -263,8 +329,41 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
         return entity;
     }
 
+    private Stream<UserSessionModel> getUserSessionsFromPersistenceProviderStream(RealmModel realm, UserModel user, boolean offline) {
+        UserSessionPersisterProvider persister = session.getProvider(UserSessionPersisterProvider.class);
+        return persister.loadUserSessionsStream(realm, user, offline, 0, null)
+                .map(persistentUserSession -> getUserSessionEntityFromCacheOrImportIfNecessary(realm, offline, persistentUserSession))
+                .filter(Objects::nonNull)
+                .map(userSessionEntity -> wrap(realm, userSessionEntity, offline));
+    }
+
+
+    protected Stream<UserSessionModel> getUserSessionsStream(RealmModel realm, UserSessionPredicate predicate, boolean offline) {
+
+        if (offline && loadOfflineSessionsStatsFromDatabase) {
+
+            // fetch the offline user-sessions from the persistence provider
+            UserSessionPersisterProvider persister = session.getProvider(UserSessionPersisterProvider.class);
+
+            UserModel user = session.users().getUserById(realm, predicate.getUserId());
+            if (user != null) {
+                return persister.loadUserSessionsStream(realm, user, offline, 0, null);
+            }
+
+            if (predicate.getBrokerSessionId() != null) {
+                // TODO add support for offline user-session lookup by brokerSessionId
+                // currently it is not possible to access the brokerSessionId in offline user-session in a database agnostic way
+                throw new ModelException("Dynamic database lookup for offline user-sessions by brokerSessionId is currently only supported for preloaded sessions.");
+            }
+
+            if (predicate.getBrokerUserId() != null) {
+                // TODO add support for offline user-session lookup by brokerUserId
+                // currently it is not possible to access the brokerUserId in offline user-session in a database agnostic way
+                throw new ModelException("Dynamic database lookup for offline user-sessions by brokerUserId is currently only supported for preloaded sessions.");
+            }
+
+        }
 
-    protected Stream<UserSessionModel> getUserSessionsStream(RealmModel realm, Predicate<Map.Entry<String, SessionEntityWrapper<UserSessionEntity>>> predicate, boolean offline) {
         Cache<String, SessionEntityWrapper<UserSessionEntity>> cache = getCache(offline);
         cache = CacheDecorators.skipCacheLoaders(cache);
 
@@ -321,6 +420,13 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
     }
 
     protected Stream<UserSessionModel> getUserSessionsStream(final RealmModel realm, ClientModel client, Integer firstResult, Integer maxResults, final boolean offline) {
+
+        if (offline && loadOfflineSessionsStatsFromDatabase) {
+            // fetch the actual offline user session count from the database
+            UserSessionPersisterProvider persister = session.getProvider(UserSessionPersisterProvider.class);
+            return persister.loadUserSessionsStream(realm, client, offline, firstResult, maxResults);
+        }
+
         final String clientUuid = client.getId();
         UserSessionPredicate predicate = UserSessionPredicate.create(realm.getId()).client(clientUuid);
 
@@ -410,6 +516,12 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
 
     @Override
     public Map<String, Long> getActiveClientSessionStats(RealmModel realm, boolean offline) {
+
+        if (offline && loadOfflineSessionsStatsFromDatabase) {
+            UserSessionPersisterProvider persister = session.getProvider(UserSessionPersisterProvider.class);
+            return persister.getUserSessionsCountsByClients(realm, offline);
+        }
+
         Cache<String, SessionEntityWrapper<UserSessionEntity>> cache = getCache(offline);
         cache = CacheDecorators.skipCacheLoaders(cache);
         return cache.entrySet().stream()
@@ -424,6 +536,13 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
     }
 
      protected long getUserSessionsCount(RealmModel realm, ClientModel client, boolean offline) {
+
+        if (offline && loadOfflineSessionsStatsFromDatabase) {
+            // fetch the actual offline user session count from the database
+            UserSessionPersisterProvider persister = session.getProvider(UserSessionPersisterProvider.class);
+            return persister.getUserSessionsCount(realm, client, offline);
+        }
+
         Cache<String, SessionEntityWrapper<UserSessionEntity>> cache = getCache(offline);
         cache = CacheDecorators.skipCacheLoaders(cache);
 
@@ -662,7 +781,12 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
 
     @Override
     public Stream<UserSessionModel> getOfflineUserSessionsStream(RealmModel realm, UserModel user) {
-        return this.getUserSessionsStream(realm, UserSessionPredicate.create(realm.getId()).user(user.getId()), true);
+
+        if (loadOfflineSessionsStatsFromDatabase) {
+            return getUserSessionsFromPersistenceProviderStream(realm, user, true);
+        }
+
+        return getUserSessionsStream(realm, UserSessionPredicate.create(realm.getId()).user(user.getId()), true);
     }
 
     @Override

model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProviderFactory.java

@@ -95,8 +95,10 @@ public class InfinispanUserSessionProviderFactory implements UserSessionProvider
         Cache<UUID, SessionEntityWrapper<AuthenticatedClientSessionEntity>> clientSessionCache = connections.getCache(InfinispanConnectionProvider.CLIENT_SESSION_CACHE_NAME);
         Cache<UUID, SessionEntityWrapper<AuthenticatedClientSessionEntity>> offlineClientSessionsCache = connections.getCache(InfinispanConnectionProvider.OFFLINE_CLIENT_SESSION_CACHE_NAME);
 
+        boolean loadOfflineSessionsStatsFromDatabase = !isPreloadingOfflineSessionsFromDatabaseEnabled();
+
         return new InfinispanUserSessionProvider(session, remoteCacheInvoker, lastSessionRefreshStore, offlineLastSessionRefreshStore,
-                persisterLastSessionRefreshStore, keyGenerator, cache, offlineSessionsCache, clientSessionCache, offlineClientSessionsCache);
+                persisterLastSessionRefreshStore, keyGenerator, cache, offlineSessionsCache, clientSessionCache, offlineClientSessionsCache, loadOfflineSessionsStatsFromDatabase);
     }
 
     @Override
@@ -145,6 +147,10 @@ public class InfinispanUserSessionProviderFactory implements UserSessionProvider
         });
     }
 
+    private boolean isPreloadingOfflineSessionsFromDatabaseEnabled() {
+        return config.getBoolean("preloadOfflineSessionsFromDatabase", true);
+    }
+
     // Max count of worker errors. Initialization will end with exception when this number is reached
     private int getMaxErrors() {
         return config.getInt("maxErrors", 20);
@@ -163,23 +169,32 @@ public class InfinispanUserSessionProviderFactory implements UserSessionProvider
 
     @Override
     public void loadPersistentSessions(final KeycloakSessionFactory sessionFactory, final int maxErrors, final int sessionsPerSegment) {
-        log.debug("Start pre-loading userSessions from persistent storage");
 
         KeycloakModelUtils.runJobInTransaction(sessionFactory, new KeycloakSessionTask() {
 
             @Override
             public void run(KeycloakSession session) {
-                InfinispanConnectionProvider connections = session.getProvider(InfinispanConnectionProvider.class);
-                Cache<String, Serializable> workCache = connections.getCache(InfinispanConnectionProvider.WORK_CACHE_NAME);
 
-                InfinispanCacheInitializer ispnInitializer = new InfinispanCacheInitializer(sessionFactory, workCache,
-                        new OfflinePersistentUserSessionLoader(sessionsPerSegment), "offlineUserSessions", sessionsPerSegment, maxErrors);
+                if (isPreloadingOfflineSessionsFromDatabaseEnabled()) {
+                    // only preload offline-sessions if necessary
+                    log.debug("Start pre-loading userSessions from persistent storage");
 
-                // DB-lock to ensure that persistent sessions are loaded from DB just on one DC. The other DCs will load them from remote cache.
-                CacheInitializer initializer = new DBLockBasedCacheInitializer(session, ispnInitializer);
+                    InfinispanConnectionProvider connections = session.getProvider(InfinispanConnectionProvider.class);
+                    Cache<String, Serializable> workCache = connections.getCache(InfinispanConnectionProvider.WORK_CACHE_NAME);
 
-                initializer.initCache();
-                initializer.loadSessions();
+                    InfinispanCacheInitializer ispnInitializer = new InfinispanCacheInitializer(sessionFactory, workCache,
+                            new OfflinePersistentUserSessionLoader(sessionsPerSegment), "offlineUserSessions", sessionsPerSegment, maxErrors);
+
+                    // DB-lock to ensure that persistent sessions are loaded from DB just on one DC. The other DCs will load them from remote cache.
+                    CacheInitializer initializer = new DBLockBasedCacheInitializer(session, ispnInitializer);
+
+                    initializer.initCache();
+                    initializer.loadSessions();
+
+                    log.debug("Pre-loading userSessions from persistent storage finished");
+                } else {
+                    log.debug("Skipping pre-loading of userSessions from persistent storage");
+                }
 
                 // Initialize persister for periodically doing bulk DB updates of lastSessionRefresh timestamps of refreshed sessions
                 persisterLastSessionRefreshStore = new PersisterLastSessionRefreshStoreFactory().createAndInit(session, true);
@@ -187,7 +202,7 @@ public class InfinispanUserSessionProviderFactory implements UserSessionProvider
 
         });
 
-        log.debug("Pre-loading userSessions from persistent storage finished");
+
     }
 
 

model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/stream/UserSessionPredicate.java

@@ -108,6 +108,22 @@ public class UserSessionPredicate implements Predicate<Map.Entry<String, Session
         return this;
     }
 
+    /**
+     * Returns the user id.
+     * @return
+     */
+    public String getUserId() {
+        return user;
+    }
+
+    public String getBrokerSessionId() {
+        return brokerSessionId;
+    }
+
+    public String getBrokerUserId() {
+        return brokerUserId;
+    }
+
     @Override
     public boolean test(Map.Entry<String, SessionEntityWrapper<UserSessionEntity>> entry) {
         UserSessionEntity entity = entry.getValue().getEntity();

model/jpa/src/main/java/org/keycloak/models/jpa/session/JpaUserSessionPersisterProvider.java

@@ -38,6 +38,7 @@ import javax.persistence.Query;
 import javax.persistence.TypedQuery;
 
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -171,8 +172,13 @@ public class JpaUserSessionPersisterProvider implements UserSessionPersisterProv
 
     @Override
     public void onRealmRemoved(RealmModel realm) {
-        int num = em.createNamedQuery("deleteClientSessionsByRealm").setParameter("realmId", realm.getId()).executeUpdate();
-        num = em.createNamedQuery("deleteUserSessionsByRealm").setParameter("realmId", realm.getId()).executeUpdate();
+        int deletedClientSessions = em.createNamedQuery("deleteClientSessionsByRealm")
+                .setParameter("realmId", realm.getId())
+                .executeUpdate();
+
+        int deletedUserSessions = em.createNamedQuery("deleteUserSessionsByRealm")
+                .setParameter("realmId", realm.getId())
+                .executeUpdate();
     }
 
     @Override
@@ -243,17 +249,114 @@ public class JpaUserSessionPersisterProvider implements UserSessionPersisterProv
     }
 
     @Override
+    public Map<String, Long> getUserSessionsCountsByClients(RealmModel realm, boolean offline) {
+
+        String offlineStr = offlineToString(offline);
+
+        Query query = em.createNamedQuery("findUserSessionsCountsByClientId");
+
+        query.setParameter("offline", offlineStr);
+        query.setParameter("realmId", realm.getId());
+
+        Map<String, Long> offlineSessionsByClient = new HashMap<>();
+
+        closing(query.getResultStream()).forEach(record -> {
+
+            Object[] row = (Object[]) record;
+
+            String clientId = String.valueOf(row[0]);
+            Long count = ((Number)row[1]).longValue();
+
+            offlineSessionsByClient.put(clientId, count);
+        });
+
+        return offlineSessionsByClient;
+    }
+
+    @Override
+    public UserSessionModel loadUserSession(RealmModel realm, String userSessionId, boolean offline) {
+
+        String offlineStr = offlineToString(offline);
+
+        TypedQuery<PersistentUserSessionEntity> userSessionQuery = em.createNamedQuery("findUserSession", PersistentUserSessionEntity.class);
+        userSessionQuery.setParameter("realmId", realm.getId());
+        userSessionQuery.setParameter("offline", offlineStr);
+        userSessionQuery.setParameter("userSessionId", userSessionId);
+        userSessionQuery.setMaxResults(1);
+
+        Stream<PersistentUserSessionAdapter> persistentUserSessions = closing(userSessionQuery.getResultStream().map(this::toAdapter));
+
+        return persistentUserSessions.findAny().map(userSession -> {
+
+            TypedQuery<PersistentClientSessionEntity> clientSessionQuery = em.createNamedQuery("findClientSessionsByUserSession", PersistentClientSessionEntity.class);
+            clientSessionQuery.setParameter("userSessionId", Collections.singleton(userSessionId));
+            clientSessionQuery.setParameter("offline", offlineStr);
+
+            Set<String> removedClientUUIDs = new HashSet<>();
+
+            clientSessionQuery.getResultStream().forEach(clientSession -> {
+                        boolean added = addClientSessionToAuthenticatedClientSessionsIfPresent(userSession, clientSession);
+                        if (!added) {
+                            // client was removed in the meantime
+                            removedClientUUIDs.add(clientSession.getClientId());
+                        }
+                    }
+            );
+
+            removedClientUUIDs.forEach(this::onClientRemoved);
+
+            return userSession;
+        }).orElse(null);
+    }
+
+    @Override
+    public Stream<UserSessionModel> loadUserSessionsStream(RealmModel realm, ClientModel client, boolean offline, Integer firstResult, Integer maxResults) {
+
+        String offlineStr = offlineToString(offline);
+
+        TypedQuery<PersistentUserSessionEntity> query = paginateQuery(
+                em.createNamedQuery("findUserSessionsByClientId", PersistentUserSessionEntity.class),
+                firstResult, maxResults);
+
+        query.setParameter("offline", offlineStr);
+        query.setParameter("realmId", realm.getId());
+        query.setParameter("clientId", client.getId());
+
+        return loadUserSessionsWithClientSessions(query, offlineStr);
+    }
+
+    @Override
+    public Stream<UserSessionModel> loadUserSessionsStream(RealmModel realm, UserModel user, boolean offline, Integer firstResult, Integer maxResults) {
+
+        String offlineStr = offlineToString(offline);
+
+        TypedQuery<PersistentUserSessionEntity> query = paginateQuery(
+                em.createNamedQuery("findUserSessionsByUserId", PersistentUserSessionEntity.class),
+                firstResult, maxResults);
+
+        query.setParameter("offline", offlineStr);
+        query.setParameter("realmId", realm.getId());
+        query.setParameter("userId", user.getId());
+
+        return loadUserSessionsWithClientSessions(query, offlineStr);
+    }
+
     public Stream<UserSessionModel> loadUserSessionsStream(Integer firstResult, Integer maxResults, boolean offline,
                                                            String lastUserSessionId) {
         String offlineStr = offlineToString(offline);
 
-        TypedQuery<PersistentUserSessionEntity> queryUserSessions = em.createNamedQuery("findUserSessionsOrderedById", PersistentUserSessionEntity.class);
-        queryUserSessions.setParameter("offline", offlineStr);
-        queryUserSessions.setParameter("lastSessionId", lastUserSessionId);
+        TypedQuery<PersistentUserSessionEntity> query = paginateQuery(em.createNamedQuery("findUserSessionsOrderedById", PersistentUserSessionEntity.class)
+            .setParameter("offline", offlineStr)
+            .setParameter("lastSessionId", lastUserSessionId), firstResult, maxResults);
 
-        List<PersistentUserSessionAdapter> userSessionAdapters = closing(paginateQuery(queryUserSessions, firstResult, maxResults).getResultStream()
-                .map(this::toAdapter))
-                .filter(Objects::nonNull)
+        return loadUserSessionsWithClientSessions(query, offlineStr);
+    }
+
+    private Stream<UserSessionModel> loadUserSessionsWithClientSessions(TypedQuery<PersistentUserSessionEntity> query, String offlineStr) {
+
+        List<PersistentUserSessionAdapter> userSessionAdapters = closing(query.getResultStream()
+                .map(this::toAdapter)
+                .filter(Objects::nonNull))
                 .collect(Collectors.toList());
 
         Map<String, PersistentUserSessionAdapter> sessionsById = userSessionAdapters.stream()
@@ -272,15 +375,10 @@ public class JpaUserSessionPersisterProvider implements UserSessionPersisterProv
 
             closing(queryClientSessions.getResultStream()).forEach(clientSession -> {
                 PersistentUserSessionAdapter userSession = sessionsById.get(clientSession.getUserSessionId());
-
-                PersistentAuthenticatedClientSessionAdapter clientSessAdapter = toAdapter(userSession.getRealm(), userSession, clientSession);
-                Map<String, AuthenticatedClientSessionModel> currentClientSessions = userSession.getAuthenticatedClientSessions();
-
-                // Case when client was removed in the meantime
-                if (clientSessAdapter.getClient() == null) {
+                boolean added = addClientSessionToAuthenticatedClientSessionsIfPresent(userSession, clientSession);
+                if (!added) {
+                    // client was removed in the meantime
                     removedClientUUIDs.add(clientSession.getClientId());
-                } else {
-                    currentClientSessions.put(clientSession.getClientId(), clientSessAdapter);
                 }
             });
         }
@@ -292,6 +390,18 @@ public class JpaUserSessionPersisterProvider implements UserSessionPersisterProv
         return userSessionAdapters.stream().map(UserSessionModel.class::cast);
     }
 
+    private boolean addClientSessionToAuthenticatedClientSessionsIfPresent(PersistentUserSessionAdapter userSession, PersistentClientSessionEntity clientSessionEntity) {
+
+        PersistentAuthenticatedClientSessionAdapter clientSessAdapter = toAdapter(userSession.getRealm(), userSession, clientSessionEntity);
+
+        if (clientSessAdapter.getClient() == null) {
+            return false;
+        }
+
+        userSession.getAuthenticatedClientSessions().put(clientSessionEntity.getClientId(), clientSessAdapter);
+        return true;
+    }
+
     private PersistentUserSessionAdapter toAdapter(PersistentUserSessionEntity entity) {
         RealmModel realm = session.realms().getRealm(entity.getRealmId());
         if (realm == null) {    // Realm has been deleted concurrently, ignore the entity
@@ -339,8 +449,21 @@ public class JpaUserSessionPersisterProvider implements UserSessionPersisterProv
     }
 
     @Override
-    public void close() {
+    public int getUserSessionsCount(RealmModel realm, ClientModel clientModel, boolean offline) {
 
+        String offlineStr = offlineToString(offline);
+
+        Query query = em.createNamedQuery("findClientSessionsCountByClient");
+        // Note, that realm is unused here, since the clientModel id already determines the offline user-sessions bound to a owning realm.
+        query.setParameter("offline", offlineStr);
+        query.setParameter("clientId", clientModel.getId());
+        Number n = (Number) query.getSingleResult();
+        return n.intValue();
+    }
+
+    @Override
+    public void close() {
+        // NOOP
     }
 
     private String offlineToString(boolean offline) {

model/jpa/src/main/java/org/keycloak/models/jpa/session/PersistentClientSessionEntity.java

@@ -38,7 +38,8 @@ import java.io.Serializable;
         @NamedQuery(name="deleteClientSessionsByUserSession", query="delete from PersistentClientSessionEntity sess where sess.userSessionId = :userSessionId and sess.offline = :offline"),
         @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="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="findClientSessionsCountByClient", query="select count(sess) from PersistentClientSessionEntity sess where sess.offline = :offline and sess.clientId = :clientId")
 })
 @Table(name="OFFLINE_CLIENT_SESSION")
 @Entity

model/jpa/src/main/java/org/keycloak/models/jpa/session/PersistentUserSessionEntity.java

@@ -40,7 +40,21 @@ import java.io.Serializable;
         @NamedQuery(name="findUserSessionsCount", query="select count(sess) from PersistentUserSessionEntity sess where sess.offline = :offline"),
         @NamedQuery(name="findUserSessionsOrderedById", query="select sess from PersistentUserSessionEntity sess, RealmEntity realm where realm.id = sess.realmId AND sess.offline = :offline" +
                 " AND sess.userSessionId > :lastSessionId" +
-                " order by sess.userSessionId")
+                " order by sess.userSessionId"),
+        @NamedQuery(name="findUserSession", query="select sess from PersistentUserSessionEntity sess where sess.offline = :offline" +
+                " AND sess.userSessionId = :userSessionId AND sess.realmId = :realmId"),
+        @NamedQuery(name="findUserSessionsByUserId", query="select sess from PersistentUserSessionEntity sess where sess.offline = :offline" +
+                " AND sess.realmId = :realmId AND sess.userId = :userId ORDER BY sess.userSessionId"),
+        @NamedQuery(name="findUserSessionsByClientId", query="SELECT sess FROM PersistentUserSessionEntity sess INNER JOIN PersistentClientSessionEntity clientSess " +
+                " ON sess.userSessionId = clientSess.userSessionId AND clientSess.clientId = :clientId WHERE sess.offline = :offline " +
+                " AND sess.userSessionId = clientSess.userSessionId AND sess.realmId = :realmId ORDER BY sess.userSessionId"),
+        @NamedQuery(name="findUserSessionsCountsByClientId", query="SELECT clientSess.clientId, count(clientSess) " +
+                " FROM PersistentUserSessionEntity sess INNER JOIN PersistentClientSessionEntity clientSess " +
+                " ON sess.userSessionId = clientSess.userSessionId " +
+                // find all available offline user-session for all clients in a realm
+                " WHERE sess.offline = :offline " +
+                " AND sess.userSessionId = clientSess.userSessionId AND sess.realmId = :realmId " +
+                " GROUP BY clientSess.clientId")
 
 })
 @Table(name="OFFLINE_USER_SESSION")

model/jpa/src/main/resources/META-INF/jpa-changelog-14.0.0.xml

@@ -17,6 +17,25 @@
   -->
 <databaseChangeLog xmlns="http://www.liquibase.org/xml/ns/dbchangelog" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://www.liquibase.org/xml/ns/dbchangelog http://www.liquibase.org/xml/ns/dbchangelog/dbchangelog-3.1.xsd">
 
+    <changeSet author="keycloak" id="14.0.0-KEYCLOAK-11019">
+        <createIndex tableName="OFFLINE_CLIENT_SESSION" indexName="IDX_OFFLINE_CSS_PRELOAD">
+            <column name="CLIENT_ID" type="VARCHAR(36)"/>
+            <column name="OFFLINE_FLAG" type="VARCHAR(4)"/>
+        </createIndex>
+
+        <createIndex tableName="OFFLINE_USER_SESSION" indexName="IDX_OFFLINE_USS_BY_USER">
+            <column name="USER_ID" type="VARCHAR(36)"/>
+            <column name="REALM_ID" type="VARCHAR(36)"/>
+            <column name="OFFLINE_FLAG" type="VARCHAR(4)"/>
+        </createIndex>
+
+        <createIndex tableName="OFFLINE_USER_SESSION" indexName="IDX_OFFLINE_USS_BY_USERSESS">
+            <column name="REALM_ID" type="VARCHAR(36)"/>
+            <column name="OFFLINE_FLAG" type="VARCHAR(4)"/>
+            <column name="USER_SESSION_ID" type="VARCHAR(36)"/>
+        </createIndex>
+    </changeSet>
+
     <changeSet author="keycloak" id="KEYCLOAK-17267-add-index-to-user-attributes">
         <createIndex indexName="IDX_USER_ATTRIBUTE_NAME" tableName="USER_ATTRIBUTE">
             <column name="NAME" type="VARCHAR(255)"/>

server-spi-private/src/main/java/org/keycloak/models/session/DisabledUserSessionPersisterProvider.java

@@ -27,6 +27,8 @@ import org.keycloak.models.UserModel;
 import org.keycloak.models.UserSessionModel;
 
 import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
 import java.util.stream.Stream;
 
 /**
@@ -109,6 +111,20 @@ public class DisabledUserSessionPersisterProvider implements UserSessionPersiste
 
     }
 
+    public UserSessionModel loadUserSession(RealmModel realm, String userSessionId, boolean offline) {
+        return null;
+    }
+
+    @Override
+    public Stream<UserSessionModel> loadUserSessionsStream(RealmModel realm, ClientModel client, boolean offline, Integer firstResult, Integer maxResults) {
+        return Stream.empty();
+    }
+
+    @Override
+    public Stream<UserSessionModel> loadUserSessionsStream(RealmModel realm, UserModel user, boolean offline, Integer firstResult, Integer maxResults) {
+        return Stream.empty();
+    }
+
     @Override
     public Stream<UserSessionModel> loadUserSessionsStream(Integer firstResult, Integer maxResults, boolean offline,
                                                            String lastUserSessionId) {
@@ -119,4 +135,14 @@ public class DisabledUserSessionPersisterProvider implements UserSessionPersiste
     public int getUserSessionsCount(boolean offline) {
         return 0;
     }
+
+    @Override
+    public int getUserSessionsCount(RealmModel realm, ClientModel clientModel, boolean offline) {
+        return 0;
+    }
+
+    @Override
+    public Map<String, Long> getUserSessionsCountsByClients(RealmModel realm, boolean offline) {
+        return Collections.emptyMap();
+    }
 }

server-spi-private/src/main/java/org/keycloak/models/session/UserSessionPersisterProvider.java

@@ -26,6 +26,7 @@ import org.keycloak.provider.Provider;
 
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -56,6 +57,37 @@ public interface UserSessionPersisterProvider extends Provider {
     // Remove userSessions and clientSessions, which are expired
     void removeExpired(RealmModel realm);
 
+    /**
+     * Loads the user session with the given userSessionId.
+     * @param userSessionId
+     * @param offline
+     * @return
+     */
+    UserSessionModel loadUserSession(RealmModel realm, String userSessionId, boolean offline);
+
+    /**
+     * Loads the user sessions for the given {@link UserModel} in the given {@link RealmModel} if present.
+     * @param realm
+     * @param user
+     * @param offline
+     * @param firstResult
+     * @param maxResults
+     * @return
+     */
+    Stream<UserSessionModel> loadUserSessionsStream(RealmModel realm, UserModel user, boolean offline, Integer firstResult, Integer maxResults);
+
+    /**
+     * Loads the user sessions for the given {@link ClientModel} in the given {@link RealmModel} if present.
+     *
+     * @param realm
+     * @param client
+     * @param offline
+     * @param firstResult
+     * @param maxResults
+     * @return
+     */
+    Stream<UserSessionModel> loadUserSessionsStream(RealmModel realm, ClientModel client, boolean offline, Integer firstResult, Integer maxResults);
+
     /**
      * Called during startup. For each userSession, it loads also clientSessions
      * @deprecated Use {@link #loadUserSessionsStream(Integer, Integer, boolean, String) loadUserSessionsStream} instead.
@@ -77,6 +109,31 @@ public interface UserSessionPersisterProvider extends Provider {
     Stream<UserSessionModel> loadUserSessionsStream(Integer firstResult, Integer maxResults, boolean offline,
                                                     String lastUserSessionId);
 
+    /**
+     * Retrieves the count of user sessions for all realms.
+     *
+     * @param offline
+     * @return
+     *
+     */
     int getUserSessionsCount(boolean offline);
 
+    /**
+     * Retrieves the count of user client-sessions for the given client
+     *
+     * @param realm
+     * @param clientModel
+     * @param offline
+     * @return
+     */
+    int getUserSessionsCount(RealmModel realm, ClientModel clientModel, boolean offline);
+
+    /**
+     * Returns a {@link Map} containing the number of user-sessions aggregated by client id for the given realm.
+     * @param realm
+     * @param offline
+     * @return the count {@link Map} with clientId as key and session count as value
+     */
+    Map<String, Long> getUserSessionsCountsByClients(RealmModel realm, boolean offline);
+
 }

testsuite/model/pom.xml

@@ -26,6 +26,7 @@
         <log4j.configuration>file:${project.build.directory}/dependency/log4j.properties</log4j.configuration>
         <jacoco.skip>true</jacoco.skip>
         <keycloak.profile.feature.map_storage>disabled</keycloak.profile.feature.map_storage>
+        <keycloak.userSessions.infinispan.preloadOfflineSessionsFromDatabase>true</keycloak.userSessions.infinispan.preloadOfflineSessionsFromDatabase>
     </properties>
     
     <dependencies>
@@ -137,6 +138,7 @@
                         <keycloak.connectionsJpa.default.url>${keycloak.connectionsJpa.url}</keycloak.connectionsJpa.default.url>
                         <log4j.configuration>file:${project.build.directory}/test-classes/log4j.properties</log4j.configuration> <!-- for the logging to properly work with tests in the 'other' module -->
                         <keycloak.profile.feature.map_storage>${keycloak.profile.feature.map_storage}</keycloak.profile.feature.map_storage>
+                        <keycloak.userSessions.infinispan.preloadOfflineSessionsFromDatabase>${keycloak.userSessions.infinispan.preloadOfflineSessionsFromDatabase}</keycloak.userSessions.infinispan.preloadOfflineSessionsFromDatabase>
                     </systemPropertyVariables>
                 </configuration>
             </plugin>
@@ -178,6 +180,14 @@
             </properties>
         </profile>
 
+        <profile>
+            <id>jpa+infinispan-sessions-preloading-disabled</id>
+            <properties>
+                <keycloak.model.parameters>Infinispan,Jpa</keycloak.model.parameters>
+                <keycloak.userSessions.infinispan.preloadOfflineSessionsFromDatabase>false</keycloak.userSessions.infinispan.preloadOfflineSessionsFromDatabase>
+            </properties>
+        </profile>
+
         <profile>
             <id>jpa-federation+infinispan</id>
             <properties>

testsuite/model/src/test/java/org/keycloak/testsuite/model/session/OfflineSessionPersistenceTest.java

@@ -26,6 +26,9 @@ import org.keycloak.models.UserModel;
 import org.keycloak.models.UserProvider;
 import org.keycloak.models.UserSessionModel;
 import org.keycloak.models.UserSessionProvider;
+import org.keycloak.models.session.UserSessionPersisterProvider;
+import org.keycloak.models.sessions.infinispan.InfinispanUserSessionProvider;
+import org.keycloak.models.sessions.infinispan.InfinispanUserSessionProviderFactory;
 import org.keycloak.services.managers.RealmManager;
 import org.keycloak.testsuite.model.KeycloakModelTest;
 import org.keycloak.testsuite.model.RequireProvider;
@@ -33,6 +36,7 @@ import java.util.Collection;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Random;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
@@ -42,6 +46,7 @@ import java.util.stream.Stream;
 import org.hamcrest.Matchers;
 import org.junit.Test;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsInAnyOrder;
 
 /**
  *
@@ -190,6 +195,82 @@ public class OfflineSessionPersistenceTest extends KeycloakModelTest {
         assertOfflineSessionsExist(realmId, offlineSessionIds);
     }
 
+    @Test
+    @RequireProvider(UserSessionPersisterProvider.class)
+    @RequireProvider(value = UserSessionProvider.class, only = InfinispanUserSessionProviderFactory.PROVIDER_ID)
+    public void testOfflineSessionLoadingAfterCacheRemoval() {
+        List<String> offlineSessionIds = createOfflineSessions(realmId, userIds);
+        assertOfflineSessionsExist(realmId, offlineSessionIds);
+
+        // Simulate server restart
+        reinitializeKeycloakSessionFactory();
+        assertOfflineSessionsExist(realmId, offlineSessionIds);
+
+        // remove sessions from the cache
+        withRealm(realmId, (session, realm) -> {
+            // Delete local user cache (persisted sessions are still kept)
+            UserSessionProvider provider = session.getProvider(UserSessionProvider.class);
+            // Remove in-memory representation of the offline sessions
+            ((InfinispanUserSessionProvider) provider).removeLocalUserSessions(realm.getId(), true);
+
+            return null;
+        });
+
+        // assert sessions are lazily loaded from DB
+        assertOfflineSessionsExist(realmId, offlineSessionIds);
+    }
+
+    @Test
+    @RequireProvider(UserSessionPersisterProvider.class)
+    @RequireProvider(value = UserSessionProvider.class, only = InfinispanUserSessionProviderFactory.PROVIDER_ID)
+    public void testLazyClientSessionStatsFetching() {
+        List<String> clientIds = withRealm(realmId, (session, realm) -> IntStream.range(0, 5)
+                .mapToObj(cid -> session.clients().addClient(realm, "client-" + cid))
+                .map(ClientModel::getId)
+                .collect(Collectors.toList()));
+
+        List<String> offlineSessionIds = createOfflineSessions(realmId, userIds);
+        assertOfflineSessionsExist(realmId, offlineSessionIds);
+
+        Random r = new Random();
+        offlineSessionIds.stream().forEach(offlineSessionId -> createOfflineClientSession(offlineSessionId, clientIds.get(r.nextInt(5))));
+
+        // Simulate server restart
+        reinitializeKeycloakSessionFactory();
+
+        // load active client sessions stats from DB
+        Map<String, Long> sessionStats = withRealm(realmId, (session, realm) -> session.sessions().getActiveClientSessionStats(realm, true));
+
+        long client1SessionCount = sessionStats.get(clientIds.get(0));
+        int clientSessionsCount = sessionStats.values().stream().reduce(0l, Long::sum).intValue();
+        assertThat(clientSessionsCount, Matchers.is(USER_COUNT * OFFLINE_SESSION_COUNT_PER_USER));
+
+        // Simulate server restart
+        reinitializeKeycloakSessionFactory();
+
+        long actualClient1SessionCount = withRealm(realmId, (session, realm) -> {
+            ClientModel client = realm.getClientById(clientIds.get(0));
+            return session.sessions().getOfflineSessionsCount(realm, client);
+        });
+        assertThat(actualClient1SessionCount, Matchers.is(client1SessionCount));
+    }
+
+    @Test
+    @RequireProvider(UserSessionPersisterProvider.class)
+    @RequireProvider(value = UserSessionProvider.class, only = InfinispanUserSessionProviderFactory.PROVIDER_ID)
+    public void testLazyOfflineUserSessionFetching() {
+        List<String> offlineSessionIds = createOfflineSessions(realmId, userIds);
+        assertOfflineSessionsExist(realmId, offlineSessionIds);
+
+        // Simulate server restart
+        reinitializeKeycloakSessionFactory();
+
+        List<String> actualOfflineSessionIds = withRealm(realmId, (session, realm) -> session.users().getUsersStream(realm).flatMap(user ->
+                session.sessions().getOfflineUserSessionsStream(realm, user)).map(UserSessionModel::getId).collect(Collectors.toList()));
+
+        assertThat(actualOfflineSessionIds, containsInAnyOrder(offlineSessionIds.toArray()));
+    }
+
     private String createOfflineClientSession(String offlineUserSessionId, String clientId) {
         return withRealm(realmId, (session, realm) -> {
             UserSessionModel offlineUserSession = session.sessions().getOfflineUserSession(realm, offlineUserSessionId);

testsuite/model/src/test/java/org/keycloak/testsuite/model/session/UserSessionInitializerTest.java

@@ -20,7 +20,6 @@ package org.keycloak.testsuite.model.session;
 import org.junit.Assert;
 import org.junit.Test;
 import org.keycloak.common.util.Time;
-import org.keycloak.connections.infinispan.InfinispanConnectionProvider;
 import org.keycloak.models.AuthenticatedClientSessionModel;
 import org.keycloak.models.ClientModel;
 import org.keycloak.models.Constants;
@@ -181,25 +180,6 @@ public class UserSessionInitializerTest extends KeycloakModelTest {
             if (provider instanceof InfinispanUserSessionProvider) {
                 // Remove in-memory representation of the offline sessions
                 ((InfinispanUserSessionProvider) provider).removeLocalUserSessions(realm.getId(), true);
-
-                // Clear ispn cache to ensure initializerState is removed as well
-                InfinispanConnectionProvider infinispan = session.getProvider(InfinispanConnectionProvider.class);
-                if (infinispan != null) {
-                    infinispan.getCache(InfinispanConnectionProvider.WORK_CACHE_NAME).clear();
-                }
-            }
-        });
-
-        inComittedTransaction(session -> {
-            // This is only valid in infinispan provider where the offline session is loaded upon start and never reloaded
-            UserSessionProvider provider = session.getProvider(UserSessionProvider.class);
-            if (provider instanceof InfinispanUserSessionProvider) {
-                RealmModel realm = session.realms().getRealm(realmId);
-
-                ClientModel testApp = realm.getClientByClientId("test-app");
-                ClientModel thirdparty = realm.getClientByClientId("third-party");
-                assertThat("Count of offline sessions for client 'test-app'", session.sessions().getOfflineSessionsCount(realm, testApp), is((long) 0));
-                assertThat("Count of offline sessions for client 'third-party'", session.sessions().getOfflineSessionsCount(realm, thirdparty), is((long) 0));
             }
         });