Avoid deadlock on CockroachDB when removing authentication sessions

Closes #14991
This commit is contained in:
Alexander Schwartz 2022-10-12 13:05:25 +02:00 committed by Hynek Mlnařík
parent 9b80bad391
commit 3a30061c44
7 changed files with 278 additions and 5 deletions

View file

@ -184,17 +184,20 @@ public abstract class JpaMapKeycloakTransaction<RE extends JpaRootEntity, E exte
}
@Override
@SuppressWarnings("unchecked")
public boolean delete(String key) {
if (key == null) return false;
UUID uuid = UUIDKey.INSTANCE.fromStringSafe(key);
if (uuid == null) return false;
cacheWithinSession.remove(key);
removeFromCache(key);
em.remove(em.getReference(entityType, uuid));
logger.tracef("tx %d: delete entity %s", hashCode(), key);
return true;
}
protected void removeFromCache(String key) {
cacheWithinSession.remove(key);
}
@Override
@SuppressWarnings("unchecked")
public long delete(QueryParameters<M> queryParameters) {

View file

@ -304,6 +304,7 @@ public class JpaMapStorageProviderFactory implements
synchronized (this) {
if (emf == null) {
this.emf = createEntityManagerFactory();
JpaMapUtils.addSpecificNamedQueries(emf);
}
}
}

View file

@ -0,0 +1,210 @@
/*
* Copyright 2016 Red Hat, Inc. and/or its affiliates
* and other contributors as indicated by the @author tags.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.keycloak.models.map.storage.jpa;
import org.hibernate.Session;
import org.hibernate.engine.query.spi.sql.NativeSQLQueryReturn;
import org.hibernate.engine.query.spi.sql.NativeSQLQuerySpecification;
import org.hibernate.engine.spi.SessionFactoryImplementor;
import org.jboss.logging.Logger;
import javax.persistence.EntityManager;
import javax.persistence.EntityManagerFactory;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.util.Collections;
import java.util.Map;
import java.util.Properties;
import java.util.regex.Pattern;
import static org.keycloak.models.map.storage.jpa.JpaMapStorageProviderFactory.HIBERNATE_DEFAULT_SCHEMA;
/**
* @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
*/
public class JpaMapUtils {
public static final String QUERY_NATIVE_SUFFIX = "[native]";
public static final String QUERY_JPQL_SUFFIX = "[jpql]";
private static final Logger logger = Logger.getLogger(JpaMapUtils.class);
public static String getSchemaForNativeQuery(EntityManager em) {
String schema = (String) em.getEntityManagerFactory().getProperties().get(HIBERNATE_DEFAULT_SCHEMA);
return (schema == null) ? "" : schema + ".";
}
/**
* Method that adds the different query variants for the database.
* The method loads the queries specified in the files
* <em>META-INF/jpa-map/queries-{dbType}.properties</em> and the default
* <em>META-INF/jpa-map/queries-default.properties</em>. At least the default file
* should exist inside the jar file. The default file contains all the
* needed queries and the specific one can overload all or some of them for
* that database type.
* @param databaseType The database type as returned by <code>getDatabaseType</code>
*/
public static Properties loadSpecificNamedQueries(String databaseType) {
URL specificUrl = JpaMapUtils.class.getClassLoader().getResource("META-INF/jpa-map/queries-" + databaseType + ".properties");
URL defaultUrl = JpaMapUtils.class.getClassLoader().getResource("META-INF/jpa-map/queries-default.properties");
if (defaultUrl == null) {
throw new IllegalStateException("META-INF/jpa-map/queries-default.properties was not found in the classpath");
}
Properties specificQueries = loadSqlProperties(specificUrl);
Properties defaultQueries = loadSqlProperties(defaultUrl);
Properties queries = new Properties();
for (String queryNameFull : defaultQueries.stringPropertyNames()) {
String querySql = defaultQueries.getProperty(queryNameFull);
String queryName = getQueryShortName(queryNameFull);
String specificQueryNameFull = getQueryFromProperties(queryName, specificQueries);
if (specificQueryNameFull != null) {
// the query is redefined in the specific database file => use it
queryNameFull = specificQueryNameFull;
querySql = specificQueries.getProperty(queryNameFull);
}
queries.put(queryNameFull, querySql);
}
return queries;
}
/**
* Returns the name of the query in the queries file. It searches for the
* three possible forms: name[native], name[jpql] or name.
* @param name The name of the query to search
* @param queries The properties file with the queries
* @return The key with the query found or null if not found
*/
private static String getQueryFromProperties(String name, Properties queries) {
if (queries == null) {
return null;
}
String nameFull = name + QUERY_NATIVE_SUFFIX;
if (queries.containsKey(nameFull)) {
return nameFull;
}
nameFull = name + QUERY_JPQL_SUFFIX;
if (queries.containsKey(nameFull)) {
return nameFull;
}
nameFull = name;
if (queries.containsKey(nameFull)) {
return nameFull;
}
return null;
}
/**
* Loads the URL as a properties file.
* @param url The url to load, it can be null
* @return A properties file with the url loaded or null
*/
public static Properties loadSqlProperties(URL url) {
if (url == null) {
return null;
}
Properties props = new Properties();
try (InputStream is = url.openStream()) {
props.load(is);
} catch (IOException e) {
throw new IllegalStateException(e);
}
return props;
}
/**
* Configures a named query to Hibernate.
*
* @param queryName the query name
* @param querySql the query SQL
* @param entityManager the entity manager
*/
public static void configureNamedQuery(String queryName, String querySql, EntityManager entityManager) {
boolean isNative = queryName.endsWith(QUERY_NATIVE_SUFFIX);
queryName = getQueryShortName(queryName);
logger.tracef("adding query from properties files native=%b %s:%s", isNative, queryName, querySql);
SessionFactoryImplementor sessionFactory = entityManager.getEntityManagerFactory().unwrap(SessionFactoryImplementor.class);
if (isNative) {
NativeSQLQuerySpecification spec = new NativeSQLQuerySpecification(querySql, new NativeSQLQueryReturn[0], Collections.emptySet());
sessionFactory.getQueryPlanCache().getNativeSQLQueryPlan(spec);
sessionFactory.addNamedQuery(queryName, entityManager.createNativeQuery(querySql));
} else {
sessionFactory.getQueryPlanCache().getHQLQueryPlan(querySql, false, Collections.emptyMap());
sessionFactory.addNamedQuery(queryName, entityManager.createQuery(querySql));
}
}
/**
* Returns the query name but removing the suffix.
* @param name The query name as it is on the key
* @return The name without the suffix
*/
private static String getQueryShortName(String name) {
if (name.endsWith(QUERY_NATIVE_SUFFIX)) {
return name.substring(0, name.length() - QUERY_NATIVE_SUFFIX.length());
} else if (name.endsWith(QUERY_JPQL_SUFFIX)) {
return name.substring(0, name.length() - QUERY_JPQL_SUFFIX.length());
} else {
return name;
}
}
public static String getDatabaseType(String productName) {
switch (productName) {
case "Microsoft SQL Server":
case "SQLOLEDB":
return "mssql";
case "EnterpriseDB":
return "postgresql";
default:
return productName.toLowerCase();
}
}
public static void addSpecificNamedQueries(EntityManagerFactory emf) {
EntityManager em = null;
try {
em = emf.createEntityManager();
String dbProductName = em.unwrap(Session.class).doReturningWork(connection -> connection.getMetaData().getDatabaseProductName());
String dbKind = getDatabaseType(dbProductName);
String schemaForNativeQuery = getSchemaForNativeQuery(em);
for (Map.Entry<Object, Object> query : loadSpecificNamedQueries(dbKind.toLowerCase()).entrySet()) {
String queryName = query.getKey().toString();
String querySql = query.getValue().toString();
querySql = querySql.replaceAll(Pattern.quote("${schemaprefix}"), schemaForNativeQuery);
configureNamedQuery(queryName, querySql, em);
}
} finally {
if (em != null) {
em.close();
}
}
}
}

View file

@ -17,24 +17,33 @@
package org.keycloak.models.map.storage.jpa.authSession;
import javax.persistence.EntityManager;
import javax.persistence.Query;
import javax.persistence.criteria.CriteriaBuilder;
import javax.persistence.criteria.Root;
import javax.persistence.criteria.Selection;
import org.hibernate.Session;
import org.hibernate.query.NativeQuery;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.ModelException;
import org.keycloak.models.map.authSession.MapRootAuthenticationSessionEntity;
import org.keycloak.models.map.authSession.MapRootAuthenticationSessionEntityDelegate;
import static org.keycloak.models.map.storage.jpa.Constants.CURRENT_SCHEMA_VERSION_AUTH_SESSION;
import org.keycloak.models.map.common.StringKeyConverter;
import org.keycloak.models.map.storage.jpa.JpaMapKeycloakTransaction;
import org.keycloak.models.map.storage.jpa.JpaModelCriteriaBuilder;
import org.keycloak.models.map.storage.jpa.JpaRootEntity;
import org.keycloak.models.map.storage.jpa.authSession.delegate.JpaRootAuthenticationSessionDelegateProvider;
import org.keycloak.models.map.storage.jpa.authSession.entity.JpaAuthenticationSessionEntity;
import org.keycloak.models.map.storage.jpa.authSession.entity.JpaRootAuthenticationSessionEntity;
import org.keycloak.sessions.RootAuthenticationSessionModel;
import java.sql.Connection;
import java.util.UUID;
public class JpaRootAuthenticationSessionMapKeycloakTransaction extends JpaMapKeycloakTransaction<JpaRootAuthenticationSessionEntity, MapRootAuthenticationSessionEntity, RootAuthenticationSessionModel> {
@SuppressWarnings("unchecked")
public JpaRootAuthenticationSessionMapKeycloakTransaction(KeycloakSession session, EntityManager em) {
super(session, JpaRootAuthenticationSessionEntity.class, RootAuthenticationSessionModel.class, em);
}
@ -65,4 +74,40 @@ public class JpaRootAuthenticationSessionMapKeycloakTransaction extends JpaMapKe
protected MapRootAuthenticationSessionEntity mapToEntityDelegate(JpaRootAuthenticationSessionEntity original) {
return new MapRootAuthenticationSessionEntityDelegate(new JpaRootAuthenticationSessionDelegateProvider(original, em));
}
@Override
public boolean delete(String key) {
int isolationLevel = em.unwrap(Session.class).doReturningWork(Connection::getTransactionIsolation);
if (isolationLevel == Connection.TRANSACTION_SERIALIZABLE) {
// If the isolation level is SERIALIZABLE, there is no need to apply the optimistic locking, as the database with its serializable checks
// takes care that no-one has modified or deleted the row sind the transaction started. On CockroachDB, using optimistic locking with the added
// version column in a delete-statement will cause a table lock, which will lead to deadlock.
// As a workaround, this is using a native query instead, without including the version for optimistic locking.
if (key == null) return false;
UUID uuid = StringKeyConverter.UUIDKey.INSTANCE.fromStringSafe(key);
if (uuid == null) return false;
removeFromCache(key);
// will throw an exception if the entity doesn't exist in the Hibernate session or in the database.
JpaRootAuthenticationSessionEntity rootAuth = em.getReference(JpaRootAuthenticationSessionEntity.class, uuid);
// will use cascading delete to all child entities
//noinspection JpaQueryApiInspection
Query deleteById =
em.createNamedQuery("deleteRootAuthenticationSessionByIdNoOptimisticLocking");
deleteById.unwrap(NativeQuery.class).addSynchronizedQuerySpace(JpaRootAuthenticationSessionEntity.TABLE_NAME,
JpaAuthenticationSessionEntity.TABLE_NAME);
deleteById.setParameter("id", key);
int deleteCount = deleteById.executeUpdate();
rootAuth.getAuthenticationSessions().forEach(e -> em.detach(e));
em.detach(rootAuth);
if (deleteCount == 1) {
return true;
} else if (deleteCount == 0) {
throw new ModelException("Unable to find root authentication session");
} else {
throw new ModelException("Deleted " + deleteCount + " root authentication session when expecting to delete one");
}
} else {
return super.delete(key);
}
}
}

View file

@ -37,6 +37,8 @@ import org.keycloak.models.map.authSession.MapAuthenticationSessionEntity;
import org.keycloak.models.map.common.DeepCloner;
import org.keycloak.models.map.common.UpdatableEntity;
import static org.keycloak.models.map.storage.jpa.Constants.CURRENT_SCHEMA_VERSION_AUTH_SESSION;
import static org.keycloak.models.map.storage.jpa.authSession.entity.JpaAuthenticationSessionEntity.TABLE_NAME;
import org.keycloak.models.map.storage.jpa.JpaChildEntity;
import org.keycloak.models.map.storage.jpa.JpaRootVersionedEntity;
import org.keycloak.models.map.storage.jpa.hibernate.jsonb.JsonbType;
@ -46,10 +48,11 @@ import org.keycloak.sessions.CommonClientSessionModel;
* Entity represents individual authentication session.
*/
@Entity
@Table(name = "kc_auth_session")
@Table(name = TABLE_NAME)
@TypeDefs({@TypeDef(name = "jsonb", typeClass = JsonbType.class)})
public class JpaAuthenticationSessionEntity extends UpdatableEntity.Impl implements MapAuthenticationSessionEntity, JpaRootVersionedEntity, JpaChildEntity<JpaRootAuthenticationSessionEntity>{
public static final String TABLE_NAME = "kc_auth_session";
@Id
@Column
@GeneratedValue

View file

@ -43,6 +43,7 @@ import org.keycloak.models.map.storage.jpa.JpaRootVersionedEntity;
import org.keycloak.models.map.storage.jpa.hibernate.jsonb.JsonbType;
import static org.keycloak.models.map.storage.jpa.JpaMapStorageProviderFactory.CLONER;
import static org.keycloak.models.map.storage.jpa.authSession.entity.JpaRootAuthenticationSessionEntity.TABLE_NAME;
/**
* Entity represents root authentication session.
@ -52,10 +53,12 @@ import static org.keycloak.models.map.storage.jpa.JpaMapStorageProviderFactory.C
* therefore marked as non-insertable and non-updatable to instruct hibernate.
*/
@Entity
@Table(name = "kc_auth_root_session")
@Table(name = TABLE_NAME)
@TypeDefs({@TypeDef(name = "jsonb", typeClass = JsonbType.class)})
public class JpaRootAuthenticationSessionEntity extends AbstractRootAuthenticationSessionEntity implements JpaRootVersionedEntity {
public static final String TABLE_NAME = "kc_auth_root_session";
@Id
@Column
private UUID id;

View file

@ -0,0 +1,8 @@
# properties file to define all default queries that are loaded separately
# in a properties file. These queries can be overloaded with a
# specific file for each database type. Queries are defined in the form:
# name[type]=sql
# type can be native (for native queries) or jpql (jpql syntax)
# if no type is defined jpql is the default
deleteRootAuthenticationSessionByIdNoOptimisticLocking[native]=delete from ${schemaprefix}kc_auth_root_session where id = :id