hmlnarik · May 13, 2021 · Apr 2, 2021
diff --git a/.../java/org/keycloak/connections/infinispan/DefaultInfinispanConnectionProviderFactory.java b/.../java/org/keycloak/connections/infinispan/DefaultInfinispanConnectionProviderFactory.java
@@ -46,6 +46,7 @@
 import org.keycloak.cluster.ClusterProvider;
 import org.keycloak.cluster.ManagedCacheManagerProvider;
 import org.keycloak.cluster.infinispan.KeycloakHotRodMarshallerFactory;
+import org.keycloak.common.util.Time;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
 
@@ -56,7 +57,14 @@
 import org.keycloak.models.utils.PostMigrationEvent;
 import org.keycloak.provider.InvalidationHandler.ObjectType;
 import org.keycloak.provider.ProviderEvent;
+import java.time.Instant;
+import java.util.concurrent.atomic.AtomicReference;
+import org.infinispan.commons.time.TimeService;
+import org.infinispan.factories.GlobalComponentRegistry;
+import org.infinispan.factories.impl.BasicComponentRegistry;
+import org.infinispan.factories.impl.ComponentRef;
 import org.infinispan.persistence.remote.configuration.RemoteStoreConfigurationBuilder;
+import org.infinispan.util.EmbeddedTimeService;
 import static org.keycloak.models.cache.infinispan.InfinispanCacheRealmProviderFactory.REALM_CLEAR_CACHE_EVENTS;
 import static org.keycloak.models.cache.infinispan.InfinispanCacheRealmProviderFactory.REALM_INVALIDATION_EVENTS;
 
@@ -191,6 +199,7 @@ protected void initEmbedded() {
 
         boolean clustered = config.getBoolean("clustered", false);
         boolean async = config.getBoolean("async", false);
+        boolean useKeycloakTimeService = config.getBoolean("useKeycloakTimeService", false);
 
         this.topologyInfo = new TopologyInfo(cacheManager, config, true);
 
@@ -209,6 +218,9 @@ protected void initEmbedded() {
         gcb.serialization().marshaller(new JBossUserMarshaller());
 
         cacheManager = new DefaultCacheManager(gcb.build());
+        if (useKeycloakTimeService) {
+            setTimeServiceToKeycloakTime(cacheManager);
+        }
         containerManaged = false;
 
         logger.debug("Started embedded Infinispan cache container");
@@ -354,6 +366,25 @@ protected void initEmbedded() {
         cacheManager.getCache(InfinispanConnectionProvider.AUTHORIZATION_REVISIONS_CACHE_NAME, true);
     }
 
+    /**
+     * Replaces the {@link TimeService} in infinispan with the one that respects Keycloak {@link Time}.
+     * @param cacheManager
+     * @return Runnable to revert replacement of the infinispan time service
+     */
+    public static Runnable setTimeServiceToKeycloakTime(EmbeddedCacheManager cacheManager) {
+        TimeService previousTimeService = replaceComponent(cacheManager, TimeService.class, KEYCLOAK_TIME_SERVICE, true);
+        AtomicReference<TimeService> ref = new AtomicReference<>(previousTimeService);
+        return () -> {
+            if (ref.get() == null) {
+                logger.warn("Calling revert of the TimeService when testing TimeService was already reverted");
+                return;
+            }
+
+            logger.info("Revert set KeycloakIspnTimeService to the infinispan cacheManager");
+
+            replaceComponent(cacheManager, TimeService.class, ref.getAndSet(null), true);
+        };
+    }
 
     private Configuration getRevisionCacheConfig(long maxEntries) {
         ConfigurationBuilder cb = createCacheConfigurationBuilder();
@@ -545,4 +576,50 @@ private void registerSystemWideListeners(KeycloakSession session) {
         });
     }
 
+
+    /**
+     * Forked from org.infinispan.test.TestingUtil class
+     *
+     * Replaces a component in a running cache manager (global component registry).
+     *
+     * @param cacheMgr       cache in which to replace component
+     * @param componentType        component type of which to replace
+     * @param replacementComponent new instance
+     * @param rewire               if true, ComponentRegistry.rewire() is called after replacing.
+     *
+     * @return the original component that was replaced
+     */
+    private static <T> T replaceComponent(EmbeddedCacheManager cacheMgr, Class<T> componentType, T replacementComponent, boolean rewire) {
+        GlobalComponentRegistry cr = cacheMgr.getGlobalComponentRegistry();
+        BasicComponentRegistry bcr = cr.getComponent(BasicComponentRegistry.class);
+        ComponentRef<T> old = bcr.getComponent(componentType);
+        bcr.replaceComponent(componentType.getName(), replacementComponent, true);
+        if (rewire) {
+            cr.rewire();
+            cr.rewireNamedRegistries();
+        }
+        return old != null ? old.wired() : null;
+    }
+
+    public static final TimeService KEYCLOAK_TIME_SERVICE = new EmbeddedTimeService() {
+
+        private long getCurrentTimeMillis() {
+            return Time.currentTimeMillis();
+        }
+
+        @Override
+        public long wallClockTime() {
+            return getCurrentTimeMillis();
+        }
+
+        @Override
+        public long time() {
+            return TimeUnit.MILLISECONDS.toNanos(getCurrentTimeMillis());
+        }
+
+        @Override
+        public Instant instant() {
+            return Instant.ofEpochMilli(getCurrentTimeMillis());
+        }
+    };
 }
diff --git a/...g/keycloak/models/sessions/infinispan/initializer/OfflinePersistentUserSessionLoader.java b/...g/keycloak/models/sessions/infinispan/initializer/OfflinePersistentUserSessionLoader.java
@@ -38,7 +38,7 @@ public class OfflinePersistentUserSessionLoader implements SessionLoader<Offline
         OfflinePersistentWorkerContext, OfflinePersistentWorkerResult>, Serializable {
 
     // Placeholder String used in the searching conditions to identify very first session
-    private static final String FIRST_SESSION_ID = "000";
+    private static final String FIRST_SESSION_ID = "00000000-0000-0000-0000-000000000000";
 
     private static final Logger log = Logger.getLogger(OfflinePersistentUserSessionLoader.class);
 
@@ -72,39 +72,36 @@ public OfflinePersistentLoaderContext computeLoaderContext(KeycloakSession sessi
 
     @Override
     public OfflinePersistentWorkerContext computeWorkerContext(OfflinePersistentLoaderContext loaderCtx, int segment, int workerId, OfflinePersistentWorkerResult previousResult) {
-        int lastCreatedOn;
         String lastSessionId;
         if (previousResult == null) {
-            lastCreatedOn = 0;
             lastSessionId = FIRST_SESSION_ID;
         } else {
-            lastCreatedOn = previousResult.getLastCreatedOn();
             lastSessionId = previousResult.getLastSessionId();
         }
 
         // We know the last loaded session. New workers iteration will start from this place
-        return new OfflinePersistentWorkerContext(segment, workerId, lastCreatedOn, lastSessionId);
+        return new OfflinePersistentWorkerContext(segment, workerId, lastSessionId);
     }
 
 
     @Override
     public OfflinePersistentWorkerResult createFailedWorkerResult(OfflinePersistentLoaderContext loaderContext, OfflinePersistentWorkerContext workerContext) {
-        return new OfflinePersistentWorkerResult(false, workerContext.getSegment(), workerContext.getWorkerId(), -1, FIRST_SESSION_ID);
+        return new OfflinePersistentWorkerResult(false, workerContext.getSegment(), workerContext.getWorkerId(), FIRST_SESSION_ID);
     }
 
 
     @Override
     public OfflinePersistentWorkerResult loadSessions(KeycloakSession session, OfflinePersistentLoaderContext loaderContext, OfflinePersistentWorkerContext ctx) {
         int first = ctx.getWorkerId() * sessionsPerSegment;
 
-        log.tracef("Loading sessions for segment=%d createdOn=%d lastSessionId=%s", ctx.getSegment(), ctx.getLastCreatedOn(), ctx.getLastSessionId());
+        log.tracef("Loading sessions for segment=%d lastSessionId=%s", ctx.getSegment(), ctx.getLastSessionId());
 
         UserSessionPersisterProvider persister = session.getProvider(UserSessionPersisterProvider.class);
         List<UserSessionModel> sessions = persister
-                .loadUserSessionsStream(first, sessionsPerSegment, true, ctx.getLastCreatedOn(), ctx.getLastSessionId())
+                .loadUserSessionsStream(first, sessionsPerSegment, true, ctx.getLastSessionId())
                 .collect(Collectors.toList());
 
-        log.tracef("Sessions loaded from DB - segment=%d createdOn=%d lastSessionId=%s", ctx.getSegment(), ctx.getLastCreatedOn(), ctx.getLastSessionId());
+        log.tracef("Sessions loaded from DB - segment=%d lastSessionId=%s", ctx.getSegment(), ctx.getLastSessionId());
 
         UserSessionModel lastSession = null;
         if (!sessions.isEmpty()) {
@@ -114,12 +111,11 @@ public OfflinePersistentWorkerResult loadSessions(KeycloakSession session, Offli
             session.sessions().importUserSessions(sessions, true);
         }
 
-        int lastCreatedOn = lastSession==null ? Time.currentTime() + 100000 : lastSession.getStarted();
         String lastSessionId = lastSession==null ? FIRST_SESSION_ID : lastSession.getId();
 
-        log.tracef("Sessions imported to infinispan - segment: %d, lastCreatedOn: %d, lastSessionId: %s", ctx.getSegment(), lastCreatedOn, lastSessionId);
+        log.tracef("Sessions imported to infinispan - segment: %d, lastSessionId: %s", ctx.getSegment(), lastSessionId);
 
-        return new OfflinePersistentWorkerResult(true, ctx.getSegment(), ctx.getWorkerId(), lastCreatedOn, lastSessionId);
+        return new OfflinePersistentWorkerResult(true, ctx.getSegment(), ctx.getWorkerId(), lastSessionId);
     }
 
 

diff --git a/...a/org/keycloak/models/sessions/infinispan/initializer/OfflinePersistentWorkerContext.java b/...a/org/keycloak/models/sessions/infinispan/initializer/OfflinePersistentWorkerContext.java
@@ -22,20 +22,13 @@
  */
 public class OfflinePersistentWorkerContext extends SessionLoader.WorkerContext {
 
-    private final int lastCreatedOn;
     private final String lastSessionId;
 
-    public OfflinePersistentWorkerContext(int segment, int workerId, int lastCreatedOn, String lastSessionId) {
+    public OfflinePersistentWorkerContext(int segment, int workerId, String lastSessionId) {
         super(segment, workerId);
-        this.lastCreatedOn = lastCreatedOn;
         this.lastSessionId = lastSessionId;
     }
 
-
-    public int getLastCreatedOn() {
-        return lastCreatedOn;
-    }
-
     public String getLastSessionId() {
         return lastSessionId;
     }

diff --git a/...va/org/keycloak/models/sessions/infinispan/initializer/OfflinePersistentWorkerResult.java b/...va/org/keycloak/models/sessions/infinispan/initializer/OfflinePersistentWorkerResult.java
@@ -22,22 +22,14 @@
  */
 public class OfflinePersistentWorkerResult extends SessionLoader.WorkerResult {
 
-    private final int lastCreatedOn;
     private final String lastSessionId;
 
 
-    public OfflinePersistentWorkerResult(boolean success, int segment, int workerId, int lastCreatedOn, String lastSessionId) {
+    public OfflinePersistentWorkerResult(boolean success, int segment, int workerId, String lastSessionId) {
         super(success, segment, workerId);
-        this.lastCreatedOn = lastCreatedOn;
         this.lastSessionId = lastSessionId;
     }
 
-
-    public int getLastCreatedOn() {
-        return lastCreatedOn;
-    }
-
-
     public String getLastSessionId() {
         return lastSessionId;
     }

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
@@ -244,31 +244,33 @@ public void removeExpired(RealmModel realm) {
 
     @Override
     public Stream<UserSessionModel> loadUserSessionsStream(Integer firstResult, Integer maxResults, boolean offline,
-                                                           Integer lastCreatedOn, String lastUserSessionId) {
+                                                           String lastUserSessionId) {
         String offlineStr = offlineToString(offline);
 
-        TypedQuery<PersistentUserSessionEntity> query = em.createNamedQuery("findUserSessions", PersistentUserSessionEntity.class);
-        query.setParameter("offline", offlineStr);
-        query.setParameter("lastCreatedOn", lastCreatedOn);
-        query.setParameter("lastSessionId", lastUserSessionId);
+        TypedQuery<PersistentUserSessionEntity> queryUserSessions = em.createNamedQuery("findUserSessionsOrderedById", PersistentUserSessionEntity.class);
+        queryUserSessions.setParameter("offline", offlineStr);
+        queryUserSessions.setParameter("lastSessionId", lastUserSessionId);
 
-        List<PersistentUserSessionAdapter> result = closing(paginateQuery(query, firstResult, maxResults).getResultStream()
+        List<PersistentUserSessionAdapter> userSessionAdapters = closing(paginateQuery(queryUserSessions, firstResult, maxResults).getResultStream()
                 .map(this::toAdapter))
                 .filter(Objects::nonNull)
                 .collect(Collectors.toList());
 
-        Map<String, PersistentUserSessionAdapter> sessionsById = result.stream()
+        Map<String, PersistentUserSessionAdapter> sessionsById = userSessionAdapters.stream()
                 .collect(Collectors.toMap(UserSessionModel::getId, Function.identity()));
 
-        Set<String> userSessionIds = sessionsById.keySet();
-
         Set<String> removedClientUUIDs = new HashSet<>();
 
-        if (!userSessionIds.isEmpty()) {
-            TypedQuery<PersistentClientSessionEntity> query2 = em.createNamedQuery("findClientSessionsByUserSessions", PersistentClientSessionEntity.class);
-            query2.setParameter("userSessionIds", userSessionIds);
-            query2.setParameter("offline", offlineStr);
-            closing(query2.getResultStream()).forEach(clientSession -> {
+        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);
+
+            closing(queryClientSessions.getResultStream()).forEach(clientSession -> {
                 PersistentUserSessionAdapter userSession = sessionsById.get(clientSession.getUserSessionId());
 
                 PersistentAuthenticatedClientSessionAdapter clientSessAdapter = toAdapter(userSession.getRealm(), userSession, clientSession);
@@ -287,7 +289,7 @@ public Stream<UserSessionModel> loadUserSessionsStream(Integer firstResult, Inte
             onClientRemoved(clientUUID);
         }
 
-        return result.stream().map(UserSessionModel.class::cast);
+        return userSessionAdapters.stream().map(UserSessionModel.class::cast);
     }
 
     private PersistentUserSessionAdapter toAdapter(PersistentUserSessionEntity entity) {

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
@@ -38,7 +38,7 @@
         @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="findClientSessionsByUserSessions", query="select sess from PersistentClientSessionEntity sess where sess.offline = :offline and sess.userSessionId IN (:userSessionIds) 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")
 })
 @Table(name="OFFLINE_CLIENT_SESSION")
 @Entity

diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/session/PersistentUserSessionEntity.java b/model/jpa/src/main/java/org/keycloak/models/jpa/session/PersistentUserSessionEntity.java
@@ -38,9 +38,9 @@
         @NamedQuery(name="updateUserSessionLastSessionRefresh", query="update PersistentUserSessionEntity sess set lastSessionRefresh = :lastSessionRefresh where sess.realmId = :realmId" +
                 " AND sess.offline = :offline AND sess.userSessionId IN (:userSessionIds)"),
         @NamedQuery(name="findUserSessionsCount", query="select count(sess) from PersistentUserSessionEntity sess where sess.offline = :offline"),
-        @NamedQuery(name="findUserSessions", query="select sess from PersistentUserSessionEntity sess, RealmEntity realm where realm.id = sess.realmId AND sess.offline = :offline" +
-                " AND (sess.createdOn > :lastCreatedOn OR (sess.createdOn = :lastCreatedOn AND sess.userSessionId > :lastSessionId))" +
-                " order by sess.createdOn,sess.userSessionId")
+        @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")
 
 })
 @Table(name="OFFLINE_USER_SESSION")

diff --git a/...ivate/src/main/java/org/keycloak/models/session/DisabledUserSessionPersisterProvider.java b/...ivate/src/main/java/org/keycloak/models/session/DisabledUserSessionPersisterProvider.java
@@ -111,7 +111,7 @@ public void removeExpired(RealmModel realm) {
 
     @Override
     public Stream<UserSessionModel> loadUserSessionsStream(Integer firstResult, Integer maxResults, boolean offline,
-                                                           Integer lastCreatedOn, String lastUserSessionId) {
+                                                           String lastUserSessionId) {
         return Stream.empty();
     }
 

diff --git a/...r-spi-private/src/main/java/org/keycloak/models/session/UserSessionPersisterProvider.java b/...r-spi-private/src/main/java/org/keycloak/models/session/UserSessionPersisterProvider.java
@@ -58,25 +58,24 @@ public interface UserSessionPersisterProvider extends Provider {
 
     /**
      * Called during startup. For each userSession, it loads also clientSessions
-     * @deprecated Use {@link #loadUserSessionsStream(Integer, Integer, boolean, Integer, String) loadUserSessionsStream} instead.
+     * @deprecated Use {@link #loadUserSessionsStream(Integer, Integer, boolean, String) loadUserSessionsStream} instead.
      */
     @Deprecated
     default List<UserSessionModel> loadUserSessions(int firstResult, int maxResults, boolean offline, int lastCreatedOn, String lastUserSessionId) {
-        return loadUserSessionsStream(firstResult, maxResults, offline, lastCreatedOn, lastUserSessionId).collect(Collectors.toList());
+        return loadUserSessionsStream(firstResult, maxResults, offline, lastUserSessionId).collect(Collectors.toList());
     }
 
     /**
      * Called during startup. For each userSession, it loads also clientSessions.
      * @param firstResult {@code Integer} Index of the first desired user session. Ignored if negative or {@code null}.
      * @param maxResults {@code Integer} Maximum number of returned user sessions. Ignored if negative or {@code null}.
      * @param offline {@code boolean} Flag to include offline sessions.
-     * @param lastCreatedOn {@code Integer} Timestamp when the user session was created. It will return only user sessions created later.
-     * @param lastUserSessionId {@code String} Id of the user session. In case of equal {@code lastCreatedOn}
+     * @param lastUserSessionId {@code String} Id of the user session. It will return only user sessions with id's lexicographically greater than this.
      * it will compare the id in dictionary order and takes only those created later.
      * @return Stream of {@link UserSessionModel}. Never returns {@code null}.
      */
     Stream<UserSessionModel> loadUserSessionsStream(Integer firstResult, Integer maxResults, boolean offline,
-                                                    Integer lastCreatedOn, String lastUserSessionId);
+                                                    String lastUserSessionId);
 
     int getUserSessionsCount(boolean offline);