diff --git a/model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapKeycloakTransaction.java b/model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapKeycloakTransaction.java index 7344975913b..2021c94db6b 100644 --- a/model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapKeycloakTransaction.java +++ b/model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapKeycloakTransaction.java @@ -19,7 +19,9 @@ package org.keycloak.models.map.storage.jpa; import java.util.HashMap; import java.util.LinkedList; import java.util.List; +import java.util.Map; import java.util.UUID; +import java.util.stream.Collectors; import java.util.stream.Stream; import javax.persistence.EntityManager; import javax.persistence.LockModeType; @@ -32,6 +34,8 @@ import javax.persistence.criteria.Predicate; import javax.persistence.criteria.Root; import javax.persistence.criteria.Selection; +import org.hibernate.Session; +import org.hibernate.internal.SessionImpl; import org.jboss.logging.Logger; import org.keycloak.common.util.Time; import org.keycloak.models.KeycloakSession; @@ -45,6 +49,7 @@ import org.keycloak.models.map.storage.chm.MapFieldPredicates; import org.keycloak.models.map.storage.chm.MapModelCriteriaBuilder; import org.keycloak.utils.LockObjectsForModification; +import static org.keycloak.common.util.StackUtil.getShortStackTrace; import static org.keycloak.models.map.common.ExpirationUtils.isExpired; import static org.keycloak.models.map.storage.jpa.JpaMapStorageProviderFactory.CLONER; import static org.keycloak.models.map.storage.jpa.PaginationUtils.paginateQuery; @@ -120,9 +125,31 @@ public abstract class JpaMapKeycloakTransaction read(QueryParameters queryParameters) { + Map, List> cache = getQueryCache(); + if (!LockObjectsForModification.isEnabled(this.session, modelType)) { + List previousResult = cache.get(queryParameters); + //noinspection resource + SessionImpl session = (SessionImpl) em.unwrap(Session.class); + // only do dirty checking if there is a previously cached result that would match the query + if (previousResult != null) { + // if the session is dirty, data has been modified, and the cache must not be used + // check if there are queued actions already, as this allows us to skip the expensive dirty check + if (!session.getActionQueue().areInsertionsOrDeletionsQueued() && session.getActionQueue().numberOfUpdates() == 0 && session.getActionQueue().numberOfCollectionUpdates() == 0 && + !session.isDirty()) { + logger.tracef("tx %d: cache hit for %s for model %s%s", hashCode(), queryParameters, modelType.getName(), getShortStackTrace()); + return closing(previousResult.stream()).map(this::mapToEntityDelegateUnique); + } else { + logger.tracef("tx %d: cache ignored due to dirty session for %s for model %s%s", hashCode(), queryParameters, modelType.getName(), getShortStackTrace()); + } + } + } + logger.tracef("tx %d: cache miss for %s for model %s%s", hashCode(), queryParameters, modelType.getName(), getShortStackTrace()); + JpaModelCriteriaBuilder mcb = queryParameters.getModelCriteriaBuilder() .flashToModelCriteriaBuilder(createJpaModelCriteriaBuilder()); @@ -161,8 +188,34 @@ public abstract class JpaMapKeycloakTransaction resultList = paginateQuery(emQuery, queryParameters.getOffset(), queryParameters.getLimit()).getResultList(); + cache.put(queryParameters, resultList); + + return closing(resultList.stream()).map(this::mapToEntityDelegateUnique); + } + + private Map, List> getQueryCache() { + //noinspection resource,unchecked + Map, Map, List>> cache = (Map, Map, List>>) em.unwrap(Session.class).getProperties().get(JPA_MAP_CACHE); + if (cache == null) { + cache = new HashMap<>(); + //noinspection resource + em.unwrap(Session.class).setProperty(JPA_MAP_CACHE, cache); + } + return cache.computeIfAbsent(modelType, k -> new HashMap<>()); + } + + public static void clearQueryCache(Session session) { + logger.tracef("query cache cleared"); + //noinspection unchecked + Map, Map> queryCache = (HashMap, Map>) session.getProperties().get(JPA_MAP_CACHE); + if (queryCache != null) { + // Can't set null as a property values as it is not serializable. Clearing each map so that the current query result might be saved. + queryCache.forEach((queryParameters, map) -> map.clear()); + } } @Override diff --git a/model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/hibernate/listeners/JpaAutoFlushListener.java b/model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/hibernate/listeners/JpaAutoFlushListener.java index c75e730df94..83ca58c2d51 100644 --- a/model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/hibernate/listeners/JpaAutoFlushListener.java +++ b/model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/hibernate/listeners/JpaAutoFlushListener.java @@ -25,12 +25,17 @@ import org.hibernate.event.spi.AutoFlushEvent; import org.hibernate.event.spi.EventSource; import org.hibernate.internal.CoreMessageLogger; import org.jboss.logging.Logger; +import org.keycloak.models.map.storage.jpa.JpaMapKeycloakTransaction; /** * Extends Hibernate's {@link DefaultAutoFlushEventListener} to always flush queued inserts to allow correct handling - * of orphans of that entities in the same transactions. + * of orphans of that entities in the same transactions, and also to clear a session-level query cache. + *

* If they wouldn't be flushed, they won't be orphaned (at least not in Hibernate 5.3.24.Final). * This class copies over all functionality of the base class that can't be overwritten via inheritance. + * This is being tracked as part of keycloak/keycloak#11666. + *

+ * This also clears the JPA map store query level cache for the {@link JpaMapKeycloakTransaction} whenever there is some data written to the database. */ public class JpaAutoFlushListener extends DefaultAutoFlushEventListener { @@ -78,18 +83,23 @@ public class JpaAutoFlushListener extends DefaultAutoFlushEventListener { } private boolean flushIsReallyNeeded(AutoFlushEvent event, final EventSource source) { - return source.getHibernateFlushMode() == FlushMode.ALWAYS + boolean flushIsReallyNeeded = source.getHibernateFlushMode() == FlushMode.ALWAYS // START OF FIX for auto-flush-mode on inserts that might later be deleted in same transaction || source.getActionQueue().numberOfInsertions() > 0 // END OF FIX || source.getActionQueue().areTablesToBeUpdated(event.getQuerySpaces()); + if (flushIsReallyNeeded) { + // clear the per-session query cache, as changing an entity might change any of the cached query results + JpaMapKeycloakTransaction.clearQueryCache(source.getSession()); + } + return flushIsReallyNeeded; } private boolean flushMightBeNeeded(final EventSource source) { return !source.getHibernateFlushMode().lessThan(FlushMode.AUTO) && source.getDontFlushFromFind() == 0 && (source.getPersistenceContext().getNumberOfManagedEntities() > 0 || - source.getPersistenceContext().getCollectionEntries().size() > 0); + source.getPersistenceContext().getCollectionEntriesSize() > 0); } } \ No newline at end of file diff --git a/model/map/src/main/java/org/keycloak/models/map/storage/QueryParameters.java b/model/map/src/main/java/org/keycloak/models/map/storage/QueryParameters.java index c422685920a..918924151b9 100644 --- a/model/map/src/main/java/org/keycloak/models/map/storage/QueryParameters.java +++ b/model/map/src/main/java/org/keycloak/models/map/storage/QueryParameters.java @@ -5,6 +5,7 @@ import org.keycloak.storage.SearchableModelField; import java.util.LinkedList; import java.util.List; +import java.util.Objects; import static org.keycloak.models.map.storage.QueryParameters.Order.ASCENDING; @@ -69,6 +70,21 @@ public class QueryParameters { return this; } + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + QueryParameters that = (QueryParameters) o; + // there is currently no equals method for the ModelCriteriaBuilder, take its String representation as a substitute. + return Objects.equals(offset, that.offset) && Objects.equals(limit, that.limit) && Objects.equals(orderBy, that.orderBy) && Objects.equals(mcb.toString(), that.mcb.toString()); + } + + @Override + public int hashCode() { + // there is currently no equals method for the ModelCriteriaBuilder, take its String representation as a substitute. + return Objects.hash(offset, limit, orderBy, mcb.toString()); + } + /** * Sets offset parameter * @@ -107,6 +123,16 @@ public class QueryParameters { return orderBy; } + @Override + public String toString() { + return "QueryParameters{" + + "offset=" + offset + + ", limit=" + limit + + ", orderBy=" + orderBy + + ", mcb=" + mcb + + '}'; + } + /** * Enum for ascending or descending ordering */ diff --git a/model/map/src/main/java/org/keycloak/models/map/user/MapUserAdapter.java b/model/map/src/main/java/org/keycloak/models/map/user/MapUserAdapter.java index 975f9cd67b1..5113a220a61 100644 --- a/model/map/src/main/java/org/keycloak/models/map/user/MapUserAdapter.java +++ b/model/map/src/main/java/org/keycloak/models/map/user/MapUserAdapter.java @@ -32,6 +32,7 @@ import org.keycloak.models.utils.RoleUtils; import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.Set; import java.util.stream.Stream; @@ -322,7 +323,7 @@ public abstract class MapUserAdapter extends AbstractUserModel { public Stream getRoleMappingsStream() { Set roles = entity.getRolesMembership(); if (roles == null || roles.isEmpty()) return Stream.empty(); - return entity.getRolesMembership().stream().map(realm::getRoleById); + return entity.getRolesMembership().stream().map(realm::getRoleById).filter(Objects::nonNull); } @Override