From 442e0535bd84acd9bccf0ec51968b0a2706f8d7b Mon Sep 17 00:00:00 2001
From: Pedro Ruivo <pruivo@redhat.com>
Date: Wed, 8 May 2024 10:54:29 +0100
Subject: [PATCH] Use cache.compute() method to improve the replace retry loop

Final commit to enable the new function.

Closes #29073

Signed-off-by: Pedro Ruivo <pruivo@redhat.com>
---
 .../InfinispanChangelogBasedTransaction.java  | 51 +++++++------------
 1 file changed, 19 insertions(+), 32 deletions(-)

diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/InfinispanChangelogBasedTransaction.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/InfinispanChangelogBasedTransaction.java
index 2df4073a2f6..b60c14093dc 100644
--- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/InfinispanChangelogBasedTransaction.java
+++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/changes/InfinispanChangelogBasedTransaction.java
@@ -223,45 +223,32 @@ private void runOperationInCluster(K key, MergedUpdate<V> task,  SessionEntityWr
 
 
     private void replace(K key, MergedUpdate<V> task, SessionEntityWrapper<V> oldVersionEntity, long lifespanMs, long maxIdleTimeMs) {
-        boolean replaced = false;
+        SessionEntityWrapper<V> oldVersion = oldVersionEntity;
+        SessionEntityWrapper<V> returnValue = null;
         int iteration = 0;
-        V session = oldVersionEntity.getEntity();
-
-        while (!replaced && iteration < InfinispanUtil.MAXIMUM_REPLACE_RETRIES) {
-            iteration++;
-
-            SessionEntityWrapper<V> newVersionEntity = generateNewVersionAndWrapEntity(session, oldVersionEntity.getLocalMetadata());
-
-            // Atomic cluster-aware replace
-            replaced = CacheDecorators.skipCacheStoreIfRemoteCacheIsEnabled(cache).replace(key, oldVersionEntity, newVersionEntity, lifespanMs, TimeUnit.MILLISECONDS, maxIdleTimeMs, TimeUnit.MILLISECONDS);
-
-            // Replace fail. Need to load latest entity from cache, apply updates again and try to replace in cache again
-            if (!replaced) {
-                if (logger.isDebugEnabled()) {
-                    logger.debugf("Replace failed for entity: %s, old version %s, new version %s. Will try again", key, oldVersionEntity.getVersion(), newVersionEntity.getVersion());
-                }
-
-                oldVersionEntity = cache.get(key);
-
-                if (oldVersionEntity == null) {
-                    logger.debugf("Entity %s not found. Maybe removed in the meantime. Replace task will be ignored", key);
-                    return;
-                }
-
-                session = oldVersionEntity.getEntity();
+        V session = oldVersion.getEntity();
+        var writeCache = CacheDecorators.skipCacheStoreIfRemoteCacheIsEnabled(cache);
+        while (iteration++ < InfinispanUtil.MAXIMUM_REPLACE_RETRIES) {
+            SessionEntityWrapper<V> newVersionEntity = generateNewVersionAndWrapEntity(session, oldVersion.getLocalMetadata());
+            returnValue = writeCache.computeIfPresent(key, new ReplaceFunction<>(oldVersion.getVersion(), newVersionEntity), lifespanMs, TimeUnit.MILLISECONDS, maxIdleTimeMs, TimeUnit.MILLISECONDS);
+
+            if (returnValue == null) {
+                logger.debugf("Entity %s not found. Maybe removed in the meantime. Replace task will be ignored", key);
+                return;
+            }
 
-                task.runUpdate(session);
-            } else {
+            if (returnValue.getVersion().equals(newVersionEntity.getVersion())) {
                 if (logger.isTraceEnabled()) {
-                    logger.tracef("Replace SUCCESS for entity: %s . old version: %s, new version: %s, Lifespan: %d ms, MaxIdle: %d ms", key, oldVersionEntity.getVersion(), newVersionEntity.getVersion(), task.getLifespanMs(), task.getMaxIdleTimeMs());
+                    logger.tracef("Replace SUCCESS for entity: %s . old version: %s, new version: %s, Lifespan: %d ms, MaxIdle: %d ms", key, oldVersion.getVersion(), newVersionEntity.getVersion(), task.getLifespanMs(), task.getMaxIdleTimeMs());
                 }
+                return;
             }
-        }
 
-        if (!replaced) {
-            logger.warnf("Failed to replace entity '%s' in cache '%s'", key, cache.getName());
+            oldVersion = returnValue;
+            session = oldVersion.getEntity();
+            task.runUpdate(session);
         }
-
+        logger.warnf("Failed to replace entity '%s' in cache '%s'. Expected: %s, Current: %s", key, cache.getName(), oldVersion, returnValue);
     }