keycloak-memoizeit

diff --git a/examples/authz/photoz/photoz-restful-api-authz-service.json b/examples/authz/photoz/photoz-restful-api-authz-service.json
index ff9ee9c..e8d8862 100644
--- a/examples/authz/photoz/photoz-restful-api-authz-service.json
+++ b/examples/authz/photoz/photoz-restful-api-authz-service.json
@@ -45,7 +45,7 @@
       "description": "Defines that only the resource owner is allowed to do something",
       "type": "drools",
       "config": {
-        "mavenArtifactVersion": "2.0.0.CR1-SNAPSHOT",
+        "mavenArtifactVersion": "2.1.0-SNAPSHOT",
         "mavenArtifactId": "photoz-authz-policy",
         "sessionName": "MainOwnerSession",
         "mavenArtifactGroupId": "org.keycloak",

diff --git a/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedPolicyStore.java b/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedPolicyStore.java
index 02800ce..10c108d 100644
--- a/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedPolicyStore.java
+++ b/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedPolicyStore.java
@@ -29,7 +29,6 @@ import org.keycloak.connections.infinispan.InfinispanConnectionProvider;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.authorization.infinispan.InfinispanStoreFactoryProvider.CacheTransaction;
 import org.keycloak.models.authorization.infinispan.entities.CachedPolicy;
-import org.keycloak.models.entities.AbstractIdentifiableEntity;
 import org.keycloak.representations.idm.authorization.DecisionStrategy;
 import org.keycloak.representations.idm.authorization.Logic;
 
@@ -64,13 +63,15 @@ public class CachedPolicyStore implements PolicyStore {
     public Policy create(String name, String type, ResourceServer resourceServer) {
         Policy policy = getDelegate().create(name, type, getStoreFactory().getResourceServerStore().findById(resourceServer.getId()));
 
+        this.transaction.whenRollback(() -> cache.remove(getCacheKeyForPolicy(policy.getId())));
+
         return createAdapter(new CachedPolicy(policy));
     }
 
     @Override
     public void delete(String id) {
         getDelegate().delete(id);
-        this.transaction.whenComplete(() -> cache.remove(getCacheKeyForPolicy(id)));
+        this.transaction.whenCommit(() -> cache.remove(getCacheKeyForPolicy(id)));
     }
 
     @Override
@@ -387,7 +388,7 @@ public class CachedPolicyStore implements PolicyStore {
                 if (this.updated == null) {
                     this.updated = getDelegate().findById(getId());
                     if (this.updated == null) throw new IllegalStateException("Not found in database");
-                    transaction.whenComplete(() -> cache.evict(getCacheKeyForPolicy(getId())));
+                    transaction.whenCommit(() -> cache.evict(getCacheKeyForPolicy(getId())));
                 }
 
                 return this.updated;

diff --git a/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedResourceServerStore.java b/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedResourceServerStore.java
index e03f3a7..33bed6a 100644
--- a/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedResourceServerStore.java
+++ b/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedResourceServerStore.java
@@ -56,13 +56,15 @@ public class CachedResourceServerStore implements ResourceServerStore {
     public ResourceServer create(String clientId) {
         ResourceServer resourceServer = getDelegate().create(clientId);
 
+        this.transaction.whenRollback(() -> cache.remove(getCacheKeyForResourceServer(resourceServer.getId())));
+
         return createAdapter(new CachedResourceServer(resourceServer));
     }
 
     @Override
     public void delete(String id) {
         getDelegate().delete(id);
-        this.transaction.whenComplete(() -> this.cache.remove(getCacheKeyForResourceServer(id)));
+        this.transaction.whenCommit(() -> this.cache.remove(getCacheKeyForResourceServer(id)));
     }
 
     @Override
@@ -167,7 +169,7 @@ public class CachedResourceServerStore implements ResourceServerStore {
                 if (this.updated == null) {
                     this.updated = getDelegate().findById(getId());
                     if (this.updated == null) throw new IllegalStateException("Not found in database");
-                    transaction.whenComplete(() -> cache.evict(getCacheKeyForResourceServer(getId())));
+                    transaction.whenCommit(() -> cache.evict(getCacheKeyForResourceServer(getId())));
                 }
 
                 return this.updated;

diff --git a/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedResourceStore.java b/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedResourceStore.java
index aa587f5..ee638f6 100644
--- a/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedResourceStore.java
+++ b/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedResourceStore.java
@@ -30,6 +30,7 @@ import org.keycloak.models.authorization.infinispan.InfinispanStoreFactoryProvid
 import org.keycloak.models.authorization.infinispan.entities.CachedResource;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map.Entry;
 import java.util.Set;
@@ -41,6 +42,7 @@ import java.util.stream.Collectors;
 public class CachedResourceStore implements ResourceStore {
 
     private static final String RESOURCE_ID_CACHE_PREFIX = "rsc-id-";
+    private static final String RESOURCE_OWNER_CACHE_PREFIX = "rsc-owner-";
 
     private final KeycloakSession session;
     private final CacheTransaction transaction;
@@ -59,6 +61,8 @@ public class CachedResourceStore implements ResourceStore {
     public Resource create(String name, ResourceServer resourceServer, String owner) {
         Resource resource = getDelegate().create(name, getStoreFactory().getResourceServerStore().findById(resourceServer.getId()), owner);
 
+        this.transaction.whenRollback(() -> cache.remove(getCacheKeyForResource(resource.getId())));
+
         return createAdapter(new CachedResource(resource));
     }
 
@@ -77,6 +81,7 @@ public class CachedResourceStore implements ResourceStore {
             Resource resource = getDelegate().findById(id);
 
             if (resource != null) {
+                updateCachedIds(getResourceOwnerCacheKey(resource.getOwner()), resource, false);
                 return createAdapter(updateResourceCache(resource));
             }
 
@@ -88,26 +93,18 @@ public class CachedResourceStore implements ResourceStore {
 
     @Override
     public List<Resource> findByOwner(String ownerId) {
-        List<Resource> cache = new ArrayList<>();
-
-        for (Entry entry : this.cache.entrySet()) {
-            String cacheKey = (String) entry.getKey();
+        List<String> cachedIds = this.cache.get(getResourceOwnerCacheKey(ownerId));
 
-            if (cacheKey.startsWith(RESOURCE_ID_CACHE_PREFIX)) {
-                List<Resource> value = (List<Resource>) entry.getValue();
-                Resource resource = value.get(0);
-
-                if (resource.getOwner().equals(ownerId)) {
-                    cache.add(findById(resource.getId()));
-                }
+        if (cachedIds == null) {
+            for (Resource resource : getDelegate().findByOwner(ownerId)) {
+                updateCachedIds(getResourceOwnerCacheKey(ownerId), resource, true);
             }
+            cachedIds = this.cache.getOrDefault(getResourceOwnerCacheKey(ownerId), Collections.emptyList());
         }
 
-        if (cache.isEmpty()) {
-            getDelegate().findByOwner(ownerId).forEach(resource -> cache.add(findById(updateResourceCache(resource).getId())));
-        }
-
-        return cache;
+        return  ((List<String>) this.cache.getOrDefault(getResourceOwnerCacheKey(ownerId), Collections.emptyList())).stream().map(this::findById)
+                        .filter(resource -> resource != null)
+                        .collect(Collectors.toList());
     }
 
     @Override
@@ -146,26 +143,7 @@ public class CachedResourceStore implements ResourceStore {
 
     @Override
     public List<Resource> findByType(String type) {
-        List<Resource> cache = new ArrayList<>();
-
-        for (Entry entry : this.cache.entrySet()) {
-            String cacheKey = (String) entry.getKey();
-
-            if (cacheKey.startsWith(RESOURCE_ID_CACHE_PREFIX)) {
-                List<Resource> value = (List<Resource>) entry.getValue();
-                Resource resource = value.get(0);
-
-                if (resource.getType().equals(type)) {
-                    cache.add(findById(resource.getId()));
-                }
-            }
-        }
-
-        if (cache.isEmpty()) {
-            getDelegate().findByType(type).forEach(resource -> cache.add(findById(updateResourceCache(resource).getId())));
-        }
-
-        return cache;
+        return  getDelegate().findByType(type).stream().map(resource -> findById(resource.getId())).collect(Collectors.toList());
     }
 
     private String getCacheKeyForResource(String id) {
@@ -278,7 +256,7 @@ public class CachedResourceStore implements ResourceStore {
                 if (this.updated == null) {
                     this.updated = getDelegate().findById(getId());
                     if (this.updated == null) throw new IllegalStateException("Not found in database");
-                    transaction.whenComplete(() -> cache.evict(getCacheKeyForResource(getId())));
+                    transaction.whenCommit(() -> cache.evict(getCacheKeyForResource(getId())));
                 }
 
                 return this.updated;
@@ -296,4 +274,24 @@ public class CachedResourceStore implements ResourceStore {
 
         return cached;
     }
+
+    private void updateCachedIds(String cacheKey, Resource resource, boolean create) {
+        List<String> cached = this.cache.get(cacheKey);
+
+        if (cached == null) {
+            if (!create) {
+                return;
+            }
+            cached = new ArrayList<>();
+            this.cache.put(getResourceOwnerCacheKey(resource.getOwner()), cached);
+        }
+
+        if (cached != null && !cached.contains(resource.getId())) {
+            cached.add(resource.getId());
+        }
+    }
+
+    private String getResourceOwnerCacheKey(String ownerId) {
+        return RESOURCE_OWNER_CACHE_PREFIX + ownerId;
+    }
 }

diff --git a/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedScopeStore.java b/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedScopeStore.java
index 5912645..af491c2 100644
--- a/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedScopeStore.java
+++ b/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/CachedScopeStore.java
@@ -56,13 +56,15 @@ public class CachedScopeStore implements ScopeStore {
     public Scope create(String name, ResourceServer resourceServer) {
         Scope scope = getDelegate().create(name, getStoreFactory().getResourceServerStore().findById(resourceServer.getId()));
 
+        this.transaction.whenRollback(() -> cache.remove(getCacheKeyForScope(scope.getId())));
+
         return createAdapter(new CachedScope(scope));
     }
 
     @Override
     public void delete(String id) {
         getDelegate().delete(id);
-        this.transaction.whenComplete(() -> cache.remove(getCacheKeyForScope(id)));
+        this.transaction.whenCommit(() -> cache.remove(getCacheKeyForScope(id)));
     }
 
     @Override
@@ -173,7 +175,7 @@ public class CachedScopeStore implements ScopeStore {
                 if (this.updated == null) {
                     this.updated = getDelegate().findById(getId());
                     if (this.updated == null) throw new IllegalStateException("Not found in database");
-                    transaction.whenComplete(() -> cache.evict(getCacheKeyForScope(getId())));
+                    transaction.whenCommit(() -> cache.evict(getCacheKeyForScope(getId())));
                 }
 
                 return this.updated;
@@ -188,7 +190,7 @@ public class CachedScopeStore implements ScopeStore {
 
         cache.add(cached);
 
-        this.transaction.whenComplete(() -> this.cache.put(getCacheKeyForScope(scope.getId()), cache));
+        this.transaction.whenCommit(() -> this.cache.put(getCacheKeyForScope(scope.getId()), cache));
 
         return cached;
     }

diff --git a/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/InfinispanStoreFactoryProvider.java b/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/InfinispanStoreFactoryProvider.java
index df9f262..cea4388 100644
--- a/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/InfinispanStoreFactoryProvider.java
+++ b/model/infinispan/src/main/java/org/keycloak/models/authorization/infinispan/InfinispanStoreFactoryProvider.java
@@ -71,6 +71,7 @@ public class InfinispanStoreFactoryProvider implements CachedStoreFactoryProvide
     static class CacheTransaction implements KeycloakTransaction {
 
         private List<Runnable> completeTasks = new ArrayList<>();
+        private List<Runnable> rollbackTasks = new ArrayList<>();
 
         @Override
         public void begin() {
@@ -84,7 +85,7 @@ public class InfinispanStoreFactoryProvider implements CachedStoreFactoryProvide
 
         @Override
         public void rollback() {
-            this.completeTasks.forEach(task -> task.run());
+            this.rollbackTasks.forEach(task -> task.run());
         }
 
         @Override
@@ -102,8 +103,12 @@ public class InfinispanStoreFactoryProvider implements CachedStoreFactoryProvide
             return false;
         }
 
-        protected void whenComplete(Runnable task) {
+        protected void whenCommit(Runnable task) {
             this.completeTasks.add(task);
         }
+
+        protected void whenRollback(Runnable task) {
+            this.rollbackTasks.add(task);
+        }
     }
 }

diff --git a/testsuite/integration-arquillian/test-apps/photoz/photoz-restful-api-authz-service.json b/testsuite/integration-arquillian/test-apps/photoz/photoz-restful-api-authz-service.json
index 91c3e0e..cc8c8f8 100644
--- a/testsuite/integration-arquillian/test-apps/photoz/photoz-restful-api-authz-service.json
+++ b/testsuite/integration-arquillian/test-apps/photoz/photoz-restful-api-authz-service.json
@@ -45,7 +45,7 @@
       "description": "Defines that only the resource owner is allowed to do something",
       "type": "drools",
       "config": {
-        "mavenArtifactVersion": "2.0.0.CR1-SNAPSHOT",
+        "mavenArtifactVersion": "2.1.0-SNAPSHOT",
         "mavenArtifactId": "photoz-authz-policy",
         "sessionName": "MainOwnerSession",
         "mavenArtifactGroupId": "org.keycloak.testsuite",