From 37af5fbffed6528f39e34ceb7fe10c32a092b698 Mon Sep 17 00:00:00 2001 From: Martin Kanis Date: Tue, 28 Feb 2023 13:42:08 +0100 Subject: [PATCH] Introduce optimistic locking for HotRod storage Closes #15402 --- model/map-hot-rod/pom.xml | 4 + .../map/storage/hotRod/HotRodMapStorage.java | 80 ++++++---- .../session/UserSessionProviderModelTest.java | 14 -- .../transaction/StorageTransactionTest.java | 148 +++++++++++++++++- 4 files changed, 196 insertions(+), 50 deletions(-) diff --git a/model/map-hot-rod/pom.xml b/model/map-hot-rod/pom.xml index a1b1b25faf..d4c1156431 100644 --- a/model/map-hot-rod/pom.xml +++ b/model/map-hot-rod/pom.xml @@ -62,6 +62,10 @@ ${project.version} true + + jakarta.persistence + jakarta.persistence-api + diff --git a/model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/HotRodMapStorage.java b/model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/HotRodMapStorage.java index df7ee9e0dd..d2a7cf1929 100644 --- a/model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/HotRodMapStorage.java +++ b/model/map-hot-rod/src/main/java/org/keycloak/models/map/storage/hotRod/HotRodMapStorage.java @@ -17,6 +17,7 @@ package org.keycloak.models.map.storage.hotRod; +import org.infinispan.client.hotrod.MetadataValue; import org.infinispan.client.hotrod.RemoteCache; import org.infinispan.client.hotrod.Search; import org.infinispan.commons.util.CloseableIterator; @@ -26,6 +27,7 @@ import org.jboss.logging.Logger; import org.keycloak.common.util.Time; import org.keycloak.models.AbstractKeycloakTransaction; import org.keycloak.models.KeycloakSession; +import org.keycloak.models.UserSessionModel; import org.keycloak.models.map.common.AbstractEntity; import org.keycloak.models.map.common.DeepCloner; import org.keycloak.models.map.common.ExpirableEntity; @@ -50,7 +52,9 @@ import org.keycloak.models.map.storage.hotRod.transaction.NoActionHotRodTransact import org.keycloak.storage.SearchableModelField; import org.keycloak.utils.LockObjectsForModification; +import javax.persistence.OptimisticLockException; import java.time.Duration; +import java.util.HashMap; import java.util.Map; import java.util.Objects; import java.util.Spliterators; @@ -79,6 +83,7 @@ public class HotRodMapStorage, MapModelCriteriaBuilder.UpdatePredicatesFunc> fieldPredicates; private final Long lockTimeout; private final RemoteCache locksCache; + private final Map entityVersionCache = new HashMap<>(); public HotRodMapStorage(KeycloakSession session, RemoteCache remoteCache, StringKeyConverter keyConverter, HotRodEntityDescriptor storedEntityDescriptor, DeepCloner cloner, AllAreasHotRodTransactionsWrapper txWrapper, Long lockTimeout) { this.session = session; @@ -146,11 +151,15 @@ public class HotRodMapStorage entityWithMetadata = remoteCache.getWithMetadata(k); + if (entityWithMetadata == null) return null; + + // store entity version + LOG.tracef("Entity %s read in version %s", key, entityWithMetadata.getVersion(), getShortStackTrace()); + entityVersionCache.put(k, entityWithMetadata.getVersion()); // Create delegate that implements Map*Entity - return delegateProducer.apply(hotRodEntity); + return entityWithMetadata.getValue() != null ? delegateProducer.apply(entityWithMetadata.getValue()) : null; } @Override @@ -160,23 +169,37 @@ public class HotRodMapStorage 0) { - previousValue = remoteCache.replace(key, value.getHotRodEntity(), lifespan, TimeUnit.MILLISECONDS); + if (!remoteCache.replaceWithVersion(key, value.getHotRodEntity(), entityVersionCache.get(key), lifespan, TimeUnit.MILLISECONDS, -1, TimeUnit.MILLISECONDS)) { + throw new OptimisticLockException("Entity " + key + " with version " + entityVersionCache.get(key) + " already changed by a different transaction."); + } } else { + if (!remoteCache.removeWithVersion(key, entityVersionCache.get(key))) { + throw new OptimisticLockException("Entity " + key + " with version " + entityVersionCache.get(key) + " already changed by a different transaction."); + } LOG.warnf("Removing entity %s from storage due to negative/zero lifespan.", key); - previousValue = remoteCache.remove(key); } - return previousValue == null ? null : delegateProducer.apply(previousValue); + + return delegateProducer.apply(value.getHotRodEntity()); } } - E previousValue = remoteCache.replace(key, value.getHotRodEntity()); - return previousValue == null ? null : delegateProducer.apply(previousValue); + if (!remoteCache.replaceWithVersion(key, value.getHotRodEntity(), entityVersionCache.get(key))) { + throw new OptimisticLockException("Entity " + key + " with version " + entityVersionCache.get(key) + " already changed by a different transaction."); + } + return delegateProducer.apply(value.getHotRodEntity()); } @Override public boolean delete(String key) { K k = keyConverter.fromStringSafe(key); + + Long entityVersion = entityVersionCache.get(k); + if (entityVersion != null) { + if (!remoteCache.removeWithVersion(k, entityVersion)) { + throw new OptimisticLockException("Entity " + key + " with version " + entityVersion + " already changed by a different transaction."); + } + return true; + } return remoteCache.remove(k) != null; } @@ -190,20 +213,6 @@ public class HotRodMapStorage read(QueryParameters queryParameters) { - if (LockObjectsForModification.isEnabled(session, storedEntityDescriptor.getModelTypeClass())) { - return pessimisticQueryRead(queryParameters); - } - - Query query = prepareQueryWithPrefixAndParameters(null, queryParameters); - CloseableIterator iterator = paginateQuery(query, queryParameters.getOffset(), - queryParameters.getLimit()).iterator(); - return closing(StreamSupport.stream(Spliterators.spliteratorUnknownSize(iterator, 0), false)) - .onClose(iterator::close) - .filter(Objects::nonNull) // see https://github.com/keycloak/keycloak/issues/9271 - .map(this.delegateProducer); - } - - private Stream pessimisticQueryRead(QueryParameters queryParameters) { DefaultModelCriteria dmc = queryParameters.getModelCriteriaBuilder(); // Optimization if the criteria contains only one id @@ -216,17 +225,33 @@ public class HotRodMapStorage read without optimistic locking. + // See https://issues.redhat.com/browse/ISPN-14537 + if (!dmc.isEmpty() && dmc.partiallyEvaluate((field, op, arg) -> field == UserSessionModel.SearchableFields.CLIENT_ID).toString().contains("__TRUE__")) { + Query query = prepareQueryWithPrefixAndParameters(null, queryParameters); + CloseableIterator iterator = paginateQuery(query, queryParameters.getOffset(), + queryParameters.getLimit()).iterator(); + return closing(StreamSupport.stream(Spliterators.spliteratorUnknownSize(iterator, 0), false)) + .onClose(iterator::close) + .filter(Objects::nonNull) // see https://github.com/keycloak/keycloak/issues/9271 + .map(this.delegateProducer); + } + + // Criteria does not contain only one id, we need to read ids without locking and then read entities one by one pessimistically or optimistically Query query = prepareQueryWithPrefixAndParameters("SELECT id ", queryParameters); CloseableIterator iterator = paginateQuery(query, queryParameters.getOffset(), queryParameters.getLimit()).iterator(); @@ -236,11 +261,10 @@ public class HotRodMapStorage a[0]) .map(String.class::cast) - // Pessimistically read + // read by id => this will register the entity in an ISPN transaction .map(this::read) // Entity can be removed in the meanwhile, we need to check for null .filter(Objects::nonNull); - } private Query prepareQueryWithPrefixAndParameters(String prefix, QueryParameters queryParameters) { diff --git a/testsuite/model/src/test/java/org/keycloak/testsuite/model/session/UserSessionProviderModelTest.java b/testsuite/model/src/test/java/org/keycloak/testsuite/model/session/UserSessionProviderModelTest.java index 661f75a486..ff0c2f2d2e 100644 --- a/testsuite/model/src/test/java/org/keycloak/testsuite/model/session/UserSessionProviderModelTest.java +++ b/testsuite/model/src/test/java/org/keycloak/testsuite/model/session/UserSessionProviderModelTest.java @@ -93,20 +93,6 @@ public class UserSessionProviderModelTest extends KeycloakModelTest { @Override public void cleanEnvironment(KeycloakSession s) { - RealmModel realm = s.realms().getRealm(realmId); - s.sessions().removeUserSessions(realm); - - UserModel user1 = s.users().getUserByUsername(realm, "user1"); - UserModel user2 = s.users().getUserByUsername(realm, "user2"); - - UserManager um = new UserManager(s); - if (user1 != null) { - um.removeUser(realm, user1); - } - if (user2 != null) { - um.removeUser(realm, user2); - } - s.realms().removeRealm(realmId); } diff --git a/testsuite/model/src/test/java/org/keycloak/testsuite/model/transaction/StorageTransactionTest.java b/testsuite/model/src/test/java/org/keycloak/testsuite/model/transaction/StorageTransactionTest.java index 6453593d09..73e88837c9 100644 --- a/testsuite/model/src/test/java/org/keycloak/testsuite/model/transaction/StorageTransactionTest.java +++ b/testsuite/model/src/test/java/org/keycloak/testsuite/model/transaction/StorageTransactionTest.java @@ -36,7 +36,9 @@ import org.keycloak.utils.LockObjectsForModification; import javax.persistence.OptimisticLockException; import javax.persistence.PessimisticLockException; +import javax.transaction.RollbackException; +import java.util.UUID; import java.util.function.Function; import static org.hamcrest.CoreMatchers.allOf; @@ -62,7 +64,6 @@ public class StorageTransactionTest extends KeycloakModelTest { RealmModel r = s.realms().createRealm("1"); r.setDefaultRole(s.roles().addRealmRole(r, Constants.DEFAULT_ROLES_ROLE_PREFIX + "-" + r.getName())); r.setAttribute("k1", "v1"); - r.setSsoSessionIdleTimeout(1000); r.setSsoSessionMaxLifespan(2000); @@ -137,14 +138,14 @@ public class StorageTransactionTest extends KeycloakModelTest { @Test // LockObjectForModification currently works only in map-jpa and map-hotrod - @RequireProvider(value = MapStorageProvider.class, only = { JpaMapStorageProviderFactory.PROVIDER_ID, HotRodMapStorageProviderFactory.PROVIDER_ID}) + @RequireProvider(value = MapStorageProvider.class, only = {JpaMapStorageProviderFactory.PROVIDER_ID, HotRodMapStorageProviderFactory.PROVIDER_ID}) public void testLockObjectForModificationById() throws Exception { testLockObjectForModification(session -> LockObjectsForModification.lockRealmsForModification(session, () -> session.realms().getRealm(realmId))); } @Test // LockObjectForModification currently works only in map-jpa and map-hotrod - @RequireProvider(value = MapStorageProvider.class, only = { JpaMapStorageProviderFactory.PROVIDER_ID, HotRodMapStorageProviderFactory.PROVIDER_ID}) + @RequireProvider(value = MapStorageProvider.class, only = {JpaMapStorageProviderFactory.PROVIDER_ID, HotRodMapStorageProviderFactory.PROVIDER_ID}) public void testLockUserSessionForModificationByQuery() throws Exception { // Create user session final String sessionId = withRealm(realmId, (session, realm) -> { @@ -200,9 +201,10 @@ public class StorageTransactionTest extends KeycloakModelTest { } @Test - // Optimistic locking works only with map-jpa - @RequireProvider(value = MapStorageProvider.class, only = JpaMapStorageProviderFactory.PROVIDER_ID) - public void testOptimisticLockingException() throws Exception { + // Optimistic locking works only with map-jpa and map-hotrod + @RequireProvider(value = MapStorageProvider.class, only = {JpaMapStorageProviderFactory.PROVIDER_ID, + HotRodMapStorageProviderFactory.PROVIDER_ID}) + public void testOptimisticLockingExceptionReadById() throws Exception { withRealm(realmId, (session, realm) -> { realm.setDisplayName("displayName1"); return null; @@ -230,8 +232,138 @@ public class StorageTransactionTest extends KeycloakModelTest { // tx2 should fail as tx1 already changed the value assertException(tx2::commit, - allOf(instanceOf(ModelException.class), - hasCause(instanceOf(OptimisticLockException.class)))); + anyOf( + allOf(instanceOf(RuntimeException.class), hasCause(instanceOf(RollbackException.class))), + allOf(instanceOf(ModelException.class), hasCause(instanceOf(OptimisticLockException.class))), + allOf(instanceOf(OptimisticLockException.class)) + )); + } + } + + @Test + // Optimistic locking works only with map-jpa and map-hotrod + @RequireProvider(value = MapStorageProvider.class, only = {JpaMapStorageProviderFactory.PROVIDER_ID, + HotRodMapStorageProviderFactory.PROVIDER_ID}) + public void testOptimisticLockingExceptionReadByQuery() throws Exception { + withRealm(realmId, (session, realm) -> { + realm.setDisplayName("displayName1"); + return null; + }); + + try (TransactionController tx1 = new TransactionController(getFactory()); + TransactionController tx2 = new TransactionController(getFactory())) { + + // tx1 acquires lock + tx1.begin(); + tx2.begin(); + + // both transactions touch the same entity + tx1.runStep(session -> { + session.realms().getRealmByName("1").setDisplayName("displayName2"); + return null; + }); + tx2.runStep(session -> { + session.realms().getRealmByName("1").setDisplayName("displayName3"); + return null; + }); + + // tx1 transaction should be successful + tx1.commit(); + + // tx2 should fail as tx1 already changed the value + assertException(tx2::commit, + anyOf( + allOf(instanceOf(RuntimeException.class), hasCause(instanceOf(RollbackException.class))), + allOf(instanceOf(ModelException.class), hasCause(instanceOf(OptimisticLockException.class))), + allOf(instanceOf(OptimisticLockException.class)) + )); + } + } + + @Test + // Optimistic locking works only with map-jpa and map-hotrod + @RequireProvider(value = MapStorageProvider.class, only = {JpaMapStorageProviderFactory.PROVIDER_ID, + HotRodMapStorageProviderFactory.PROVIDER_ID}) + public void testOptimisticLockingDeleteWhenReadingByQuery() throws Exception { + withRealm(realmId, (session, realm) -> { + session.users().addUser(realm, "user", "user", false, false); + return null; + }); + + try (TransactionController tx1 = new TransactionController(getFactory()); + TransactionController tx2 = new TransactionController(getFactory())) { + + tx1.begin(); + tx2.begin(); + + // both transactions touch the same entity + tx1.runStep(session -> { + // read by criteria + session.users().getUserByUsername(session.realms().getRealm(realmId), "user").setFirstName("firstName"); + return null; + }); + tx2.runStep(session -> { + RealmModel realm = session.realms().getRealm(realmId); + + // remove by id + session.users().removeUser(realm, session.users().getUserByUsername(realm, "user")); + return null; + }); + + // tx1 transaction should be successful + tx1.commit(); + + // tx2 should fail as tx1 already changed the value + assertException(tx2::commit, + anyOf( + allOf(instanceOf(RuntimeException.class), hasCause(instanceOf(RollbackException.class))), + allOf(instanceOf(ModelException.class), hasCause(instanceOf(OptimisticLockException.class))), + allOf(instanceOf(OptimisticLockException.class)) + )); + } + } + + @Test + // Optimistic locking works only with map-jpa and map-hotrod + @RequireProvider(value = MapStorageProvider.class, only = {JpaMapStorageProviderFactory.PROVIDER_ID, + HotRodMapStorageProviderFactory.PROVIDER_ID}) + public void testOptimisticLockingDeleteWhenReadingById() throws Exception { + String userId = UUID.randomUUID().toString(); + withRealm(realmId, (session, realm) -> { + session.users().addUser(realm, userId, "user", false, false); + return null; + }); + + try (TransactionController tx1 = new TransactionController(getFactory()); + TransactionController tx2 = new TransactionController(getFactory())) { + + tx1.begin(); + tx2.begin(); + + // both transactions touch the same entity + tx1.runStep(session -> { + // read by id + session.users().getUserById(session.realms().getRealm(realmId), userId).setFirstName("firstName"); + return null; + }); + tx2.runStep(session -> { + RealmModel realm = session.realms().getRealm(realmId); + + // remove by id after read by id + session.users().removeUser(realm, session.users().getUserById(realm, userId)); + return null; + }); + + // tx1 transaction should be successful + tx1.commit(); + + // tx2 should fail as tx1 already changed the value + assertException(tx2::commit, + anyOf( + allOf(instanceOf(RuntimeException.class), hasCause(instanceOf(RollbackException.class))), + allOf(instanceOf(ModelException.class), hasCause(instanceOf(OptimisticLockException.class))), + allOf(instanceOf(OptimisticLockException.class)) + )); } } }