External-Mirrors/keycloak
2
0
Fork 0
mirror of https://github.com/keycloak/keycloak.git synced 2026-01-09 23:12:06 -03:30
Code Issues Packages Projects Releases Wiki Activity

First naïve per-session caching for JPA map store

Browse Source 
Closes #14938
This commit is contained in:
Alexander Schwartz 2022-10-17 18:45:44 +02:00 committed by Hynek Mlnařík
parent 8f9c3cdeab
commit e494649a4e
4 changed files with 96 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

57
model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/JpaMapKeycloakTransaction.java
View File

@ -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<RE extends JpaRootEntity, E exte
return e != null && isExpirableEntity && isExpired((ExpirableEntity) e, true) ? null : e;
}
private final static String JPA_MAP_CACHE = "keycloak.jpamap.cache";
@Override
@SuppressWarnings("unchecked")
public Stream<E> read(QueryParameters<M> queryParameters) {
Map<QueryParameters<M>, List<RE>> cache = getQueryCache();
if (!LockObjectsForModification.isEnabled(this.session, modelType)) {
List<RE> 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<RE extends JpaRootEntity, E exte
if (LockObjectsForModification.isEnabled(session, modelType)) {
emQuery = emQuery.setLockMode(LockModeType.PESSIMISTIC_WRITE);
}
return closing(paginateQuery(emQuery, queryParameters.getOffset(), queryParameters.getLimit()).getResultStream())
.map(this::mapToEntityDelegateUnique);
// In order to cache the result, the full result needs to be retrieved.
// There is also no difference to that in Hibernate, as Hibernate will first retrieve all elements from the ResultSet.
List<RE> resultList = paginateQuery(emQuery, queryParameters.getOffset(), queryParameters.getLimit()).getResultList();
cache.put(queryParameters, resultList);
return closing(resultList.stream()).map(this::mapToEntityDelegateUnique);
}
private Map<QueryParameters<M>, List<RE>> getQueryCache() {
//noinspection resource,unchecked
Map<Class<?>, Map<QueryParameters<M>, List<RE>>> cache = (Map<Class<?>, Map<QueryParameters<M>, List<RE>>>) 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<Class<?>, Map<?,?>> queryCache = (HashMap<Class<?>, 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

16
model/map-jpa/src/main/java/org/keycloak/models/map/storage/jpa/hibernate/listeners/JpaAutoFlushListener.java
View File

@ -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.
* <p />
* 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 <a href="https://github.com/keycloak/keycloak/issues/11666">keycloak/keycloak#11666</a>.
* <p />
* 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);
}
}

26
model/map/src/main/java/org/keycloak/models/map/storage/QueryParameters.java
View File

@ -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<M> {
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<M> {
return orderBy;
}
@Override
public String toString() {
return "QueryParameters{" +
"offset=" + offset +
", limit=" + limit +
", orderBy=" + orderBy +
", mcb=" + mcb +
'}';
}
/**
* Enum for ascending or descending ordering
*/

3
model/map/src/main/java/org/keycloak/models/map/user/MapUserAdapter.java
View File

@ -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<MapUserEntity> {
public Stream<RoleModel> getRoleMappingsStream() {
Set<String> 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