Add structured concurrency wrapper for CompletableFuture

temporal-sdk/src/main/java/io/temporal/common/CancellationToken.java

@@ -0,0 +1,57 @@
+package io.temporal.common;
+
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CompletableFuture;
+
+/** Token that allows asynchronous code to observe cancellation requests. */
+@Experimental
+public interface CancellationToken {
+  CancellationToken NONE =
+      new CancellationToken() {
+        @Override
+        public boolean isCancellationRequested() {
+          return false;
+        }
+
+        @Override
+        public void throwIfCancellationRequested() throws CancellationException {}
+
+        @Override
+        public Registration onCancel(Runnable callback) {
+          return () -> {};
+        }
+      };
+
+  /** Returns true after cancellation has been requested. */
+  boolean isCancellationRequested();
+
+  /** Throws {@link CancellationException} if cancellation has been requested. */
+  void throwIfCancellationRequested() throws CancellationException;
+
+  /**
+   * Future that completes normally when cancellation has been requested.
+   *
+   * <p>Code waiting on external work can attach a callback to this future to abort in-flight
+   * requests when cancellation is requested.
+   */
+  default CompletableFuture<Void> getCancellationFuture() {
+    CompletableFuture<Void> result = new CompletableFuture<>();
+    Registration registration = onCancel(() -> result.complete(null));
+    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();
+  }
+}

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 <T> List<T> flatten(List<? extends List<? extends T>> lists) {
+    List<T> result = new ArrayList<>();
+    for (List<? extends T> list : lists) {
+      result.addAll(list);
+    }
+    return result;
+  }
+}

temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/AsyncTask.java

@@ -0,0 +1,72 @@
+package io.temporal.internal.concurrent.structured;
+
+import io.temporal.common.CancellationToken;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+/**
+ * Internal task handle over a {@link CompletableFuture} that manages cancellation and tracks
+ * derived tasks.
+ *
+ * <p><strong>Continuation style.</strong> {@link #map}, {@link #recover}, and {@link #whenSettled}
+ * chain follow-on work, mirroring {@code thenApply}/{@code exceptionally}/{@code whenComplete}.
+ *
+ * <p><strong>Cancellation is downstream by default.</strong> {@link #cancel()} settles this task
+ * and every task derived from it (its {@code map}/{@code recover} children). It does <em>not</em>
+ * cancel the task this one was derived <em>from</em>.
+ *
+ * @param <T> the result type
+ */
+interface AsyncTask<T> extends TaskChain<T> {
+
+  @Override
+  <R> AsyncTask<R> map(Function<? super T, ? extends R> fn);
+
+  @Override
+  AsyncTask<T> recover(Function<? super Throwable, ? extends T> fn);
+
+  @Override
+  default AsyncTask<Void> thenAccept(Consumer<? super T> fn) {
+    return map(
+        value -> {
+          fn.accept(value);
+          return null;
+        });
+  }
+
+  /**
+   * Runs a side effect when this task settles. Unlike {@code CompletableFuture.whenComplete}, the
+   * callback receives the unwrapped throwable ({@code null} on success).
+   */
+  AsyncTask<T> whenSettled(BiConsumer<? super T, ? super Throwable> cb);
+
+  /**
+   * 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<T> joinSettled();
+
+  /**
+   * @return the read-only cancellation token for this task.
+   */
+  CancellationToken token();
+
+  /** Escape hatch to the underlying future for interop with existing APIs. */
+  CompletableFuture<T> toCompletableFuture();
+}

temporal-sdk/src/main/java/io/temporal/internal/concurrent/structured/CancelSource.java

@@ -0,0 +1,89 @@
+package io.temporal.internal.concurrent.structured;
+
+import io.temporal.common.CancellationToken;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * The <em>write</em> side of cancellation. Whoever holds the {@code CancelSource} can request
+ * cancellation while everyone else observes using {@link #token()}.
+ */
+public final class CancelSource {
+
+  private final Object lock = new Object();
+  private volatile boolean cancelled = false;
+
+  /** Pending callbacks. Set to {@code null} by {@link #cancel()} once it takes ownership. */
+  private List<Runnable> callbacks = new ArrayList<>();
+
+  private final CancellationToken token =
+      new CancellationToken() {
+        @Override
+        public boolean isCancellationRequested() {
+          return cancelled;
+        }
+
+        @Override
+        public void throwIfCancellationRequested() {
+          if (cancelled) throw new java.util.concurrent.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. Idempotent; fires each registered callback exactly once. */
+  public void cancel() {
+    List<Runnable> toRun;
+    synchronized (lock) {
+      if (cancelled) {
+        return;
+      }
+      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 whenever any {@code parent} token is cancelled. */
+  public static CancelSource linkedTo(CancellationToken... parents) {
+    CancelSource s = new CancelSource();
+    for (CancellationToken p : parents) p.onCancel(s::cancel);
+    return s;
+  }
+}