KEYCLOAK-18298 ClearExpiredUserSessions timeouts with large number of sessions

This commit is contained in:
Martin Kanis 2021-05-27 14:41:30 +02:00 committed by Hynek Mlnařík
parent 669556af71
commit 122fbe1bc6
4 changed files with 39 additions and 26 deletions

View file

@ -275,34 +275,28 @@ public final class KeycloakModelUtils {
* @param timeoutInSeconds
*/
public static void runJobInTransactionWithTimeout(KeycloakSessionFactory factory, KeycloakSessionTask task, int timeoutInSeconds) {
JtaTransactionManagerLookup lookup = (JtaTransactionManagerLookup)factory.getProviderFactory(JtaTransactionManagerLookup.class);
try {
if (lookup != null) {
if (lookup.getTransactionManager() != null) {
try {
lookup.getTransactionManager().setTransactionTimeout(timeoutInSeconds);
} catch (SystemException e) {
throw new RuntimeException(e);
}
}
}
setTransactionLimit(factory, timeoutInSeconds);
runJobInTransaction(factory, task);
} finally {
if (lookup != null) {
if (lookup.getTransactionManager() != null) {
try {
// Reset to default transaction timeout
lookup.getTransactionManager().setTransactionTimeout(0);
} catch (SystemException e) {
// Shouldn't happen for Wildfly transaction manager
throw new RuntimeException(e);
}
setTransactionLimit(factory, 0);
}
}
public static void setTransactionLimit(KeycloakSessionFactory factory, int timeoutInSeconds) {
JtaTransactionManagerLookup lookup = (JtaTransactionManagerLookup) factory.getProviderFactory(JtaTransactionManagerLookup.class);
if (lookup != null) {
if (lookup.getTransactionManager() != null) {
try {
// If timeout is set to 0, reset to default transaction timeout
lookup.getTransactionManager().setTransactionTimeout(timeoutInSeconds);
} catch (SystemException e) {
// Shouldn't happen for Wildfly transaction manager
throw new RuntimeException(e);
}
}
}
}
public static Function<KeycloakSessionFactory, ComponentModel> componentModelGetter(String realmId, String componentId) {

View file

@ -20,6 +20,7 @@ package org.keycloak.services.scheduled;
import org.jboss.logging.Logger;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.KeycloakSessionFactory;
import org.keycloak.models.utils.KeycloakModelUtils;
import org.keycloak.services.ServicesLogger;
import org.keycloak.timer.ScheduledTask;
@ -32,22 +33,34 @@ public class ScheduledTaskRunner implements Runnable {
protected final KeycloakSessionFactory sessionFactory;
protected final ScheduledTask task;
private int transactionLimit;
public ScheduledTaskRunner(KeycloakSessionFactory sessionFactory, ScheduledTask task) {
this.sessionFactory = sessionFactory;
this.task = task;
}
public ScheduledTaskRunner(KeycloakSessionFactory sessionFactory, ScheduledTask task, int transactionLimit) {
this(sessionFactory, task);
this.transactionLimit = transactionLimit;
}
@Override
public void run() {
KeycloakSession session = sessionFactory.create();
try {
if (transactionLimit != 0) {
KeycloakModelUtils.setTransactionLimit(sessionFactory, transactionLimit);
}
runTask(session);
} catch (Throwable t) {
ServicesLogger.LOGGER.failedToRunScheduledTask(t, task.getClass().getSimpleName());
session.getTransactionManager().rollback();
} finally {
if (transactionLimit != 0) {
KeycloakModelUtils.setTransactionLimit(sessionFactory, 0);
}
try {
session.close();
} catch (Throwable t) {

View file

@ -35,11 +35,13 @@ public class BasicTimerProvider implements TimerProvider {
private final KeycloakSession session;
private final Timer timer;
private final int transactionTimeout;
private final BasicTimerProviderFactory factory;
public BasicTimerProvider(KeycloakSession session, Timer timer, BasicTimerProviderFactory factory) {
public BasicTimerProvider(KeycloakSession session, Timer timer, int transactionTimeout, BasicTimerProviderFactory factory) {
this.session = session;
this.timer = timer;
this.transactionTimeout = transactionTimeout;
this.factory = factory;
}
@ -65,7 +67,7 @@ public class BasicTimerProvider implements TimerProvider {
@Override
public void scheduleTask(ScheduledTask scheduledTask, long intervalMillis, String taskName) {
ScheduledTaskRunner scheduledTaskRunner = new ScheduledTaskRunner(session.getKeycloakSessionFactory(), scheduledTask);
ScheduledTaskRunner scheduledTaskRunner = new ScheduledTaskRunner(session.getKeycloakSessionFactory(), scheduledTask, transactionTimeout);
this.schedule(scheduledTaskRunner, intervalMillis, taskName);
}

View file

@ -24,7 +24,6 @@ import org.keycloak.timer.TimerProvider;
import org.keycloak.timer.TimerProviderFactory;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
@ -35,15 +34,20 @@ public class BasicTimerProviderFactory implements TimerProviderFactory {
private Timer timer;
private int transactionTimeout;
public static final String TRANSACTION_TIMEOUT = "transactionTimeout";
private ConcurrentMap<String, TimerTaskContextImpl> scheduledTasks = new ConcurrentHashMap<>();
@Override
public TimerProvider create(KeycloakSession session) {
return new BasicTimerProvider(session, timer, this);
return new BasicTimerProvider(session, timer, transactionTimeout, this);
}
@Override
public void init(Config.Scope config) {
transactionTimeout = config.getInt(TRANSACTION_TIMEOUT, 0);
timer = new Timer();
}