libre.sh/keycloak-scim
3
0
Fork 0
Code Issues 15 Pull requests Releases 1 Packages Activity Actions

[KEYCLOAK-16780] - Allow batching writes to storage when running migration (#7717)

Browse source 
Co-authored-by: Hynek Mlnařík <hmlnarik@users.noreply.github.com>

Co-authored-by: Hynek Mlnařík <hmlnarik@users.noreply.github.com>
This commit is contained in:
Pedro Igor 2021-01-29 09:35:19 -03:00 committed by GitHub
parent 8432513daa
commit cdf0ead957
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 42 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
model/jpa/src/main/java/org/keycloak/connections/jpa/DefaultJpaConnectionProviderFactory.java
View file

@ -88,7 +88,7 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
em = emf.createEntityManager(SynchronizationType.SYNCHRONIZED);
}
em = PersistenceExceptionConverter.create(em);
em = PersistenceExceptionConverter.create(session, em);
if (!jtaEnabled) session.getTransactionManager().enlist(new JpaKeycloakTransaction(em));
return new DefaultJpaConnectionProvider(em);
}

31
model/jpa/src/main/java/org/keycloak/connections/jpa/PersistenceExceptionConverter.java
View file

@ -18,6 +18,8 @@
package org.keycloak.connections.jpa;
import org.hibernate.exception.ConstraintViolationException;
import org.keycloak.models.Constants;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.ModelDuplicateException;
import org.keycloak.models.ModelException;
@ -27,31 +29,52 @@ import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.util.regex.Pattern;
/**
* @author <a href="mailto:sthorger@redhat.com">Stian Thorgersen</a>
*/
public class PersistenceExceptionConverter implements InvocationHandler {
private EntityManager em;
private static final Pattern WRITE_METHOD_NAMES = Pattern.compile("persist|merge");
public static EntityManager create(EntityManager em) {
return (EntityManager) Proxy.newProxyInstance(EntityManager.class.getClassLoader(), new Class[]{EntityManager.class}, new PersistenceExceptionConverter(em));
private final EntityManager em;
private final boolean batchEnabled;
private final int batchSize;
private int changeCount = 0;
public static EntityManager create(KeycloakSession session, EntityManager em) {
return (EntityManager) Proxy.newProxyInstance(EntityManager.class.getClassLoader(), new Class[]{EntityManager.class}, new PersistenceExceptionConverter(session, em));
}
private PersistenceExceptionConverter(EntityManager em) {
private PersistenceExceptionConverter(KeycloakSession session, EntityManager em) {
batchEnabled = session.getAttributeOrDefault(Constants.STORAGE_BATCH_ENABLED, false);
batchSize = session.getAttributeOrDefault(Constants.STORAGE_BATCH_SIZE, 100);
this.em = em;
}
@Override
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
try {
flushInBatchIfEnabled(method);
return method.invoke(em, args);
} catch (InvocationTargetException e) {
throw convert(e.getCause());
}
}
private void flushInBatchIfEnabled(Method method) {
if (batchEnabled) {
if (WRITE_METHOD_NAMES.matcher(method.getName()).matches()) {
if (changeCount++ > batchSize) {
em.flush();
em.clear();
changeCount = 0;
}
}
}
}
public static ModelException convert(Throwable t) {
if (t.getCause() != null && t.getCause() instanceof ConstraintViolationException) {
throw new ModelDuplicateException(t);

2
quarkus/runtime/src/main/java/org/keycloak/connections/jpa/QuarkusJpaConnectionProviderFactory.java
View file

@ -112,7 +112,7 @@ public class QuarkusJpaConnectionProviderFactory implements JpaConnectionProvide
em = emf.createEntityManager(SynchronizationType.SYNCHRONIZED);
}
em = PersistenceExceptionConverter.create(em);
em = PersistenceExceptionConverter.create(session, em);
if (!jtaEnabled) session.getTransactionManager().enlist(new JpaKeycloakTransaction(em));
return new DefaultJpaConnectionProvider(em);
}

3
server-spi-private/src/main/java/org/keycloak/migration/MigrationModelManager.java
View file

@ -52,6 +52,7 @@ import org.keycloak.migration.migrators.MigrateTo8_0_2;
import org.keycloak.migration.migrators.MigrateTo9_0_0;
import org.keycloak.migration.migrators.MigrateTo9_0_4;
import org.keycloak.migration.migrators.Migration;
import org.keycloak.models.Constants;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.RealmModel;
import org.keycloak.representations.idm.RealmRepresentation;
@ -96,6 +97,8 @@ public class MigrationModelManager {
};
public static void migrate(KeycloakSession session) {
session.setAttribute(Constants.STORAGE_BATCH_ENABLED, Boolean.getBoolean("keycloak.migration.batch-enabled"));
session.setAttribute(Constants.STORAGE_BATCH_SIZE, Integer.getInteger("keycloak.migration.batch-size"));
MigrationModel model = session.realms().getMigrationModel();
ModelVersion currentVersion = new ModelVersion(Version.VERSION_KEYCLOAK);

10
server-spi-private/src/main/java/org/keycloak/models/Constants.java
View file

@ -108,4 +108,14 @@ public final class Constants {
public static final Pattern CFG_DELIMITER_PATTERN = Pattern.compile("\\s*" + CFG_DELIMITER + "\\s*");
public static final String OFFLINE_ACCESS_SCOPE_CONSENT_TEXT = "${offlineAccessScopeConsentText}";
/**
* If set as an attribute in the {@link KeycloakSession}, indicates that the storage should batch write operations.
*/
public static final String STORAGE_BATCH_ENABLED = "org.keycloak.storage.batch_enabled";
/**
* If {@code #STORAGE_BATCH_ENABLED} is set, indicates the batch size.
*/
public static final String STORAGE_BATCH_SIZE = "org.keycloak.storage.batch_size";
}