keycloak-aplcache

diff --git a/model/infinispan/src/main/java/org/keycloak/cluster/infinispan/InfinispanNotificationsManager.java b/model/infinispan/src/main/java/org/keycloak/cluster/infinispan/InfinispanNotificationsManager.java
index 847e144..4cd8cf1 100644
--- a/model/infinispan/src/main/java/org/keycloak/cluster/infinispan/InfinispanNotificationsManager.java
+++ b/model/infinispan/src/main/java/org/keycloak/cluster/infinispan/InfinispanNotificationsManager.java
@@ -52,7 +52,6 @@ import org.keycloak.cluster.ClusterProvider;
 import org.keycloak.common.util.ConcurrentMultivaluedHashMap;
 import org.keycloak.executors.ExecutorsProvider;
 import org.keycloak.models.KeycloakSession;
-import org.keycloak.models.sessions.infinispan.util.InfinispanUtil;
 
 /**
  * Impl for sending infinispan messages across cluster and listening to them
@@ -154,8 +153,7 @@ public class InfinispanNotificationsManager {
                     .put(eventKey, wrappedEvent, 120, TimeUnit.SECONDS);
         } else {
             // Add directly to remoteCache. Will notify remote listeners on all nodes in all DCs
-            RemoteCache remoteCache = InfinispanUtil.getRemoteCache(workCache);
-            remoteCache.put(eventKey, wrappedEvent, 120, TimeUnit.SECONDS);
+            workRemoteCache.put(eventKey, wrappedEvent, 120, TimeUnit.SECONDS);
         }
     }
 

diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remotestore/RemoteCacheSessionsLoader.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remotestore/RemoteCacheSessionsLoader.java
index 1201563..4ecd445 100644
--- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remotestore/RemoteCacheSessionsLoader.java
+++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/remotestore/RemoteCacheSessionsLoader.java
@@ -42,30 +42,40 @@ public class RemoteCacheSessionsLoader implements SessionLoader {
     private static final Logger log = Logger.getLogger(RemoteCacheSessionsLoader.class);
 
 
-    // Javascript to be executed on remote infinispan server (Flag CACHE_MODE_LOCAL assumes that remoteCache is replicated)
+    // Javascript to be executed on remote infinispan server.
+    // Flag CACHE_MODE_LOCAL is optimization used just when remoteCache is replicated as all the entries are available locally. For distributed caches, it can't be used
     private static final String REMOTE_SCRIPT_FOR_LOAD_SESSIONS =
             "function loadSessions() {" +
-            "  var flagClazz = cache.getClass().getClassLoader().loadClass(\"org.infinispan.context.Flag\"); \n" +
-            "  var localFlag = java.lang.Enum.valueOf(flagClazz, \"CACHE_MODE_LOCAL\"); \n" +
-            "  var cacheStream = cache.getAdvancedCache().withFlags([ localFlag ]).entrySet().stream();\n" +
-            "  var result = cacheStream.skip(first).limit(max).collect(java.util.stream.Collectors.toMap(\n" +
-            "    new java.util.function.Function() {\n" +
-            "      apply: function(entry) {\n" +
-            "        return entry.getKey();\n" +
-            "      }\n" +
-            "    },\n" +
-            "    new java.util.function.Function() {\n" +
-            "      apply: function(entry) {\n" +
-            "        return entry.getValue();\n" +
-            "      }\n" +
-            "    }\n" +
-            "  ));\n" +
-            "\n" +
-            "  cacheStream.close();\n" +
-            "  return result;\n" +
-            "};\n" +
-            "\n" +
-            "loadSessions();";
+                    "  var flagClazz = cache.getClass().getClassLoader().loadClass(\"org.infinispan.context.Flag\"); \n" +
+                    "  var localFlag = java.lang.Enum.valueOf(flagClazz, \"CACHE_MODE_LOCAL\"); \n" +
+                    "  var cacheMode = cache.getCacheConfiguration().clustering().cacheMode(); \n" +
+                    "  var canUseLocalFlag = !cacheMode.isClustered() || cacheMode.isReplicated(); \n" +
+
+                    "  var cacheStream; \n" +
+                    "  if (canUseLocalFlag) { \n" +
+                    "      cacheStream = cache.getAdvancedCache().withFlags([ localFlag ]).entrySet().stream();\n" +
+                    "  } else { \n" +
+                    "      cacheStream = cache.getAdvancedCache().withFlags([ ]).entrySet().stream();\n" +
+                    "  }; \n" +
+
+                    "  var result = cacheStream.skip(first).limit(max).collect(java.util.stream.Collectors.toMap(\n" +
+                    "    new java.util.function.Function() {\n" +
+                    "      apply: function(entry) {\n" +
+                    "        return entry.getKey();\n" +
+                    "      }\n" +
+                    "    },\n" +
+                    "    new java.util.function.Function() {\n" +
+                    "      apply: function(entry) {\n" +
+                    "        return entry.getValue();\n" +
+                    "      }\n" +
+                    "    }\n" +
+                    "  ));\n" +
+                    "\n" +
+                    "  cacheStream.close();\n" +
+                    "  return result;\n" +
+                    "};\n" +
+                    "\n" +
+                    "loadSessions();";