keycloak-aplcache

diff --git a/common/src/main/java/org/keycloak/common/util/Environment.java b/common/src/main/java/org/keycloak/common/util/Environment.java
index 94a4d45..0b0fb13 100644
--- a/common/src/main/java/org/keycloak/common/util/Environment.java
+++ b/common/src/main/java/org/keycloak/common/util/Environment.java
@@ -24,4 +24,15 @@ public class Environment {
 
     public static final boolean IS_IBM_JAVA = System.getProperty("java.vendor").contains("IBM");
 
+    public static final int DEFAULT_JBOSS_AS_STARTUP_TIMEOUT = 300;
+
+    public static int getServerStartupTimeout() {
+        String timeout = System.getProperty("jboss.as.management.blocking.timeout");
+        if (timeout != null) {
+            return Integer.parseInt(timeout);
+        } else {
+            return DEFAULT_JBOSS_AS_STARTUP_TIMEOUT;
+        }
+    }
+
 }

diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanAuthenticationSessionProviderFactory.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanAuthenticationSessionProviderFactory.java
index 72891fa..e580bbd 100644
--- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanAuthenticationSessionProviderFactory.java
+++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanAuthenticationSessionProviderFactory.java
@@ -31,6 +31,7 @@ import org.keycloak.models.sessions.infinispan.events.AbstractAuthSessionCluster
 import org.keycloak.models.sessions.infinispan.events.ClientRemovedSessionEvent;
 import org.keycloak.models.sessions.infinispan.events.RealmRemovedSessionEvent;
 import org.keycloak.models.sessions.infinispan.util.InfinispanKeyGenerator;
+import org.keycloak.models.utils.KeycloakModelUtils;
 import org.keycloak.models.utils.PostMigrationEvent;
 import org.keycloak.provider.ProviderEvent;
 import org.keycloak.provider.ProviderEventListener;
@@ -73,7 +74,12 @@ public class InfinispanAuthenticationSessionProviderFactory implements Authentic
             @Override
             public void onEvent(ProviderEvent event) {
                 if (event instanceof PostMigrationEvent) {
-                    registerClusterListeners(((PostMigrationEvent) event).getSession());
+
+                    KeycloakModelUtils.runJobInTransaction(factory, (KeycloakSession session) -> {
+
+                        registerClusterListeners(session);
+
+                    });
                 }
             }
         });

diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProviderFactory.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProviderFactory.java
index dc6e261..9e47c41 100755
--- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProviderFactory.java
+++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProviderFactory.java
@@ -23,6 +23,7 @@ import org.infinispan.persistence.remote.RemoteStore;
 import org.jboss.logging.Logger;
 import org.keycloak.Config;
 import org.keycloak.cluster.ClusterProvider;
+import org.keycloak.common.util.Environment;
 import org.keycloak.connections.infinispan.InfinispanConnectionProvider;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
@@ -109,13 +110,19 @@ public class InfinispanUserSessionProviderFactory implements UserSessionProvider
             @Override
             public void onEvent(ProviderEvent event) {
                 if (event instanceof PostMigrationEvent) {
-                    KeycloakSession session = ((PostMigrationEvent) event).getSession();
 
-                    keyGenerator = new InfinispanKeyGenerator();
-                    checkRemoteCaches(session);
-                    loadPersistentSessions(factory, getMaxErrors(), getSessionsPerSegment());
-                    registerClusterListeners(session);
-                    loadSessionsFromRemoteCaches(session);
+                    int preloadTransactionTimeout = getTimeoutForPreloadingSessionsSeconds();
+                    log.debugf("Will preload sessions with transaction timeout %d seconds", preloadTransactionTimeout);
+
+                    KeycloakModelUtils.runJobInTransactionWithTimeout(factory, (KeycloakSession session) -> {
+
+                        keyGenerator = new InfinispanKeyGenerator();
+                        checkRemoteCaches(session);
+                        loadPersistentSessions(factory, getMaxErrors(), getSessionsPerSegment());
+                        registerClusterListeners(session);
+                        loadSessionsFromRemoteCaches(session);
+
+                    }, preloadTransactionTimeout);
 
                 } else if (event instanceof UserModel.UserRemovedEvent) {
                     UserModel.UserRemovedEvent userRemovedEvent = (UserModel.UserRemovedEvent) event;
@@ -137,6 +144,11 @@ public class InfinispanUserSessionProviderFactory implements UserSessionProvider
         return config.getInt("sessionsPerSegment", 100);
     }
 
+    private int getTimeoutForPreloadingSessionsSeconds() {
+        Integer timeout = config.getInt("sessionsPreloadTimeoutInSeconds", null);
+        return timeout != null ? timeout : Environment.getServerStartupTimeout();
+    }
+
 
     @Override
     public void loadPersistentSessions(final KeycloakSessionFactory sessionFactory, final int maxErrors, final int sessionsPerSegment) {

diff --git a/server-spi-private/src/main/java/org/keycloak/models/utils/KeycloakModelUtils.java b/server-spi-private/src/main/java/org/keycloak/models/utils/KeycloakModelUtils.java
index 96d0cb3..3771363 100755
--- a/server-spi-private/src/main/java/org/keycloak/models/utils/KeycloakModelUtils.java
+++ b/server-spi-private/src/main/java/org/keycloak/models/utils/KeycloakModelUtils.java
@@ -244,6 +244,45 @@ public final class KeycloakModelUtils {
     }
 
 
+    /**
+     * Wrap given runnable job into KeycloakTransaction. Set custom timeout for the JTA transaction (in case we're in the environment with JTA enabled)
+     *
+     * @param factory
+     * @param task
+     * @param timeoutInSeconds
+     */
+    public static void runJobInTransactionWithTimeout(KeycloakSessionFactory factory, KeycloakSessionTask task, int timeoutInSeconds) {
+        JtaTransactionManagerLookup lookup = (JtaTransactionManagerLookup)factory.getProviderFactory(JtaTransactionManagerLookup.class);
+        try {
+            if (lookup != null) {
+                if (lookup.getTransactionManager() != null) {
+                    try {
+                        lookup.getTransactionManager().setTransactionTimeout(timeoutInSeconds);
+                    } catch (SystemException e) {
+                        throw new RuntimeException(e);
+                    }
+                }
+            }
+
+            runJobInTransaction(factory, task);
+
+        } finally {
+            if (lookup != null) {
+                if (lookup.getTransactionManager() != null) {
+                    try {
+                        // Reset to default transaction timeout
+                        lookup.getTransactionManager().setTransactionTimeout(0);
+                    } catch (SystemException e) {
+                        // Shouldn't happen for Wildfly transaction manager
+                        throw new RuntimeException(e);
+                    }
+                }
+            }
+        }
+
+    }
+
+
     public static String getMasterRealmAdminApplicationClientId(String realmName) {
         return realmName + "-realm";
     }

diff --git a/server-spi-private/src/main/java/org/keycloak/models/utils/PostMigrationEvent.java b/server-spi-private/src/main/java/org/keycloak/models/utils/PostMigrationEvent.java
index 430d895..73889e0 100644
--- a/server-spi-private/src/main/java/org/keycloak/models/utils/PostMigrationEvent.java
+++ b/server-spi-private/src/main/java/org/keycloak/models/utils/PostMigrationEvent.java
@@ -17,7 +17,6 @@
 
 package org.keycloak.models.utils;
 
-import org.keycloak.models.KeycloakSession;
 import org.keycloak.provider.ProviderEvent;
 
 /**
@@ -26,14 +25,4 @@ import org.keycloak.provider.ProviderEvent;
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 public class PostMigrationEvent implements ProviderEvent {
-
-    private final KeycloakSession session;
-
-    public PostMigrationEvent(KeycloakSession session) {
-        this.session = session;
-    }
-
-    public KeycloakSession getSession() {
-        return session;
-    }
 }

diff --git a/services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java b/services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java
index f100148..cd6daf5 100644
--- a/services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java
+++ b/services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java
@@ -162,12 +162,12 @@ public class KeycloakApplication extends Application {
                 public void run(KeycloakSession session) {
                     boolean shouldBootstrapAdmin = new ApplianceBootstrap(session).isNoMasterUser();
                     bootstrapAdminUser.set(shouldBootstrapAdmin);
-
-                    sessionFactory.publish(new PostMigrationEvent(session));
                 }
 
             });
 
+            sessionFactory.publish(new PostMigrationEvent());
+
             singletons.add(new WelcomeResource(bootstrapAdminUser.get()));
 
             setupScheduledTasks(sessionFactory);