Workflow streams user publish executor

contrib/temporal-workflowstreams/README.md

@@ -112,6 +112,18 @@ Items are buffered and flushed automatically every batch interval (default 2s),
 when the buffer reaches the max batch size, on `forceFlush`, on an explicit
 `flush()`, or on `close()`.
 
+Background flushes run on the client's publish executor: a single daemon thread
+owned by each client by default. Applications running many clients can supply a
+shared executor via `publishExecutor` (see the options table); it is never shut
+down by the client.
+
+If a flush retry exceeds `maxRetryDuration`, background flushing stops for that
+client — neither the periodic tick nor a `forceFlush`/max-batch-size trigger
+sends again. Later items stay buffered until an explicit `flush()` or `close()`
+drains them. The two report the dropped batch at different points: `flush()`
+rethrows the `FlushTimeoutException` before sending anything (call it again to
+drain), while `close()` drains first and rethrows afterwards.
+
 ## Subscribing
 
 There are two subscriber APIs over one shared poll engine: a non-blocking
@@ -196,6 +208,7 @@ unrecoverable poll failure is rethrown from `hasNext()`.
 | `maxRetryDuration` | 10m | Max time to retry a failed flush before `FlushTimeoutException`. Must be < the workflow's publisher TTL (15m) to preserve exactly-once delivery |
 | `payloadConverters` | standard set | Per-item serialization. Payload conversion only — the client's codec chain runs once on the envelope, never per item |
 | `pollExecutor` | 2 daemon threads, client-owned | Scheduler shared by the client's subscriptions. It runs the short update-admission and delivery steps and poll cooldowns — never held during the long poll itself. A user-supplied executor is never shut down by the client; supply a bigger pool for many subscriptions against slow workflows |
+| `publishExecutor` | 1 daemon thread, client-owned | Scheduler driving the client's background flushes (periodic ticks and full-buffer/`forceFlush` triggers). A flush occupies a thread while signaling the workflow. A user-supplied executor is never shut down by the client; share one across clients instead of paying a thread per client |
 | `SubscribeOptions.pollCooldown` | 100ms | Min interval between polls |
 
 ## Cross-language protocol

contrib/temporal-workflowstreams/src/main/java/io/temporal/workflowstreams/WorkflowStreamClient.java

@@ -101,7 +101,8 @@ private WorkflowStreamClient(
             dataConverter,
             options.getBatchInterval(),
             options.getMaxBatchSize(),
-            options.getMaxRetryDuration());
+            options.getMaxRetryDuration(),
+            options.getPublishExecutor());
   }
 
   /**
@@ -206,28 +207,45 @@ private ScheduledExecutorService pollExecutor() {
    *
    * <p>Also stops this client's live subscriptions (their done futures complete normally, without
    * {@link WorkflowStreamListener#onCompleted}) and, if the client owns the default poll executor,
-   * shuts it down. A user-supplied poll executor is never shut down.
+   * shuts it down. That teardown runs even when the final flush fails, so a thrown {@link
+   * FlushTimeoutException} never leaves subscriptions polling. A user-supplied poll or publish
+   * executor is never shut down — only this client's own tasks on it are stopped.
    */
   @Override
   public void close() {
-    publisher.close();
+    // The final flush can throw (a deferred FlushTimeoutException, or a failing signal), and the
+    // rest of the teardown must still run: otherwise live subscriptions keep polling forever and
+    // the owned poll executor is never shut down.
+    try {
+      publisher.close();
+    } finally {
+      stopSubscriptions();
+      shutdownOwnedPollExecutor();
+    }
+  }
+
+  private void stopSubscriptions() {
     for (SubscriptionDriver driver : liveSubscriptions.toArray(new SubscriptionDriver[0])) {
       driver.close();
     }
+  }
+
+  private void shutdownOwnedPollExecutor() {
     ScheduledExecutorService owned;
     synchronized (this) {
       owned = ownedPollExecutor;
     }
-    if (owned != null) {
-      owned.shutdown();
-      try {
-        if (!owned.awaitTermination(1, TimeUnit.SECONDS)) {
-          owned.shutdownNow();
-        }
-      } catch (InterruptedException e) {
-        Thread.currentThread().interrupt();
+    if (owned == null) {
+      return;
+    }
+    owned.shutdown();
+    try {
+      if (!owned.awaitTermination(1, TimeUnit.SECONDS)) {
         owned.shutdownNow();
       }
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      owned.shutdownNow();
     }
   }
 

contrib/temporal-workflowstreams/src/main/java/io/temporal/workflowstreams/WorkflowStreamClientOptions.java

@@ -24,18 +24,21 @@ public static WorkflowStreamClientOptions getDefaultInstance() {
   private final Duration maxRetryDuration;
   private final PayloadConverter[] payloadConverters;
   @Nullable private final ScheduledExecutorService pollExecutor;
+  @Nullable private final ScheduledExecutorService publishExecutor;
 
   private WorkflowStreamClientOptions(
       Duration batchInterval,
       int maxBatchSize,
       Duration maxRetryDuration,
       PayloadConverter[] payloadConverters,
-      @Nullable ScheduledExecutorService pollExecutor) {
+      @Nullable ScheduledExecutorService pollExecutor,
+      @Nullable ScheduledExecutorService publishExecutor) {
     this.batchInterval = batchInterval;
     this.maxBatchSize = maxBatchSize;
     this.maxRetryDuration = maxRetryDuration;
     this.payloadConverters = payloadConverters.clone();
     this.pollExecutor = pollExecutor;
+    this.publishExecutor = publishExecutor;
   }
 
   public Duration getBatchInterval() {
@@ -59,12 +62,18 @@ public ScheduledExecutorService getPollExecutor() {
     return pollExecutor;
   }
 
+  @Nullable
+  public ScheduledExecutorService getPublishExecutor() {
+    return publishExecutor;
+  }
+
   public static final class Builder {
     private Duration batchInterval = WorkflowStreamConstants.DEFAULT_BATCH_INTERVAL;
     private int maxBatchSize;
     private Duration maxRetryDuration = WorkflowStreamConstants.DEFAULT_MAX_RETRY_DURATION;
     private PayloadConverter[] payloadConverters = new PayloadConverter[0];
     @Nullable private ScheduledExecutorService pollExecutor;
+    @Nullable private ScheduledExecutorService publishExecutor;
 
     private Builder() {}
 
@@ -128,9 +137,29 @@ public Builder setPollExecutor(ScheduledExecutorService pollExecutor) {
       return this;
     }
 
+    /**
+     * Executor that drives the client's background publish path: the periodic flushes, and the
+     * flushes triggered by a full buffer or {@code forceFlush}. The caller owns its lifecycle; it
+     * is shared across all publishes of this client and must have at least one thread. A flush
+     * blocks while signaling the workflow, so it occupies an executor thread for the duration of
+     * each send — supply a pool sized for the number of clients that may flush concurrently.
+     *
+     * <p>Default: a single-thread daemon executor created lazily and owned by the client's
+     * publisher (shut down by {@link WorkflowStreamClient#close}).
+     */
+    public Builder setPublishExecutor(ScheduledExecutorService publishExecutor) {
+      this.publishExecutor = publishExecutor;
+      return this;
+    }
+
     public WorkflowStreamClientOptions build() {
       return new WorkflowStreamClientOptions(
-          batchInterval, maxBatchSize, maxRetryDuration, payloadConverters, pollExecutor);
+          batchInterval,
+          maxBatchSize,
+          maxRetryDuration,
+          payloadConverters,
+          pollExecutor,
+          publishExecutor);
     }
   }
 }

contrib/temporal-workflowstreams/src/main/java/io/temporal/workflowstreams/internal/StreamPublisher.java

@@ -10,14 +10,19 @@
 import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.Executors;
+import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
 
 /**
  * Owns the client-side publish path: it buffers published values, batches them, and sends each
  * batch to the workflow via the injected signal function. It assigns the per-publisher dedup key (a
  * stable publisher ID plus a monotonic sequence advanced only on a confirmed send) so the workflow
- * can drop duplicates, and it retries a failed batch until the max retry duration elapses.
+ * can drop duplicates, and it retries a failed batch until the max retry duration elapses. Once a
+ * background flush exceeds that duration the background loop stops for good and the resulting
+ * {@link FlushTimeoutException} is deferred to the next {@link #flush} or {@link #close}.
  *
  * <p>The signal function is injected (rather than holding a client) so the publish path can be
  * exercised in isolation. Internal to the workflow streams module.
@@ -36,6 +41,8 @@ public interface SignalFunction {
   private final long batchIntervalMs;
   private final int maxBatchSize;
   private final long maxRetryDurationMs;
+  // When null, the publisher creates a single-thread executor it owns and shuts down in close().
+  @Nullable private final ScheduledExecutorService userExecutor;
 
   private final Object stateLock = new Object();
   private List<PublishEntry> buffer = new ArrayList<>();
@@ -45,8 +52,16 @@ public interface SignalFunction {
   private long pendingStartNanos;
   private boolean started;
   private boolean closed;
+  // Set when a background flush timed out: the loop is stopped for good, so no background send may
+  // run before flush()/close() surfaces the deferred error. Guarded by stateLock.
+  private boolean loopStopped;
   private FlushTimeoutException deferredError;
+  // The executor driving the flush loop once started; the owned one when no user executor was
+  // supplied. Guarded by stateLock.
   private ScheduledExecutorService scheduler;
+  // The periodic flush tick, tracked so it can be cancelled without shutting down a user-supplied
+  // executor. Guarded by stateLock.
+  private ScheduledFuture<?> flushTask;
 
   /** Serializes doFlush so concurrent callers send sequentially. */
   private final Object flushLock = new Object();
@@ -57,12 +72,31 @@ public StreamPublisher(
       Duration batchInterval,
       int maxBatchSize,
       Duration maxRetryDuration) {
+    this(signal, dataConverter, batchInterval, maxBatchSize, maxRetryDuration, null);
+  }
+
+  /**
+   * @param executor drives the background flush loop (the periodic ticks and the flushes triggered
+   *     by a full buffer or {@code forceFlush}). When non-null the caller owns its lifecycle and it
+   *     is never shut down by this publisher, so many publishers can share one executor; when null
+   *     a single-thread executor is created lazily, owned by this publisher, and shut down by
+   *     {@link #close}. Flushes block while signaling the workflow, so each in-flight flush
+   *     occupies an executor thread for the duration of the send.
+   */
+  public StreamPublisher(
+      SignalFunction signal,
+      DataConverter dataConverter,
+      Duration batchInterval,
+      int maxBatchSize,
+      Duration maxRetryDuration,
+      @Nullable ScheduledExecutorService executor) {
     this.signal = signal;
     this.dataConverter = dataConverter;
     this.publisherId = UUID.randomUUID().toString().replace("-", "").substring(0, 16);
     this.batchIntervalMs = batchInterval.toMillis();
     this.maxBatchSize = maxBatchSize;
     this.maxRetryDurationMs = maxRetryDuration.toMillis();
+    this.userExecutor = executor;
   }
 
   /**
@@ -73,6 +107,13 @@ public StreamPublisher(
    * publish} call itself instead of poisoning the buffer and silently wedging every later item
    * behind it in the background flush loop.
    *
+   * <p>After a background flush exceeds the max retry duration the background loop is stopped
+   * permanently — neither the periodic tick nor a {@code forceFlush}/max-batch-size trigger sends
+   * again. Items published afterwards stay buffered until {@link #flush} or {@link #close} drains
+   * them, and that call surfaces the deferred {@link FlushTimeoutException} first (flush) or after
+   * the final drain (close). This keeps a caller-owned executor untouched without letting more data
+   * ship before the failure is reported.
+   *
    * @throws RuntimeException if no configured payload converter accepts {@code value}
    */
   public void publish(String topic, Object value, boolean forceFlush) {
@@ -81,14 +122,19 @@ public void publish(String topic, Object value, boolean forceFlush) {
     ScheduledExecutorService toTrigger = null;
     synchronized (stateLock) {
       buffer.add(entry);
-      trigger = forceFlush || (maxBatchSize > 0 && buffer.size() >= maxBatchSize);
+      trigger = (forceFlush || (maxBatchSize > 0 && buffer.size() >= maxBatchSize)) && !loopStopped;
       if (!closed) {
         ensureStartedLocked();
         toTrigger = scheduler;
       }
     }
     if (trigger && toTrigger != null) {
-      toTrigger.execute(this::backgroundFlush);
+      try {
+        toTrigger.execute(this::backgroundFlush);
+      } catch (RejectedExecutionException e) {
+        // The executor stopped between reading it and submitting (close(), or a user executor
+        // shut down by its owner). The item stays buffered for flush()/close() to drain.
+      }
     }
   }
 
@@ -97,27 +143,54 @@ private void ensureStartedLocked() {
       return;
     }
     started = true;
-    scheduler =
-        Executors.newSingleThreadScheduledExecutor(
-            r -> {
-              Thread t = new Thread(r, "temporal-workflow-stream-publisher");
-              t.setDaemon(true);
-              return t;
-            });
-    scheduler.scheduleWithFixedDelay(
-        this::backgroundFlush, batchIntervalMs, batchIntervalMs, TimeUnit.MILLISECONDS);
+    if (userExecutor != null) {
+      scheduler = userExecutor;
+    } else {
+      scheduler =
+          Executors.newSingleThreadScheduledExecutor(
+              r -> {
+                Thread t = new Thread(r, "temporal-workflow-stream-publisher");
+                t.setDaemon(true);
+                return t;
+              });
+    }
+    try {
+      flushTask =
+          scheduler.scheduleWithFixedDelay(
+              this::backgroundFlush, batchIntervalMs, batchIntervalMs, TimeUnit.MILLISECONDS);
+    } catch (RejectedExecutionException e) {
+      // A user-supplied executor was already shut down. Don't fail the publish call with the
+      // executor's own exception: items stay buffered for flush()/close() to drain on the
+      // caller's thread, as they do after a flush timeout stops the loop.
+    }
   }
 
   private void backgroundFlush() {
+    synchronized (stateLock) {
+      if (loopStopped || closed) {
+        // The loop is stopped (a timed-out flush) or the publisher is closed. A task already
+        // queued at that point must not send: with a user-supplied executor nothing purges the
+        // queue, so this is the only thing keeping a flush from running after close() returned.
+        return;
+      }
+    }
     try {
       doFlush();
     } catch (FlushTimeoutException e) {
       // The pending batch was dropped and can't be recovered. Stash the error so
-      // flush/close surface it and stop the loop.
+      // flush/close surface it and stop the loop for good: with a user-supplied executor
+      // cancelling the periodic task is not enough, since publish() can still submit a
+      // triggered flush onto the still-live executor.
+      ScheduledFuture<?> toCancel;
       ScheduledExecutorService toStop;
       synchronized (stateLock) {
         deferredError = e;
-        toStop = scheduler;
+        loopStopped = true;
+        toCancel = flushTask;
+        toStop = ownedSchedulerLocked();
+      }
+      if (toCancel != null) {
+        toCancel.cancel(false);
       }
       if (toStop != null) {
         toStop.shutdown();
@@ -226,18 +299,24 @@ public void flush() {
 
   /**
    * Stops the background flush loop and drains any remaining items, surfacing a deferred {@link
-   * FlushTimeoutException} from a prior background failure.
+   * FlushTimeoutException} from a prior background failure. A user-supplied executor is never shut
+   * down; only the periodic flush task is cancelled, leaving the executor free for its other work.
    */
   public void close() {
+    ScheduledFuture<?> toCancel;
     ScheduledExecutorService toStop;
     synchronized (stateLock) {
       if (closed) {
         return;
       }
       closed = true;
-      toStop = scheduler;
+      toCancel = flushTask;
+      toStop = ownedSchedulerLocked();
     }
 
+    if (toCancel != null) {
+      toCancel.cancel(false);
+    }
     if (toStop != null) {
       toStop.shutdownNow();
       try {
@@ -259,6 +338,11 @@ public void close() {
     throwDeferred();
   }
 
+  /** Returns the executor to shut down on stop, or null when a user executor must be left alone. */
+  private ScheduledExecutorService ownedSchedulerLocked() {
+    return userExecutor == null ? scheduler : null;
+  }
+
   private void throwDeferred() {
     synchronized (stateLock) {
       if (deferredError != null) {