2626import io .serverlessworkflow .impl .lifecycle .WorkflowStatusEvent ;
2727import io .serverlessworkflow .impl .lifecycle .WorkflowSuspendedEvent ;
2828import java .time .Instant ;
29+ import java .util .ArrayList ;
30+ import java .util .Collection ;
2931import java .util .Map ;
3032import java .util .Optional ;
3133import java .util .concurrent .CancellationException ;
@@ -54,6 +56,8 @@ public class WorkflowMutableInstance implements WorkflowInstance {
5456 private Lock statusLock = new ReentrantLock ();
5557 private Map <CompletableFuture <TaskContext >, TaskContext > suspended ;
5658
59+ private Collection <CompletableFuture <?>> cancelables = new ArrayList <>();
60+
5761 protected WorkflowMutableInstance (WorkflowDefinition definition , String id , WorkflowModel input ) {
5862 this .id = id ;
5963 this .input = input ;
@@ -280,7 +284,8 @@ public boolean cancel() {
280284 try {
281285 statusLock .lock ();
282286 if (TaskExecutorHelper .isActive (status .get ())) {
283- futureRef .get ().cancel (true );
287+ cancelables .forEach (t -> t .cancel (true ));
288+ cancelables .clear ();
284289 status (WorkflowStatus .CANCELLED );
285290 publishEvent (
286291 workflowContext ,
@@ -294,6 +299,29 @@ public boolean cancel() {
294299 }
295300 }
296301
302+ public void addCancelable (CompletableFuture <?> cancelable ) {
303+ try {
304+ statusLock .lock ();
305+ if (status .get () == WorkflowStatus .CANCELLED ) {
306+ statusLock .unlock ();
307+ cancelable .cancel (true );
308+ } else {
309+ cancelables .add (cancelable );
310+ statusLock .unlock ();
311+ cancelable .thenAccept (
312+ __ -> {
313+ try {
314+ cancelables .remove (cancelable );
315+ } finally {
316+ statusLock .unlock ();
317+ }
318+ });
319+ }
320+ } finally {
321+ statusLock .unlock ();
322+ }
323+ }
324+
297325 public <T > T additionalObject (String key , Supplier <T > supplier ) {
298326 return (T ) additionalObjects .computeIfAbsent (key , k -> supplier .get ());
299327 }
0 commit comments