Make RunWorkflowTask aware of executor cancellation due to timeout

Closes #45175

Signed-off-by: Stefan Guilhen <sguilhen@redhat.com>
This commit is contained in:
Stefan Guilhen 2026-02-11 07:30:49 -03:00 committed by Pedro Igor
parent 04d8886678
commit 1384d3b72a
4 changed files with 29 additions and 5 deletions

View file

@ -1,6 +1,8 @@
package org.keycloak.models.workflow;
import java.util.concurrent.TimeoutException;
import org.keycloak.common.util.DurationConverter;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.utils.KeycloakModelUtils;
@ -72,7 +74,11 @@ class RunWorkflowTask extends WorkflowTransactionalTask {
String nextStepId = KeycloakModelUtils.runJobInTransactionWithResult(s.getKeycloakSessionFactory(), s.getContext(), session -> {
// we need a copy of the context with the new session to run the step provider
DefaultWorkflowExecutionContext stepContext = new DefaultWorkflowExecutionContext(session, context, step);
// check if the workflow execution was cancelled before running the step
checkExecutionCancelled(step);
getStepProvider(session, step).run(stepContext);
// now check again if the workflow execution was cancelled after running the step, as the step provider might have taken a long time to execute
checkExecutionCancelled(step);
WorkflowStep nextStep = stepContext.getNextStep();
return nextStep != null ? nextStep.getId() : null;
}, "Workflow step execution task");
@ -98,4 +104,15 @@ class RunWorkflowTask extends WorkflowTransactionalTask {
throw e;
}
}
private void checkExecutionCancelled(WorkflowStep step) {
Throwable throwable = super.futureCancelled.get();
if (super.futureCancelled.get() != null) {
if (throwable instanceof TimeoutException || throwable.getCause() instanceof TimeoutException) {
throw new RuntimeException("Workflow executor timed out during execution of step " + step.getProviderId(), throwable);
} else {
throw new RuntimeException("Workflow executor was cancelled during execution of step " + step.getProviderId(), throwable);
}
}
}
}

View file

@ -24,7 +24,7 @@ final class WorkflowExecutor {
this.taskTimeout = taskTimeout;
}
void runTask(KeycloakSession session, Runnable task) {
void runTask(KeycloakSession session, WorkflowTransactionalTask task) {
enlistTransaction(session, new WorkflowTask(this, task));
}
@ -35,7 +35,7 @@ final class WorkflowExecutor {
if (error instanceof TimeoutException) {
log.warnf("Timeout occurred while processing workflow task: %s", task);
}
task.cancel();
task.cancel(error);
});
if (blocking) {

View file

@ -10,14 +10,14 @@ import org.keycloak.models.utils.KeycloakModelUtils;
public final class WorkflowTask extends AbstractKeycloakTransaction implements Runnable {
private final Runnable task;
private final WorkflowTransactionalTask task;
private final WorkflowExecutor executor;
private final String id;
private CompletableFuture<Void> future;
private long startTime;
private AtomicReference<Thread> thread;
WorkflowTask(WorkflowExecutor executor, Runnable task) {
WorkflowTask(WorkflowExecutor executor, WorkflowTransactionalTask task) {
Objects.requireNonNull(executor, "executor");
Objects.requireNonNull(task, "task");
this.executor = executor;
@ -78,7 +78,8 @@ public final class WorkflowTask extends AbstractKeycloakTransaction implements R
return "id: " + id + ", executionTime: " + (System.currentTimeMillis() - startTime) + "ms , status: " + status + ", task: [" + task.toString() + "]";
}
public void cancel() {
public void cancel(Throwable error) {
this.task.cancel(error);
if (future != null) {
future.cancel(true);
if (thread.get() != null) {

View file

@ -1,6 +1,7 @@
package org.keycloak.models.workflow;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicReference;
import org.keycloak.models.KeycloakContext;
import org.keycloak.models.KeycloakSession;
@ -13,6 +14,7 @@ public abstract class WorkflowTransactionalTask implements Runnable, KeycloakSes
private final KeycloakSessionFactory sessionFactory;
private final KeycloakContext context;
protected AtomicReference<Throwable> futureCancelled = new AtomicReference<Throwable>(null);
public WorkflowTransactionalTask(KeycloakSession session) {
Objects.requireNonNull(session, "KeycloakSession must not be null");
@ -24,4 +26,8 @@ public abstract class WorkflowTransactionalTask implements Runnable, KeycloakSes
public void run() {
runJobInTransaction(sessionFactory, context, this);
}
public void cancel(Throwable error) {
futureCancelled.compareAndSet(null, error);
}
}