diff --git a/temporal-sdk/src/main/java/io/temporal/activity/ActivityCancellationToken.java b/temporal-sdk/src/main/java/io/temporal/activity/ActivityCancellationToken.java deleted file mode 100644 index 612aeaae2..000000000 --- a/temporal-sdk/src/main/java/io/temporal/activity/ActivityCancellationToken.java +++ /dev/null @@ -1,51 +0,0 @@ -package io.temporal.activity; - -import io.temporal.client.ActivityCanceledException; -import io.temporal.common.Experimental; -import java.util.concurrent.CompletableFuture; - -/** Token that allows an Activity implementation to observe cancellation requests. */ -@Experimental -public interface ActivityCancellationToken { - - ActivityCancellationToken NONE = - new ActivityCancellationToken() { - @Override - public boolean isCancellationRequested() { - return false; - } - - @Override - public void throwIfCancellationRequested() throws ActivityCanceledException {} - - @Override - public CompletableFuture getCancellationFuture() { - return new CompletableFuture<>(); - } - }; - - /** - * Returns true after cancellation has been requested for this Activity Execution. - * - *

If this method returns true, the Activity implementation should stop its work and usually - * call {@link #throwIfCancellationRequested()} to report successful cancellation to Temporal. - */ - boolean isCancellationRequested(); - - /** - * Throws {@link ActivityCanceledException} if cancellation has been requested for this Activity - * Execution. - * - *

Rethrowing this exception from Activity code reports successful cancellation to Temporal. - */ - void throwIfCancellationRequested() throws ActivityCanceledException; - - /** - * Future that completes exceptionally with {@link ActivityCanceledException} when cancellation - * has been requested for this Activity Execution. - * - *

Activity code should still call {@link #throwIfCancellationRequested()} or otherwise report - * cancellation if it wants the Activity Execution to complete as canceled. - */ - CompletableFuture getCancellationFuture(); -} diff --git a/temporal-sdk/src/main/java/io/temporal/activity/ActivityExecutionContext.java b/temporal-sdk/src/main/java/io/temporal/activity/ActivityExecutionContext.java index 0d8c1be79..ac656dfa1 100644 --- a/temporal-sdk/src/main/java/io/temporal/activity/ActivityExecutionContext.java +++ b/temporal-sdk/src/main/java/io/temporal/activity/ActivityExecutionContext.java @@ -1,8 +1,10 @@ package io.temporal.activity; import com.uber.m3.tally.Scope; +import io.temporal.client.ActivityCanceledException; import io.temporal.client.ActivityCompletionException; import io.temporal.client.WorkflowClient; +import io.temporal.common.CancellationToken; import io.temporal.common.Experimental; import io.temporal.serviceclient.WorkflowServiceStubsOptions; import io.temporal.worker.WorkerOptions; @@ -95,7 +97,7 @@ public interface ActivityExecutionContext { * recording Heartbeats. */ @Experimental - ActivityCancellationToken getCancellationToken(); + CancellationToken getCancellationToken(); /** * If this method is called during an Activity Execution then the Activity Execution is not going diff --git a/temporal-sdk/src/main/java/io/temporal/common/CancellationToken.java b/temporal-sdk/src/main/java/io/temporal/common/CancellationToken.java new file mode 100644 index 000000000..0c07c0a92 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/common/CancellationToken.java @@ -0,0 +1,77 @@ +package io.temporal.common; + +import java.util.concurrent.CompletableFuture; + +/** + * Token that allows asynchronous code to observe cancellation requests. + * + * @param the exception type surfaced by {@link #throwIfCancellationRequested()} and {@link + * #getCancellationFuture()} when cancellation is requested + */ +@Experimental +public interface CancellationToken { + + /** Returns true after cancellation has been requested. */ + boolean isCancellationRequested(); + + /** Throws {@code E} if cancellation has been requested. */ + void throwIfCancellationRequested() throws E; + + /** + * Future that completes exceptionally with {@code E} when cancellation has been requested. + * + *

Code waiting on external work can chain off this future to abort in-flight requests when + * cancellation is requested. + */ + default CompletableFuture getCancellationFuture() { + CompletableFuture result = new CompletableFuture<>(); + Registration registration = + onCancel( + () -> { + try { + throwIfCancellationRequested(); + } catch (RuntimeException e) { + result.completeExceptionally(e); + } + }); + result.whenComplete((ignored, error) -> registration.close()); + return result; + } + + /** + * Registers a callback to run when cancellation is requested, or immediately if already + * cancelled. + * + * @return a handle that removes the callback if cancellation has not happened yet. + */ + Registration onCancel(Runnable callback); + + /** Handle for removing a previously registered cancellation callback. */ + interface Registration extends AutoCloseable { + @Override + void close(); + } + + /** A token that is never cancelled. */ + static CancellationToken none() { + return new CancellationToken() { + @Override + public boolean isCancellationRequested() { + return false; + } + + @Override + public void throwIfCancellationRequested() {} + + @Override + public CompletableFuture getCancellationFuture() { + return new CompletableFuture<>(); + } + + @Override + public Registration onCancel(Runnable callback) { + return () -> {}; + } + }; + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/common/interceptors/ActivityExecutionContextBase.java b/temporal-sdk/src/main/java/io/temporal/common/interceptors/ActivityExecutionContextBase.java index 3ce86cec3..adf2e8353 100644 --- a/temporal-sdk/src/main/java/io/temporal/common/interceptors/ActivityExecutionContextBase.java +++ b/temporal-sdk/src/main/java/io/temporal/common/interceptors/ActivityExecutionContextBase.java @@ -1,12 +1,13 @@ package io.temporal.common.interceptors; import com.uber.m3.tally.Scope; -import io.temporal.activity.ActivityCancellationToken; import io.temporal.activity.ActivityExecutionContext; import io.temporal.activity.ActivityInfo; import io.temporal.activity.ManualActivityCompletionClient; +import io.temporal.client.ActivityCanceledException; import io.temporal.client.ActivityCompletionException; import io.temporal.client.WorkflowClient; +import io.temporal.common.CancellationToken; import java.lang.reflect.Type; import java.util.Optional; @@ -54,7 +55,7 @@ public byte[] getTaskToken() { } @Override - public ActivityCancellationToken getCancellationToken() { + public CancellationToken getCancellationToken() { return next.getCancellationToken(); } diff --git a/temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityCancellationTokenImpl.java b/temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityCancellationTokenImpl.java deleted file mode 100644 index b78fcdafa..000000000 --- a/temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityCancellationTokenImpl.java +++ /dev/null @@ -1,51 +0,0 @@ -package io.temporal.internal.activity; - -import io.temporal.activity.ActivityCancellationToken; -import io.temporal.client.ActivityCanceledException; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionException; - -final class ActivityCancellationTokenImpl implements ActivityCancellationToken { - private final CompletableFuture cancellationFuture = new CompletableFuture<>(); - private volatile ActivityCanceledException cancellationException; - - @Override - public boolean isCancellationRequested() { - return cancellationException != null; - } - - @Override - public void throwIfCancellationRequested() throws ActivityCanceledException { - ActivityCanceledException exception = cancellationException; - if (exception != null) { - throw exception; - } - } - - @Override - public CompletableFuture getCancellationFuture() { - CompletableFuture result = new CompletableFuture<>(); - cancellationFuture.whenComplete( - (ignored, exception) -> { - if (exception == null) { - result.complete(null); - } else { - result.completeExceptionally(unwrapCompletionException(exception)); - } - }); - return result; - } - - synchronized void requestCancel(ActivityCanceledException exception) { - if (cancellationException == null) { - cancellationException = exception; - cancellationFuture.completeExceptionally(exception); - } - } - - private static Throwable unwrapCompletionException(Throwable exception) { - return exception instanceof CompletionException && exception.getCause() != null - ? exception.getCause() - : exception; - } -} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityExecutionContextImpl.java b/temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityExecutionContextImpl.java index db8943138..40fe45c32 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityExecutionContextImpl.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityExecutionContextImpl.java @@ -1,12 +1,13 @@ package io.temporal.internal.activity; import com.uber.m3.tally.Scope; -import io.temporal.activity.ActivityCancellationToken; import io.temporal.activity.ActivityExecutionContext; import io.temporal.activity.ActivityInfo; import io.temporal.activity.ManualActivityCompletionClient; +import io.temporal.client.ActivityCanceledException; import io.temporal.client.ActivityCompletionException; import io.temporal.client.WorkflowClient; +import io.temporal.common.CancellationToken; import io.temporal.common.converter.DataConverter; import io.temporal.internal.client.external.ManualActivityCompletionClientFactory; import io.temporal.payload.context.ActivitySerializationContext; @@ -110,7 +111,7 @@ public byte[] getTaskToken() { } @Override - public ActivityCancellationToken getCancellationToken() { + public CancellationToken getCancellationToken() { return heartbeatContext.getCancellationToken(); } diff --git a/temporal-sdk/src/main/java/io/temporal/internal/activity/HeartbeatContext.java b/temporal-sdk/src/main/java/io/temporal/internal/activity/HeartbeatContext.java index c960f70ba..faef0d795 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/activity/HeartbeatContext.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/activity/HeartbeatContext.java @@ -1,7 +1,8 @@ package io.temporal.internal.activity; -import io.temporal.activity.ActivityCancellationToken; +import io.temporal.client.ActivityCanceledException; import io.temporal.client.ActivityCompletionException; +import io.temporal.common.CancellationToken; import java.lang.reflect.Type; import java.util.Optional; @@ -24,7 +25,7 @@ interface HeartbeatContext { Object getLatestHeartbeatDetails(); - ActivityCancellationToken getCancellationToken(); + CancellationToken getCancellationToken(); /** Mark this activity as canceled by an external worker command. */ void cancelFromWorkerCommand(); diff --git a/temporal-sdk/src/main/java/io/temporal/internal/activity/HeartbeatContextImpl.java b/temporal-sdk/src/main/java/io/temporal/internal/activity/HeartbeatContextImpl.java index 477780a97..91da94ab0 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/activity/HeartbeatContextImpl.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/activity/HeartbeatContextImpl.java @@ -3,16 +3,17 @@ import com.uber.m3.tally.Scope; import io.grpc.Status; import io.grpc.StatusRuntimeException; -import io.temporal.activity.ActivityCancellationToken; import io.temporal.activity.ActivityExecutionContext; import io.temporal.activity.ActivityInfo; import io.temporal.api.common.v1.Payloads; import io.temporal.api.enums.v1.TimeoutType; import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatResponse; import io.temporal.client.*; +import io.temporal.common.CancellationToken; import io.temporal.common.converter.DataConverter; import io.temporal.failure.TimeoutFailure; import io.temporal.internal.client.ActivityClientHelper; +import io.temporal.internal.concurrent.structured.CancelSource; import io.temporal.payload.context.ActivitySerializationContext; import io.temporal.serviceclient.WorkflowServiceStubs; import java.lang.reflect.Type; @@ -76,8 +77,8 @@ static long getLocalHeartbeatTimeoutBufferMillis() { private boolean rejectNewHeartbeats; private ActivityCompletionException lastException; - private final ActivityCancellationTokenImpl cancellationToken = - new ActivityCancellationTokenImpl(); + private final CancelSource cancellationSource = + new CancelSource<>(ActivityCanceledException::new); public HeartbeatContextImpl( WorkflowServiceStubs service, @@ -166,7 +167,7 @@ public void heartbeat(V details) throws ActivityCompletionException { if (lastException != null) { throw lastException; } - cancellationToken.throwIfCancellationRequested(); + cancellationSource.token().throwIfCancellationRequested(); } finally { lock.unlock(); } @@ -257,8 +258,8 @@ public void asyncCompletionStarted() { } @Override - public ActivityCancellationToken getCancellationToken() { - return cancellationToken; + public CancellationToken getCancellationToken() { + return cancellationSource.token(); } private void doHeartBeatLocked(Object details) { @@ -363,7 +364,7 @@ private void sendHeartbeatRequest(Object details) { private void requestCancelLocked() { ActivityCanceledException exception = new ActivityCanceledException(info); lastException = exception; - cancellationToken.requestCancel(exception); + cancellationSource.cancel(exception); } private static long getHeartbeatIntervalMs( diff --git a/temporal-sdk/src/main/java/io/temporal/internal/activity/LocalActivityExecutionContextImpl.java b/temporal-sdk/src/main/java/io/temporal/internal/activity/LocalActivityExecutionContextImpl.java index 0f6636424..0e8282c58 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/activity/LocalActivityExecutionContextImpl.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/activity/LocalActivityExecutionContextImpl.java @@ -1,11 +1,12 @@ package io.temporal.internal.activity; import com.uber.m3.tally.Scope; -import io.temporal.activity.ActivityCancellationToken; import io.temporal.activity.ActivityInfo; import io.temporal.activity.ManualActivityCompletionClient; +import io.temporal.client.ActivityCanceledException; import io.temporal.client.ActivityCompletionException; import io.temporal.client.WorkflowClient; +import io.temporal.common.CancellationToken; import java.lang.reflect.Type; import java.util.Optional; @@ -59,8 +60,8 @@ public byte[] getTaskToken() { } @Override - public ActivityCancellationToken getCancellationToken() { - return ActivityCancellationToken.NONE; + public CancellationToken getCancellationToken() { + return CancellationToken.none(); } @Override diff --git a/temporal-sdk/src/main/java/io/temporal/internal/common/ListUtils.java b/temporal-sdk/src/main/java/io/temporal/internal/common/ListUtils.java new file mode 100644 index 000000000..ae70a08ab --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/common/ListUtils.java @@ -0,0 +1,18 @@ +package io.temporal.internal.common; + +import java.util.ArrayList; +import java.util.List; + +public final class ListUtils { + + private ListUtils() {} + + /** Concatenates a list of lists into a single list, preserving order. */ + public static List flatten(List> lists) { + List result = new ArrayList<>(); + for (List list : lists) { + result.addAll(list); + } + return result; + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/AsyncTask.java b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/AsyncTask.java new file mode 100644 index 000000000..230eeb147 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/AsyncTask.java @@ -0,0 +1,63 @@ +package io.temporal.internal.concurrent.structured; + +import io.temporal.common.CancellationToken; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.function.Consumer; +import java.util.function.Function; + +/** + * Internal task handle over a {@link CompletableFuture} that manages cancellation and tracks + * derived tasks. + * + *

Cancellation is downstream by default. {@link #cancel()} settles this task + * and every task derived from it (its {@code map}/{@code recover} children). It does not + * cancel the task this one was derived from. + * + * @param the result type + */ +interface AsyncTask extends TaskChain { + + @Override + AsyncTask map(Function fn); + + @Override + AsyncTask recover(Function fn); + + @Override + default AsyncTask thenAccept(Consumer fn) { + return map( + value -> { + fn.accept(value); + return null; + }); + } + + /** + * Cancels this task and everything derived from it. + * + * @return {@code true} if this call initiated cancellation (the task had not already settled). + */ + boolean cancel(); + + boolean isDone(); + + boolean isCancelled(); + + /** + * Blocks for the value; throws on failure, or {@link java.util.concurrent.CancellationException} + * on cancel. + */ + T join(); + + /** Blocks until settled and returns the outcome as a {@link Result}; never throws. */ + Result joinSettled(); + + /** + * @return the read-only cancellation token for this task. + */ + CancellationToken token(); + + /** Escape hatch to the underlying future for interop with existing APIs. */ + CompletableFuture toCompletableFuture(); +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/CancelSource.java b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/CancelSource.java new file mode 100644 index 000000000..7cfd5e277 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/CancelSource.java @@ -0,0 +1,117 @@ +package io.temporal.internal.concurrent.structured; + +import io.temporal.common.CancellationToken; +import java.util.ArrayList; +import java.util.List; +import java.util.function.Supplier; + +/** + * The write side of cancellation. Whoever holds the {@code CancelSource} can request + * cancellation while everyone else observes using {@link #token()}. + * + *

Sources can be linked: a source created via {@link #linkedTo} is cancelled + * automatically when any of its parent tokens is cancelled, without giving the child the power to + * cancel the parent. + * + * @param the exception surfaced by the token when cancellation is requested + */ +public final class CancelSource { + + private final Object lock = new Object(); + private final Supplier defaultException; + private volatile boolean cancelled = false; + private volatile E cancellationException; + + /** Pending callbacks. Set to {@code null} by {@link #cancel} once it takes ownership. */ + private List callbacks = new ArrayList<>(); + + public CancelSource(Supplier defaultException) { + this.defaultException = defaultException; + } + + private final CancellationToken token = + new CancellationToken() { + @Override + public boolean isCancellationRequested() { + return cancelled; + } + + @Override + public void throwIfCancellationRequested() { + if (cancelled) { + throw cancellationException; + } + } + + @Override + public Registration onCancel(Runnable cb) { + synchronized (lock) { + if (!cancelled) { + callbacks.add(cb); + return () -> { + synchronized (lock) { + if (callbacks != null) { + callbacks.remove(cb); + } + } + }; + } + } + runSafely(cb); + return () -> {}; + } + }; + + /** The read-only token to hand to code that observes cancellation. */ + public CancellationToken token() { + return token; + } + + public boolean isCancelled() { + return cancelled; + } + + /** Requests cancellation with a freshly created exception. */ + public void cancel() { + cancel(defaultException.get()); + } + + /** Requests cancellation, surfacing {@code exception} from the token. Idempotent. */ + public void cancel(E exception) { + List toRun; + synchronized (lock) { + if (cancelled) { + return; + } + cancellationException = exception; + cancelled = true; + toRun = callbacks; + callbacks = null; + } + for (Runnable cb : toRun) { + runSafely(cb); + } + } + + private static void runSafely(Runnable cb) { + try { + cb.run(); + } catch (Throwable ignored) { + /* a bad callback must not block others */ + } + } + + /** + * Creates a source whose token is cancelled (via {@link #cancel()}) whenever any {@code parent} + * token is cancelled. + */ + @SafeVarargs + public static CancelSource linkedTo( + Supplier defaultException, CancellationToken... parents) { + CancelSource s = new CancelSource<>(defaultException); + for (CancellationToken p : parents) { + p.onCancel(s::cancel); + } + return s; + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/DefaultAsyncTask.java b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/DefaultAsyncTask.java new file mode 100644 index 000000000..f48576343 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/DefaultAsyncTask.java @@ -0,0 +1,184 @@ +package io.temporal.internal.concurrent.structured; + +import io.temporal.common.CancellationToken; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutionException; +import java.util.function.BiConsumer; +import java.util.function.Consumer; +import java.util.function.Function; + +/** Reference implementation of {@link AsyncTask}. */ +final class DefaultAsyncTask implements AsyncTask { + + final CompletableFuture cf; + final CancelSource source; + private final CompletableFuture terminated; + private final Runnable cancellationHook; + private final Consumer> taskRegistrar; + private final BiConsumer, DefaultAsyncTask> resultRegistrar; + private final List> children = new CopyOnWriteArrayList<>(); + + DefaultAsyncTask( + CompletableFuture cf, + CancelSource source, + Runnable cancellationHook, + Consumer> taskRegistrar, + BiConsumer, DefaultAsyncTask> resultRegistrar) { + this.cf = cf; + this.source = source; + this.cancellationHook = cancellationHook; + this.taskRegistrar = taskRegistrar; + this.resultRegistrar = resultRegistrar; + this.terminated = cf.handle((v, e) -> (Void) null); + source + .token() + .onCancel( + () -> { + if (cancellationHook != null) { + cancellationHook.run(); + } + cf.completeExceptionally(new CancellationException()); + }); + } + + @Override + public AsyncTask map(Function fn) { + CompletableFuture next = new CompletableFuture<>(); + cf.whenComplete( + (value, error) -> { + if (error != null) { + next.completeExceptionally(error); + return; + } + if (source.token().isCancellationRequested()) { + next.completeExceptionally(new CancellationException()); + return; + } + try { + next.complete(fn.apply(value)); + } catch (Throwable t) { + next.completeExceptionally(t); + } + }); + return derive(next); + } + + @Override + public AsyncTask recover(Function fn) { + CompletableFuture next = new CompletableFuture<>(); + cf.whenComplete( + (value, error) -> { + if (error == null) { + if (source.token().isCancellationRequested()) { + next.completeExceptionally(new CancellationException()); + } else { + next.complete(value); + } + return; + } + + Throwable unwrapped = unwrap(error); + if (unwrapped instanceof CancellationException) { + next.completeExceptionally(unwrapped); + return; + } + + try { + next.complete(fn.apply(unwrapped)); + } catch (Throwable t) { + next.completeExceptionally(t); + } + }); + return derive(next); + } + + private DefaultAsyncTask derive(CompletableFuture next) { + DefaultAsyncTask child = + new DefaultAsyncTask<>( + next, + CancelSource.linkedTo(CancellationException::new, source.token()), + null, + taskRegistrar, + resultRegistrar); + children.add(child); + taskRegistrar.accept(child); + resultRegistrar.accept(this, child); + return child; + } + + @Override + public boolean cancel() { + boolean first = !cf.isDone(); + source.cancel(); + cf.completeExceptionally(new CancellationException()); + for (DefaultAsyncTask c : children) c.cancel(); + return first; + } + + @Override + public boolean isDone() { + return cf.isDone(); + } + + @Override + public boolean isCancelled() { + if (!cf.isDone()) { + return source.isCancelled(); + } + return joinSettled().isCancelled(); + } + + @Override + public T join() { + try { + return cf.join(); + } catch (CompletionException e) { + throw rethrow(unwrap(e)); + } + } + + @Override + public Result joinSettled() { + try { + return Result.success(cf.join()); + } catch (CancellationException e) { + return Result.cancelled(); + } catch (CompletionException e) { + Throwable c = unwrap(e); + return (c instanceof CancellationException) ? Result.cancelled() : Result.failure(c); + } + } + + @Override + public CancellationToken token() { + return source.token(); + } + + @Override + public CompletableFuture toCompletableFuture() { + return cf; + } + + /** Completes (normally, never exceptionally) once this task's future has settled. */ + CompletableFuture terminated() { + return terminated; + } + + static Throwable unwrap(Throwable t) { + while ((t instanceof CompletionException || t instanceof ExecutionException) + && t.getCause() != null) { + t = t.getCause(); + } + return t; + } + + static RuntimeException rethrow(Throwable t) { + if (t instanceof RuntimeException) return (RuntimeException) t; + if (t instanceof Error) throw (Error) t; + return new CompletionException(t); + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/DefaultTaskScope.java b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/DefaultTaskScope.java new file mode 100644 index 000000000..d330882c2 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/DefaultTaskScope.java @@ -0,0 +1,223 @@ +package io.temporal.internal.concurrent.structured; + +import io.temporal.common.CancellationToken; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import java.util.function.Function; +import javax.annotation.Nonnull; + +/** + * Reference implementation of {@link TaskScope}. Each task's {@link CancelSource} is linked to the + * scope's token, so {@link #cancelAll()} trips every task at once. + */ +final class DefaultTaskScope implements TaskScope { + + private final CancelSource scope = + new CancelSource<>(CancellationException::new); + private final AtomicBoolean closed = new AtomicBoolean(false); + private final List> ownedTasks = new CopyOnWriteArrayList<>(); + private final List> resultTasks = new CopyOnWriteArrayList<>(); + + @Override + public CancellationToken token() { + return scope.token(); + } + + @Override + public AsyncTask attach(@Nonnull CompletableFuture future) { + Objects.requireNonNull(future, "future"); + ensureOpen(); + CancelSource childSrc = + CancelSource.linkedTo(CancellationException::new, scope.token()); + DefaultAsyncTask task = + new DefaultAsyncTask<>( + future, childSrc, () -> future.cancel(true), ownedTasks::add, this::replaceResultTask); + ownedTasks.add(task); + resultTasks.add(task); + return task; + } + + private void replaceResultTask(DefaultAsyncTask parent, DefaultAsyncTask child) { + int index = resultTasks.indexOf(parent); + if (index >= 0) { + resultTasks.set(index, child); + } else { + resultTasks.add(child); + } + } + + private void ensureOpen() { + if (closed.get()) { + throw new IllegalStateException("TaskScope is closed"); + } + } + + @Override + public void cancelAll() { + scope.cancel(); + } + + @Override + public CompletableFuture awaitAll(Function, R> resultTransformer) { + CompletableFuture result = new CompletableFuture<>(); + AtomicReference errorRef = new AtomicReference<>(); + awaitTermination(0, failFastWatch(errorRef)) + .whenComplete( + (ignored, terminationError) -> { + Throwable error = errorRef.get(); + if (error != null) { + result.completeExceptionally(error); + return; + } + try { + List> collected = resultTaskSnapshot(); + List values = new ArrayList<>(collected.size()); + for (DefaultAsyncTask task : collected) { + values.add(task.join()); + } + result.complete(resultTransformer.apply(values)); + } catch (Throwable t) { + result.completeExceptionally(t); + } + }); + propagateCancellation(result); + return result; + } + + @Override + public CompletableFuture>> awaitAllSettled() { + CompletableFuture>> result = new CompletableFuture<>(); + // All-settled never fails fast, so it collects outcomes without cancelling siblings. + awaitTermination(0, null) + .whenComplete( + (ignored, terminationError) -> { + List> collected = resultTaskSnapshot(); + List> results = new ArrayList<>(collected.size()); + for (DefaultAsyncTask task : collected) { + results.add(task.joinSettled()); + } + result.complete(results); + }); + propagateCancellation(result); + return result; + } + + /** + * Completes once every owned task has settled. The owned-task list grows as tasks are derived, so + * each round re-checks and waits for any that appeared meanwhile. Never completes exceptionally. + * {@code onDiscover} runs once per task as it is first observed. + */ + private CompletableFuture awaitTermination( + int from, Consumer> onDiscover) { + List> snapshot = new ArrayList<>(ownedTasks); + if (from >= snapshot.size()) { + return CompletableFuture.completedFuture(null); + } + CompletableFuture[] terminations = new CompletableFuture[snapshot.size() - from]; + for (int i = from; i < snapshot.size(); i++) { + DefaultAsyncTask task = snapshot.get(i); + if (onDiscover != null) { + onDiscover.accept(task); + } + terminations[i - from] = task.terminated(); + } + int next = snapshot.size(); + return CompletableFuture.allOf(terminations) + .thenCompose(ignored -> awaitTermination(next, onDiscover)); + } + + /** + * Records the first failure and cancels the whole sibling group. Registered on each task as it is + * discovered, so tasks derived mid-wait fail fast too. + */ + private Consumer> failFastWatch(AtomicReference errorRef) { + return task -> + task.toCompletableFuture() + .whenComplete( + (ignored, error) -> { + if (error != null + && errorRef.compareAndSet(null, DefaultAsyncTask.unwrap(error))) { + cancelAll(); + } + }); + } + + private void propagateCancellation(CompletableFuture result) { + result.whenComplete( + (ignored, error) -> { + if (result.isCancelled()) { + cancelAll(); + } + }); + } + + @SuppressWarnings("unchecked") + private List> resultTaskSnapshot() { + return (List>) (List) new ArrayList<>(resultTasks); + } + + @Override + public void close() { + if (!closed.compareAndSet(false, true)) { + return; + } + // Cancelling settles every attached and derived future synchronously; there are no threads to + // wait for. + cancelAll(); + } + + /** + * Non-blocking completion for {@link TaskScope#withScope}: when {@code body} settles, cancels the + * group on failure, waits for every task to settle, then delivers the outcome. The returned + * future does not settle until all tasks have, so no task outlives the scope. + */ + CompletableFuture closeWhenDone(CompletableFuture body) { + CompletableFuture delivered = new CompletableFuture<>(); + body.whenComplete( + (value, error) -> { + if (error != null) { + cancelAll(); + } + allTerminated() + .whenComplete( + (ignored, terminationError) -> { + closed.set(true); + if (error != null) { + delivered.completeExceptionally(DefaultAsyncTask.unwrap(error)); + } else { + delivered.complete(value); + } + }); + }); + propagateCancellation(delivered); + return delivered; + } + + private CompletableFuture allTerminated() { + return awaitTermination(0, null); + } + + /** + * Runs {@code body} against {@code scope} and ties the scope's lifetime to the returned future. + */ + static CompletableFuture run( + DefaultTaskScope scope, Function, CompletableFuture> body) { + CompletableFuture future; + try { + future = Objects.requireNonNull(body.apply(scope), "future"); + } catch (Throwable t) { + scope.close(); + CompletableFuture failed = new CompletableFuture<>(); + failed.completeExceptionally(t); + return failed; + } + return scope.closeWhenDone(future); + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/README.md b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/README.md new file mode 100644 index 000000000..33be27804 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/README.md @@ -0,0 +1,163 @@ +# Structured Concurrency + +Minimal, non-blocking, "structured concurrency" wrapper for `CompletableFuture`. No executors or +thread management. Just a wrapper that provides correctness guarantees and convenience. + +## Usage + +```java +CompletableFuture> result = + TaskScope.withScope( // 1. create a scope + scope -> { + scope.attach(startWork()); // 2. add tasks to the scope + scope.attach(startMoreWork()) + .map(value -> doDownstreamWork(value)); + + return scope.awaitAll(); // 3. collect results + }); + +result.get(); // 4. get results a usual +// or +result.cancel(true); // 5. cancel (propagates to all child tasks) +``` + +What it provides: + +1. A `withScope` boundary whose result waits for every attached task to settle before completing + normally or exceptionally, so no task outlives the scope. +2. Fan-in over a group of async operations: `awaitAll` (results in order), `awaitAll(transformer)` + (reshape the group), and `awaitAllSettled` (per-task `Result`). +3. Fail-fast: the first failure completes the scope's result with that error and cancels the rest. +4. Scope cancellation (`cancelAll`, or cancelling the returned future) propagates to every attached + task; task-chain cancellation propagates from parent stages to derived stages. +5. Cooperative cancellation via `CancellationToken`. +6. `TaskScope` methods return a `TaskChain`, not raw task handles, so tasks cannot be joined, + cancelled, or converted back to futures outside the owning scope. + +## Execution model + +`TaskScope` is an ownership boundary, not an executor. It owns task lifetime, cancellation, result +collection, and the guarantee that scoped work settles before the returned future does. It does not fork +new threads, use executors under the hood, or create any virtual thread abstractions. It's just a wrapper +over CompletableFuture. + +## Why + +```java +List> batches = getPayloadBatches(); +CompletableFuture> result = + TaskScope.withScope( scope -> { + StorageDriverStoreContext context = new StorageDriverStoreContextImpl(target, scope.token()); + + for (Batch batch : batches) { + scope + .attach(batch.driver.store(context, batch.values())) + .map(claims -> createReferencePayloads(batch, claims)); + } + return scope.awaitAll(ListUtils::flatten); + }); + +result.cancel(true); +``` + +The equivalent code using `CompletableFuture` directly has to rebuild the same ownership and cancellation rules by hand: + +```java +List> batches = getPayloadBatches(); +CancelSource cancellation = new CancelSource(); +StorageDriverStoreContext context = + new StorageDriverStoreContextImpl(target, cancellation.token()); + +List> upstream = new ArrayList<>(); +List>> downstream = new ArrayList<>(); +CompletableFuture> result = new CompletableFuture<>(); + +for (Batch batch : batches) { + CompletableFuture> storeFuture = + batch.driver.store(context, batch.values()); + upstream.add(storeFuture); + + CompletableFuture> downstreamFuture = + storeFuture.thenApply(claims -> createReferencePayloads(batch, claims)); + downstream.add(downstreamFuture); +} + +if (downstream.isEmpty()) { + result.complete(Collections.emptyList()); + return result; +} + +AtomicBoolean completed = new AtomicBoolean(false); +AtomicInteger remaining = new AtomicInteger(downstream.size()); + +Runnable cancelTracked = + () -> { + cancellation.cancel(); + for (CompletableFuture upstreamFuture : upstream) { + upstreamFuture.cancel(true); + } + for (CompletableFuture downstreamFuture : downstream) { + downstreamFuture.cancel(true); + } + }; + +Supplier> allTrackedSettled = + () -> { + List> tracked = new ArrayList<>(upstream.size() + downstream.size()); + tracked.addAll(upstream); + tracked.addAll(downstream); + return CompletableFuture + .allOf(tracked.toArray(new CompletableFuture[tracked.size()])) + .handle((ignored, terminationError) -> null); + }; + +for (CompletableFuture> downstreamFuture : downstream) { + downstreamFuture.whenComplete( + (ignored, err) -> { + if (err != null) { + if (completed.compareAndSet(false, true)) { + cancelTracked.run(); + allTrackedSettled + .get() + .whenComplete( + (unused, terminationError) -> result.completeExceptionally(err)); + } + return; + } + + if (remaining.decrementAndGet() == 0 && completed.compareAndSet(false, true)) { + allTrackedSettled + .get() + .whenComplete( + (unused, terminationError) -> { + try { + List> values = new ArrayList<>(); + for (CompletableFuture> completedDownstream : downstream) { + values.add(completedDownstream.join()); + } + result.complete(ListUtils.flatten(values)); + } catch (Throwable resultError) { + result.completeExceptionally(resultError); + } + }); + } + }); +} + +result.whenComplete( + (ignored, err) -> { + if (result.isCancelled() && completed.compareAndSet(false, true)) { + cancelTracked.run(); + } + }); + +return result; +``` + +With direct `CompletableFuture` composition, fan-out/fan-in code usually degrades into: + +- ad-hoc cancellation wiring per call site +- separate upstream and downstream tracking so scope cancellation reaches both driver work and derived stages +- manual propagation to cooperative cancellation tokens as well as `CompletableFuture.cancel(true)` +- no built-in parent boundary that owns the child set +- duplicated fail-fast, all-settled, and group-termination orchestration logic diff --git a/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/Result.java b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/Result.java new file mode 100644 index 000000000..c06f2de08 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/Result.java @@ -0,0 +1,105 @@ +package io.temporal.internal.concurrent.structured; + +import java.util.concurrent.CancellationException; +import java.util.function.Function; + +/** + * A settled outcome of a task: {@code SUCCESS}, {@code FAILURE}, or {@code CANCELLED}. + * + *

This is the element type returned by {@link TaskScope#awaitAllSettled}. + * + * @param the success value type + */ +public final class Result { + + public enum Status { + SUCCESS, + FAILURE, + CANCELLED + } + + private final Status status; + private final T value; + private final Throwable cause; + + private Result(Status status, T value, Throwable cause) { + this.status = status; + this.value = value; + this.cause = cause; + } + + public static Result success(T value) { + return new Result<>(Status.SUCCESS, value, null); + } + + public static Result failure(Throwable t) { + return new Result<>(Status.FAILURE, null, t); + } + + public static Result cancelled() { + return new Result<>(Status.CANCELLED, null, new CancellationException()); + } + + public Status status() { + return status; + } + + public boolean isSuccess() { + return status == Status.SUCCESS; + } + + public boolean isFailure() { + return status == Status.FAILURE; + } + + public boolean isCancelled() { + return status == Status.CANCELLED; + } + + /** + * @return the value on success; otherwise throws {@link IllegalStateException} (cause attached). + */ + public T get() { + if (status == Status.SUCCESS) return value; + throw new IllegalStateException("Result.get() attempted on unsuccessful result", cause); + } + + /** + * @return the value on success, else {@code other}. + */ + public T orElse(T other) { + return status == Status.SUCCESS ? value : other; + } + + /** + * @return the throwable for FAILURE/CANCELLED, or {@code null} for SUCCESS. + */ + public Throwable cause() { + return cause; + } + + /** Maps the success value, leaving FAILURE/CANCELLED untouched. */ + @SuppressWarnings("unchecked") + public Result map(Function fn) { + return status == Status.SUCCESS ? Result.success(fn.apply(value)) : (Result) this; + } + + /** Collapses both branches into one value. CANCELLED is routed to {@code onFailure}. */ + public R fold( + Function onSuccess, + Function onFailure) { + return status == Status.SUCCESS ? onSuccess.apply(value) : onFailure.apply(cause); + } + + @Override + public String toString() { + switch (status) { + case SUCCESS: + return "Success(" + value + ")"; + case CANCELLED: + return "Cancelled"; + default: + return "Failure(" + cause + ")"; + } + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/TaskChain.java b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/TaskChain.java new file mode 100644 index 000000000..3557779a0 --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/TaskChain.java @@ -0,0 +1,27 @@ +package io.temporal.internal.concurrent.structured; + +import java.util.function.Consumer; +import java.util.function.Function; + +/** A scope-owned continuation chain. Results are observed through the owning {@link TaskScope}. */ +public interface TaskChain { + + /** Transforms the success value once available. Failures/cancellation pass through untouched. */ + TaskChain map(Function fn); + + /** + * Supplies a fallback value if this chain fails. Unlike {@code CompletableFuture.exceptionally}, + * a cancelled chain is not recovered (cancellation propagates) and the fallback receives the + * unwrapped cause. + */ + TaskChain recover(Function fn); + + /** Runs a side effect on success and yields a {@code Void} chain. */ + default TaskChain thenAccept(Consumer fn) { + return map( + value -> { + fn.accept(value); + return null; + }); + } +} diff --git a/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/TaskScope.java b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/TaskScope.java new file mode 100644 index 000000000..11628503d --- /dev/null +++ b/temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/TaskScope.java @@ -0,0 +1,85 @@ +package io.temporal.internal.concurrent.structured; + +import io.temporal.common.CancellationToken; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; +import javax.annotation.Nonnull; + +/** + * A structured-concurrency scope that owns a group of sibling tasks and a shared cancellation + * signal, and guarantees that none of those tasks outlive the scope's completion. + * + * 

{@code
+ * CompletableFuture> report = TaskScope.withScope(scope -> {
+ *     scope.attach(fetchData(1));
+ *     scope.attach(fetchData(2));
+ *
+ *     return scope.awaitAll();
+ * });
+ * }
+ */ +public interface TaskScope extends AutoCloseable { + + /** + * @return the scope-wide cancellation token (tripped by {@link #cancelAll()} or {@link + * #close()}). + */ + CancellationToken token(); + + /** + * Attaches an existing asynchronous task to this scope. Scope cancellation requests cancellation + * on the attached future. + */ + TaskChain attach(@Nonnull CompletableFuture future); + + /** Cancels every task in this scope. */ + void cancelAll(); + + /** + * Non-blocking fail-fast wait that completes with the collected task results in the order the + * tasks were attached. + * + *

This method does not close the scope. Scope lifetime remains caller-owned. + */ + default CompletableFuture> awaitAll() { + return awaitAll(Function.identity()); + } + + /** + * Non-blocking fail-fast wait that collects all task results and passes them to a transformer. + * + *

The returned future completes after all collected tasks complete successfully. Attached + * tasks are collected by default; when a collected task is transformed, the transformed child + * replaces its parent in the collected result list. On first failure or cancellation it completes + * exceptionally and cancels unfinished tasks. + * + *

This method does not close the scope. Scope lifetime remains caller-owned. + */ + CompletableFuture awaitAll(Function, R> resultTransformer); + + /** + * Non-blocking wait that completes with each collected task's settled outcome in the order the + * tasks were attached. + * + *

Task failures and cancellations are returned as {@link Result} values instead of completing + * the returned future exceptionally. + * + *

This method does not close the scope. Scope lifetime remains caller-owned. + */ + CompletableFuture>> awaitAllSettled(); + + /** Cancels all tasks; idempotent. Cancellation settles every attached and derived future. */ + @Override + void close(); + + /** + * Runs async work in a lexical scope and closes the scope when the returned future settles. The + * returned future does not complete until every task attached in {@code body} has settled, so no + * task outlives the scope. + */ + static CompletableFuture withScope(Function, CompletableFuture> body) { + return DefaultTaskScope.run(new DefaultTaskScope<>(), body); + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/internal/activity/HeartbeatContextImplTest.java b/temporal-sdk/src/test/java/io/temporal/internal/activity/HeartbeatContextImplTest.java index 9b04bb5cd..1379aed15 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/activity/HeartbeatContextImplTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/activity/HeartbeatContextImplTest.java @@ -7,7 +7,6 @@ import com.uber.m3.tally.NoopScope; import io.grpc.Status; import io.grpc.StatusRuntimeException; -import io.temporal.activity.ActivityCancellationToken; import io.temporal.activity.ActivityInfo; import io.temporal.api.enums.v1.TimeoutType; import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatRequest; @@ -16,6 +15,7 @@ import io.temporal.client.ActivityCanceledException; import io.temporal.client.ActivityCompletionException; import io.temporal.client.WorkflowClient; +import io.temporal.common.CancellationToken; import io.temporal.common.converter.GlobalDataConverter; import io.temporal.failure.TimeoutFailure; import io.temporal.serviceclient.WorkflowServiceStubs; @@ -367,7 +367,7 @@ private HeartbeatContextImpl createHeartbeatContext( } private static ActivityCanceledException assertCancellationFutureCompletedExceptionally( - ActivityCancellationToken cancellationToken) { + CancellationToken cancellationToken) { CompletableFuture cancellationFuture = cancellationToken.getCancellationFuture(); assertTrue(cancellationFuture.isDone()); assertTrue(cancellationFuture.isCompletedExceptionally()); diff --git a/temporal-sdk/src/test/java/io/temporal/internal/common/ListUtilsTest.java b/temporal-sdk/src/test/java/io/temporal/internal/common/ListUtilsTest.java new file mode 100644 index 000000000..08cabd020 --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/common/ListUtilsTest.java @@ -0,0 +1,75 @@ +package io.temporal.internal.common; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertTrue; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedList; +import java.util.List; +import org.junit.Test; + +public class ListUtilsTest { + + @Test + public void flattensNestedCollectionsInOrder() { + List flat = + ListUtils.flatten( + Arrays.asList(Arrays.asList(1, 2), Arrays.asList(3), Arrays.asList(4, 5))); + + assertEquals(Arrays.asList(1, 2, 3, 4, 5), flat); + } + + @Test + public void preservesDuplicates() { + List flat = ListUtils.flatten(Arrays.asList(Arrays.asList(1, 1), Arrays.asList(1))); + + assertEquals(Arrays.asList(1, 1, 1), flat); + } + + @Test + public void skipsEmptyInnerCollections() { + List flat = + ListUtils.flatten( + Arrays.asList( + Collections.emptyList(), + Arrays.asList(1), + Collections.emptyList())); + + assertEquals(Arrays.asList(1), flat); + } + + @Test + public void emptyOuterCollectionYieldsEmptyList() { + assertTrue(ListUtils.flatten(Collections.>emptyList()).isEmpty()); + } + + @Test + public void acceptsMixedListImplementations() { + List> nested = new ArrayList<>(); + nested.add(new LinkedList<>(Arrays.asList(1, 2))); + nested.add(new ArrayList<>(Arrays.asList(3))); + + assertEquals(Arrays.asList(1, 2, 3), ListUtils.flatten(nested)); + } + + @Test + public void returnsIndependentCopy() { + List inner = new ArrayList<>(Arrays.asList(1, 2)); + List> outer = new ArrayList<>(); + outer.add(inner); + + List flat = ListUtils.flatten(outer); + assertNotSame(inner, flat); + + // Mutating an input after flattening must not change the result. + inner.add(3); + assertEquals(Arrays.asList(1, 2), flat); + + // Mutating the result must not change the inputs. + flat.add(99); + assertEquals(Arrays.asList(1, 2, 3), inner); + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/internal/concurrent/structured/AsyncTaskTest.java b/temporal-sdk/src/test/java/io/temporal/internal/concurrent/structured/AsyncTaskTest.java new file mode 100644 index 000000000..62809042f --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/concurrent/structured/AsyncTaskTest.java @@ -0,0 +1,81 @@ +package io.temporal.internal.concurrent.structured; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Test; + +public class AsyncTaskTest { + + @Test + public void recoverDoesNotHandleCancellation() { + DefaultTaskScope scope = new DefaultTaskScope<>(); + + CompletableFuture upstream = new CompletableFuture<>(); + AsyncTask task = scope.attach(upstream); + AsyncTask recovered = task.recover(error -> 99); + + task.cancel(); + + assertTrue(task.joinSettled().isCancelled()); + assertTrue(recovered.joinSettled().isCancelled()); + } + + @Test + public void cancelReturnsFalseWhenAlreadySettledAndPreservesValue() { + DefaultTaskScope scope = new DefaultTaskScope<>(); + + CompletableFuture upstream = new CompletableFuture<>(); + AsyncTask task = scope.attach(upstream); + upstream.complete(1); + + assertFalse(task.cancel()); + assertEquals(1, task.join().intValue()); + assertFalse(task.isCancelled()); + } + + @Test + public void cancelPropagatesToUnsettledDerivedChild() { + DefaultTaskScope scope = new DefaultTaskScope<>(); + + CompletableFuture upstream = new CompletableFuture<>(); + AsyncTask parent = scope.attach(upstream); + AsyncTask child = parent.map(value -> value + 1); + + parent.cancel(); + + assertTrue(parent.joinSettled().isCancelled()); + assertTrue(child.joinSettled().isCancelled()); + } + + @Test + public void mapRecoverChainTransformsFailureIntoSuccess() { + DefaultTaskScope scope = new DefaultTaskScope<>(); + + CompletableFuture upstream = new CompletableFuture<>(); + AsyncTask task = + scope.attach(upstream).map(value -> value / 0).recover(error -> 7).map(value -> value * 2); + + upstream.complete(1); + + assertEquals(14, task.join().intValue()); + } + + @Test + public void thenAcceptRunsSideEffectAndCompletesVoid() { + DefaultTaskScope scope = new DefaultTaskScope<>(); + + AtomicReference seen = new AtomicReference<>(); + CompletableFuture upstream = new CompletableFuture<>(); + AsyncTask done = scope.attach(upstream).thenAccept(seen::set); + + upstream.complete(7); + + assertNull(done.join()); + assertEquals(Integer.valueOf(7), seen.get()); + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/internal/concurrent/structured/CancelSourceTest.java b/temporal-sdk/src/test/java/io/temporal/internal/concurrent/structured/CancelSourceTest.java new file mode 100644 index 000000000..c923db8e1 --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/concurrent/structured/CancelSourceTest.java @@ -0,0 +1,112 @@ +package io.temporal.internal.concurrent.structured; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import io.temporal.common.CancellationToken; +import java.util.concurrent.CancellationException; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; + +public class CancelSourceTest { + + @Test + public void cancelRunsAllCallbacksEvenIfOneThrows() { + CancelSource source = new CancelSource<>(CancellationException::new); + AtomicInteger callbacksRan = new AtomicInteger(); + + source + .token() + .onCancel( + () -> { + callbacksRan.incrementAndGet(); + throw new IllegalStateException("boom"); + }); + source.token().onCancel(callbacksRan::incrementAndGet); + + source.cancel(); + source.cancel(); + + assertEquals(2, callbacksRan.get()); + } + + @Test + public void closingRegistrationPreventsCallback() { + CancelSource source = new CancelSource<>(CancellationException::new); + AtomicInteger callbacksRan = new AtomicInteger(); + + CancellationToken.Registration registration = + source.token().onCancel(callbacksRan::incrementAndGet); + registration.close(); + + source.cancel(); + + assertEquals(0, callbacksRan.get()); + } + + @Test + public void onCancelRunsImmediatelyWhenAlreadyCancelled() { + CancelSource source = new CancelSource<>(CancellationException::new); + AtomicInteger callbacksRan = new AtomicInteger(); + + source.cancel(); + source.token().onCancel(callbacksRan::incrementAndGet); + + assertEquals(1, callbacksRan.get()); + } + + @Test + public void linkedCancellationFlowsDownstreamOnly() { + CancelSource parent = new CancelSource<>(CancellationException::new); + CancelSource child = + CancelSource.linkedTo(CancellationException::new, parent.token()); + + child.cancel(); + + assertTrue(child.isCancelled()); + assertFalse(parent.isCancelled()); + + parent.cancel(); + + assertTrue(child.isCancelled()); + assertTrue(parent.isCancelled()); + } + + @Test + public void linkedToCancelsWhenAnyParentCancels() { + CancelSource a = new CancelSource<>(CancellationException::new); + CancelSource b = new CancelSource<>(CancellationException::new); + CancelSource child = + CancelSource.linkedTo(CancellationException::new, a.token(), b.token()); + + assertFalse(child.isCancelled()); + + b.cancel(); + + assertTrue(child.isCancelled()); + assertFalse(a.isCancelled()); + } + + @Test + public void linkedToAlreadyCancelledParentCancelsImmediately() { + CancelSource parent = new CancelSource<>(CancellationException::new); + parent.cancel(); + + CancelSource child = + CancelSource.linkedTo(CancellationException::new, parent.token()); + + assertTrue(child.isCancelled()); + } + + @Test + public void registrationCloseAfterCancelIsSafe() { + CancelSource source = new CancelSource<>(CancellationException::new); + CancellationToken.Registration registration = + source.token().onCancel(() -> {}); // registered before cancel + + source.cancel(); + + registration.close(); // must not throw even though cancel() dropped the callback list + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/internal/concurrent/structured/TaskScopeAndResultTest.java b/temporal-sdk/src/test/java/io/temporal/internal/concurrent/structured/TaskScopeAndResultTest.java new file mode 100644 index 000000000..06216b2f8 --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/concurrent/structured/TaskScopeAndResultTest.java @@ -0,0 +1,183 @@ +package io.temporal.internal.concurrent.structured; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import org.junit.Test; + +public class TaskScopeAndResultTest { + + @Test + public void awaitAllAndAwaitAllSettledHandleEmptyScopes() { + try (TaskScope scope = new DefaultTaskScope<>()) { + assertTrue(scope.awaitAll().join().isEmpty()); + assertTrue(scope.awaitAllSettled().join().isEmpty()); + } + } + + @Test + public void awaitAllReturnsValuesInAttachOrder() { + try (DefaultTaskScope scope = new DefaultTaskScope<>()) { + CompletableFuture first = new CompletableFuture<>(); + CompletableFuture second = new CompletableFuture<>(); + scope.attach(first); + scope.attach(second); + + CompletableFuture> out = scope.awaitAll(); + + // Completion order does not affect collection order. + second.complete(2); + first.complete(1); + + assertEquals(Arrays.asList(1, 2), out.join()); + } + } + + @Test + public void awaitAllSettledPreservesPerTaskOutcome() { + try (DefaultTaskScope scope = new DefaultTaskScope<>()) { + CompletableFuture ok = new CompletableFuture<>(); + CompletableFuture bad = new CompletableFuture<>(); + CompletableFuture cancelled = new CompletableFuture<>(); + scope.attach(ok); + scope.attach(bad); + scope.attach(cancelled); + + ok.complete(1); + bad.completeExceptionally(new IllegalArgumentException("bad")); + cancelled.cancel(true); + + List> results = scope.awaitAllSettled().join(); + + assertTrue(results.get(0).isSuccess()); + assertTrue(results.get(1).isFailure()); + assertTrue(results.get(2).isCancelled()); + } + } + + @Test + public void resultHelpersBehaveAcrossStates() { + Result success = Result.success(3); + Result failure = Result.failure(new IllegalArgumentException("bad")); + Result cancelled = Result.cancelled(); + + assertEquals(4, success.map(value -> value + 1).get().intValue()); + assertEquals(9, success.fold(value -> value * 3, error -> -1).intValue()); + assertEquals(7, failure.orElse(7).intValue()); + assertEquals(8, cancelled.orElse(8).intValue()); + assertFalse(failure.map(value -> value + 1).isSuccess()); + assertEquals("cancelled", cancelled.fold(value -> "success", error -> "cancelled")); + } + + @Test + public void awaitAllAppliesTransformerAcrossResults() { + try (DefaultTaskScope scope = new DefaultTaskScope<>()) { + CompletableFuture first = new CompletableFuture<>(); + CompletableFuture second = new CompletableFuture<>(); + scope.attach(first); + scope.attach(second); + + CompletableFuture sum = scope.awaitAll(values -> values.get(0) + values.get(1)); + + first.complete(2); + second.complete(3); + + assertEquals(5, sum.join().intValue()); + } + } + + @Test + public void awaitAllCollectsTransformedChildInsteadOfParent() { + try (DefaultTaskScope scope = new DefaultTaskScope<>()) { + CompletableFuture upstream = new CompletableFuture<>(); + scope.attach(upstream).map(value -> value * 10); + + upstream.complete(1); + + // The transformed child replaces its parent in the collected results. + assertEquals(Arrays.asList(10), scope.awaitAll().join()); + } + } + + @Test + public void awaitAllSettledDoesNotCancelSiblingsOnFailure() { + try (DefaultTaskScope scope = new DefaultTaskScope<>()) { + CompletableFuture good = new CompletableFuture<>(); + CompletableFuture bad = new CompletableFuture<>(); + AsyncTask goodTask = scope.attach(good); + scope.attach(bad); + + CompletableFuture>> out = scope.awaitAllSettled(); + bad.completeExceptionally(new IllegalStateException("boom")); + + // All-settled never fails fast, so the still-pending sibling is not cancelled. + assertFalse(goodTask.isCancelled()); + assertFalse(out.isDone()); + + good.complete(1); + + List> results = out.join(); + assertTrue(results.get(0).isSuccess()); + assertEquals(Integer.valueOf(1), results.get(0).get()); + assertTrue(results.get(1).isFailure()); + } + } + + @Test + public void attachMapCollectsTransformedValue() { + try (DefaultTaskScope scope = new DefaultTaskScope<>()) { + CompletableFuture upstream = new CompletableFuture<>(); + scope.attach(upstream).map(value -> value + 100); + + upstream.complete(5); + + // Mirrors the production pattern: attach an external future, then transform it in-scope. + assertEquals(Arrays.asList(105), scope.awaitAll().join()); + } + } + + @Test + public void awaitAllPropagatesTransformerException() { + try (DefaultTaskScope scope = new DefaultTaskScope<>()) { + CompletableFuture upstream = new CompletableFuture<>(); + scope.attach(upstream); + upstream.complete(1); + + CompletableFuture result = + scope.awaitAll( + values -> { + throw new IllegalStateException("transform"); + }); + + try { + result.join(); + fail("Expected the transformer failure to propagate"); + } catch (CompletionException e) { + assertTrue(e.getCause() instanceof IllegalStateException); + } + } + } + + @Test + public void largeFanOutCollectsAllResultsInOrder() { + int taskCount = 200; + try (DefaultTaskScope scope = new DefaultTaskScope<>()) { + for (int i = 0; i < taskCount; i++) { + scope.attach(CompletableFuture.completedFuture(i)); + } + + List results = scope.awaitAll().join(); + + assertEquals(taskCount, results.size()); + for (int i = 0; i < taskCount; i++) { + assertEquals(Integer.valueOf(i), results.get(i)); + } + } + } +} diff --git a/temporal-sdk/src/test/java/io/temporal/internal/concurrent/structured/TaskScopeTest.java b/temporal-sdk/src/test/java/io/temporal/internal/concurrent/structured/TaskScopeTest.java new file mode 100644 index 000000000..5025afefd --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/concurrent/structured/TaskScopeTest.java @@ -0,0 +1,165 @@ +package io.temporal.internal.concurrent.structured; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import org.junit.Test; + +public class TaskScopeTest { + + @Test + public void awaitAllFailsFastAndCancelsSiblings() { + try (DefaultTaskScope scope = new DefaultTaskScope<>()) { + CompletableFuture slow = new CompletableFuture<>(); + CompletableFuture failing = new CompletableFuture<>(); + scope.attach(slow); + scope.attach(failing); + + CompletableFuture> out = scope.awaitAll(); + failing.completeExceptionally(new IllegalStateException("boom")); + + try { + out.join(); + fail("Expected awaitAll to fail fast"); + } catch (CompletionException e) { + assertTrue(e.getCause() instanceof IllegalStateException); + } + assertTrue(slow.isCancelled()); + } + } + + @Test + public void awaitAllSettledWaitsForAllEvenWhenSomeFail() { + try (DefaultTaskScope scope = new DefaultTaskScope<>()) { + CompletableFuture failing = new CompletableFuture<>(); + CompletableFuture slow = new CompletableFuture<>(); + scope.attach(failing); + scope.attach(slow); + + CompletableFuture>> out = scope.awaitAllSettled(); + failing.completeExceptionally(new IllegalStateException("boom")); + + // Must keep waiting for the still-pending sibling. + assertFalse(out.isDone()); + + slow.complete(1); + + List> results = out.join(); + assertTrue(results.get(0).isFailure()); + assertTrue(results.get(1).isSuccess()); + } + } + + @Test + public void awaitAllWaitsForTaskDerivedWhileWaiting() { + try (DefaultTaskScope scope = new DefaultTaskScope<>()) { + CompletableFuture root = new CompletableFuture<>(); + TaskChain rootTask = scope.attach(root); + + CompletableFuture> out = scope.awaitAll(); + assertFalse(out.isDone()); + + // Derive a stage after awaitAll has already snapshotted the initial task set. + rootTask.map(value -> value + 1); + root.complete(1); + + assertEquals(Arrays.asList(2), out.join()); + } + } + + @Test + public void cancellingReturnedFutureCancelsScope() { + try (DefaultTaskScope scope = new DefaultTaskScope<>()) { + CompletableFuture pending = new CompletableFuture<>(); + scope.attach(pending); + + CompletableFuture awaiting = scope.awaitAll(); + awaiting.cancel(true); + + // Cancelling the future returned by awaitAll cancels the whole scope. + assertTrue(scope.token().isCancellationRequested()); + assertTrue(pending.isCancelled()); + } + } + + @Test + public void attachAfterCloseThrows() { + DefaultTaskScope scope = new DefaultTaskScope<>(); + scope.close(); + + try { + scope.attach(new CompletableFuture<>()); + fail("Expected attach on a closed scope to throw"); + } catch (IllegalStateException expected) { + // expected + } + } + + @Test + public void withScopeCompletesWithTransformedResult() throws Exception { + CompletableFuture first = new CompletableFuture<>(); + CompletableFuture second = new CompletableFuture<>(); + + CompletableFuture result = + TaskScope.withScope( + scope -> { + scope.attach(first); + scope.attach(second); + return scope.awaitAll(values -> values.get(0) + values.get(1)); + }); + + first.complete(2); + second.complete(3); + + assertEquals(5, result.get(2, TimeUnit.SECONDS).intValue()); + } + + @Test + public void withScopeFailsFastAndSettlesAfterAllTasksSettle() throws Exception { + CompletableFuture failing = new CompletableFuture<>(); + CompletableFuture sibling = new CompletableFuture<>(); + + CompletableFuture result = + TaskScope.withScope( + scope -> { + scope.attach(failing); + scope.attach(sibling); + return scope.awaitAll(ignored -> 0); + }); + + failing.completeExceptionally(new IllegalStateException("boom")); + + try { + result.get(2, TimeUnit.SECONDS); + fail("Expected the failure to propagate"); + } catch (ExecutionException e) { + assertTrue(e.getCause() instanceof IllegalStateException); + } + // Fail-fast cancels the sibling; the scope's result only settles once it has. + assertTrue(sibling.isCancelled()); + } + + @Test + public void withScopeBodyThrowingSynchronouslyCompletesFutureExceptionally() throws Exception { + CompletableFuture result = + TaskScope.withScope( + scope -> { + throw new IllegalStateException("boom"); + }); + + try { + result.get(2, TimeUnit.SECONDS); + fail("Expected the body failure to propagate"); + } catch (ExecutionException e) { + assertTrue(e.getCause() instanceof IllegalStateException); + } + } +}