From 1384d3b72a43ea6c5258cd46192f288de7344fa4 Mon Sep 17 00:00:00 2001 From: Stefan Guilhen Date: Wed, 11 Feb 2026 07:30:49 -0300 Subject: [PATCH] Make RunWorkflowTask aware of executor cancellation due to timeout Closes #45175 Signed-off-by: Stefan Guilhen --- .../models/workflow/RunWorkflowTask.java | 17 +++++++++++++++++ .../models/workflow/WorkflowExecutor.java | 4 ++-- .../keycloak/models/workflow/WorkflowTask.java | 7 ++++--- .../workflow/WorkflowTransactionalTask.java | 6 ++++++ 4 files changed, 29 insertions(+), 5 deletions(-) diff --git a/model/jpa/src/main/java/org/keycloak/models/workflow/RunWorkflowTask.java b/model/jpa/src/main/java/org/keycloak/models/workflow/RunWorkflowTask.java index e00f2dd7246..5d1009a6fd0 100644 --- a/model/jpa/src/main/java/org/keycloak/models/workflow/RunWorkflowTask.java +++ b/model/jpa/src/main/java/org/keycloak/models/workflow/RunWorkflowTask.java @@ -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); + } + } + } } diff --git a/server-spi-private/src/main/java/org/keycloak/models/workflow/WorkflowExecutor.java b/server-spi-private/src/main/java/org/keycloak/models/workflow/WorkflowExecutor.java index e7e1647114e..497c16a747c 100644 --- a/server-spi-private/src/main/java/org/keycloak/models/workflow/WorkflowExecutor.java +++ b/server-spi-private/src/main/java/org/keycloak/models/workflow/WorkflowExecutor.java @@ -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) { diff --git a/server-spi-private/src/main/java/org/keycloak/models/workflow/WorkflowTask.java b/server-spi-private/src/main/java/org/keycloak/models/workflow/WorkflowTask.java index 2033a0505c8..739872906e0 100644 --- a/server-spi-private/src/main/java/org/keycloak/models/workflow/WorkflowTask.java +++ b/server-spi-private/src/main/java/org/keycloak/models/workflow/WorkflowTask.java @@ -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 future; private long startTime; private AtomicReference 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) { diff --git a/server-spi-private/src/main/java/org/keycloak/models/workflow/WorkflowTransactionalTask.java b/server-spi-private/src/main/java/org/keycloak/models/workflow/WorkflowTransactionalTask.java index 1ae16eeaed5..ef82d923fc8 100644 --- a/server-spi-private/src/main/java/org/keycloak/models/workflow/WorkflowTransactionalTask.java +++ b/server-spi-private/src/main/java/org/keycloak/models/workflow/WorkflowTransactionalTask.java @@ -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 futureCancelled = new AtomicReference(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); + } }