From 04445c8a23a86edb4756588fe3e540c1ac362576 Mon Sep 17 00:00:00 2001
From: mposolda <mposolda@gmail.com>
Date: Mon, 10 Dec 2018 14:38:12 +0100
Subject: [PATCH] KEYCLOAK-8904 Backpressure in RemoteCacheSessionsLoader

---
 .../infinispan/remotestore/RemoteCacheSessionsLoader.java   | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

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 2e927df9b3..730bd6ab3e 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
@@ -18,6 +18,7 @@
 package org.keycloak.models.sessions.infinispan.remotestore;
 
 import java.io.Serializable;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -114,6 +115,7 @@ public class RemoteCacheSessionsLoader implements SessionLoader<RemoteCacheSessi
 
         log.debugf("Will do bulk load of sessions from remote cache '%s' . Segment: %d", cache.getName(), ctx.getSegment());
 
+        Map<Object, Object> remoteEntries = new HashMap<>();
         CloseableIterator<Map.Entry> iterator = null;
         int countLoaded = 0;
         try {
@@ -121,7 +123,7 @@ public class RemoteCacheSessionsLoader implements SessionLoader<RemoteCacheSessi
             while (iterator.hasNext()) {
                 countLoaded++;
                 Map.Entry entry = iterator.next();
-                decoratedCache.putAsync(entry.getKey(), entry.getValue());
+                remoteEntries.put(entry.getKey(), entry.getValue());
             }
         } catch (RuntimeException e) {
             log.warnf(e, "Error loading sessions from remote cache '%s' for segment '%d'", remoteCache.getName(), ctx.getSegment());
@@ -132,6 +134,8 @@ public class RemoteCacheSessionsLoader implements SessionLoader<RemoteCacheSessi
             }
         }
 
+        decoratedCache.putAll(remoteEntries);
+
         log.debugf("Successfully finished loading sessions from cache '%s' . Segment: %d, Count of sessions loaded: %d", cache.getName(), ctx.getSegment(), countLoaded);
 
         return new WorkerResult(true, ctx.getSegment(), ctx.getWorkerId());