Don't rely on DefaultModeLCriteria in equals/hashCode

Instead, map this to JPA query and then create the cache lookup key from there.

Closes #14938
This commit is contained in:
Alexander Schwartz 2022-10-19 21:50:43 +02:00 committed by Hynek Mlnařík
parent e494649a4e
commit 9fb9780f02
2 changed files with 86 additions and 45 deletions

View file

@ -20,11 +20,12 @@ import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.UUID;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import javax.persistence.EntityManager;
import javax.persistence.LockModeType;
import javax.persistence.Parameter;
import javax.persistence.TypedQuery;
import javax.persistence.criteria.CriteriaBuilder;
import javax.persistence.criteria.CriteriaDelete;
@ -36,9 +37,11 @@ import javax.persistence.criteria.Selection;
import org.hibernate.Session;
import org.hibernate.internal.SessionImpl;
import org.hibernate.query.spi.QueryImplementor;
import org.jboss.logging.Logger;
import org.keycloak.common.util.Time;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.ModelException;
import org.keycloak.models.map.common.AbstractEntity;
import org.keycloak.models.map.common.ExpirableEntity;
import org.keycloak.models.map.common.StringKeyConverter;
@ -130,26 +133,6 @@ public abstract class JpaMapKeycloakTransaction<RE extends JpaRootEntity, E exte
@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());
@ -184,37 +167,59 @@ public abstract class JpaMapKeycloakTransaction<RE extends JpaRootEntity, E exte
}
if (predicateFunc != null) query.where(predicateFunc.apply(cb, query::subquery, root));
TypedQuery<RE> emQuery = em.createQuery(query);
TypedQuery<RE> emQuery = paginateQuery(em.createQuery(query), queryParameters.getOffset(), queryParameters.getLimit());
Map<QueryCacheKey, List<RE>> cache = getQueryCache();
QueryCacheKey queryCacheKey = new QueryCacheKey(emQuery, modelType);
if (!LockObjectsForModification.isEnabled(this.session, modelType)) {
List<RE> previousResult = cache.get(queryCacheKey);
//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/%s%s", hashCode(), queryParameters, queryCacheKey, getShortStackTrace());
return closing(previousResult.stream()).map(this::mapToEntityDelegateUnique);
} else {
logger.tracef("tx %d: cache ignored due to dirty session", hashCode());
}
}
}
logger.tracef("tx %d: cache miss for %s/%s%s", hashCode(), queryParameters, queryCacheKey, getShortStackTrace());
if (LockObjectsForModification.isEnabled(session, modelType)) {
emQuery = emQuery.setLockMode(LockModeType.PESSIMISTIC_WRITE);
}
// 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);
List<RE> resultList = emQuery.getResultList();
cache.put(queryCacheKey, resultList);
return closing(resultList.stream()).map(this::mapToEntityDelegateUnique);
}
private Map<QueryParameters<M>, List<RE>> getQueryCache() {
private Map<QueryCacheKey, 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);
Map<QueryCacheKey, List<RE>> cache = (Map<QueryCacheKey, 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<>());
return cache;
}
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);
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());
queryCache.clear();
}
}
@ -321,4 +326,56 @@ public abstract class JpaMapKeycloakTransaction<RE extends JpaRootEntity, E exte
return cb.or(cb.greaterThan(root.get("expiration"), Time.currentTimeMillis()),
cb.isNull(root.get("expiration")));
}
private static class QueryCacheKey {
private final String queryString;
private final Integer queryMaxResults;
private final Integer queryFirstResult;
private final HashMap<String, Object> queryParameters;
private final Class<?> modelType;
public QueryCacheKey(TypedQuery<?> emQuery, Class<?> modelType) {
// copy over all fields from the query that relevant for caching
QueryImplementor<?> query = emQuery.unwrap(QueryImplementor.class);
this.queryString = query.getQueryString();
this.queryParameters = new HashMap<>();
for (Parameter<?> parameter : query.getParameters()) {
if (parameter.getName() == null) {
throw new ModelException("Can't prepare query for caching as parameter doesn't have a name");
}
this.queryParameters.put(parameter.getName(), query.getParameterValue(parameter.getName()));
}
this.queryMaxResults = emQuery.getMaxResults();
this.queryFirstResult = emQuery.getFirstResult();
this.modelType = modelType;
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
QueryCacheKey that = (QueryCacheKey) o;
return Objects.equals(queryString, that.queryString)
&& Objects.equals(queryMaxResults, that.queryMaxResults)
&& Objects.equals(queryFirstResult, that.queryFirstResult)
&& Objects.equals(queryParameters, that.queryParameters)
&& Objects.equals(modelType, that.modelType);
}
@Override
public int hashCode() {
return Objects.hash(queryString, queryMaxResults, queryFirstResult, queryParameters, modelType);
}
@Override
public String toString() {
return "QueryCacheKey{" +
"queryString='" + queryString + '\'' +
", queryMaxResults=" + queryMaxResults +
", queryFirstResult=" + queryFirstResult +
", queryParameters=" + queryParameters +
", modelType=" + modelType.getName() +
'}';
}
}
}

View file

@ -5,7 +5,6 @@ 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;
@ -70,21 +69,6 @@ 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
*