Expand config with executor decorators (#1213)

related:
- #1186

Author: pivovarit

README.md

@@ -106,6 +106,8 @@ Additionally, you can customize:
 - grouping by key via `parallelBy(...)` / `parallelToStreamBy(...)` methods
 - ordered streaming via the `ordered()` configurer option (streaming collectors only)
 - a custom downstream `Collector` (`ParallelCollectors.parallel` only)
+- executor decoration via `executorDecorator()` to wrap the resolved executor
+- task decoration via `taskDecorator()` to wrap each individual task
 
 All configuration is done via the `CollectingConfigurer` (for `parallel`/`parallelBy`) or `StreamingConfigurer` (for `parallelToStream`/`parallelToStreamBy`) passed as a `Consumer`:
 
@@ -170,6 +172,40 @@ The `parallelBy(...)` and `parallelToStreamBy(...)` methods allow you to classif
 
 The `Group<K, V>` record provides `key()` and `values()` accessors, plus a `map()` method for transforming values while preserving the grouping key.
 
+#### Decorators
+
+Two decorator options let you add cross-cutting behavior without replacing the executor:
+
+**`executorDecorator(UnaryOperator<Executor>)`** wraps the resolved executor (the virtual-thread default or a custom one) and returns a replacement. It is invoked once per collector, before any tasks are submitted. This is a natural fit for intercepting every `execute()` call, for example to plug in a monitoring layer.
+
+The returned executor must not drop or discard tasks — doing so will cause the collector to wait indefinitely for results that will never arrive.
+
+    list.stream()
+      .collect(parallel(i -> foo(i), c -> c
+        .executorDecorator(exec -> task -> {
+            metrics.incrementAndGet();
+            exec.execute(task);
+        }),
+      toList()));
+
+**`taskDecorator(UnaryOperator<Runnable>)`** wraps each individual task before it is handed to the executor. Unlike the executor decorator, it runs on the worker thread and is re-applied for every element. This makes it the right tool for propagating thread-local context (MDC, OpenTelemetry spans, `SecurityContext`) into worker threads:
+
+    var snapshot = MDC.getCopyOfContextMap();
+
+    list.stream()
+      .collect(parallel(i -> foo(i), c -> c
+        .taskDecorator(task -> () -> {
+            MDC.setContextMap(snapshot);
+            try {
+                task.run();
+            } finally {
+                MDC.clear();
+            }
+        }),
+      toList()));
+
+Both decorators can be combined and each may be specified at most once per configurer.
+
 ### Leveraging CompletableFuture
 
 Parallel Collectors expose results wrapped in `CompletableFuture` instances, which provides great flexibility and the possibility of working with them in a non-blocking fashion:
@@ -243,6 +279,34 @@ What's more, since JDK9, [you can even provide your own timeout easily](https://
         .batching(),
       toList()));
 
+##### 8. Propagate MDC context into worker threads via `taskDecorator`
+
+    var snapshot = MDC.getCopyOfContextMap();
+
+    CompletableFuture<List<String>> result = list.stream()
+      .collect(parallel(i -> foo(i), c -> c
+        .taskDecorator(task -> () -> {
+            MDC.setContextMap(snapshot);
+            try {
+                task.run();
+            } finally {
+                MDC.clear();
+            }
+        }),
+      toList()));
+
+##### 9. Instrument every task submission via `executorDecorator`
+
+    var submitted = new AtomicInteger();
+
+    CompletableFuture<List<String>> result = list.stream()
+      .collect(parallel(i -> foo(i), c -> c
+        .executorDecorator(exec -> task -> {
+            submitted.incrementAndGet();
+            exec.execute(task);
+        }),
+      toList()));
+
 ## Rationale
 
 Stream API is a great tool for collection processing, especially if you need to parallelize the execution of CPU-intensive tasks, for example:

src/main/java/com/pivovarit/collectors/CollectingConfigurer.java

@@ -19,8 +19,10 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.Executor;
+import java.util.function.UnaryOperator;
 
 /**
  * Fluent configuration builder for collectors that <em>collect</em> all results (i.e. non-streaming).
@@ -88,6 +90,51 @@ public CollectingConfigurer executor(Executor executor) {
         return this;
     }
 
+    /**
+     * Decorates the executor used for running tasks.
+     * <p>
+     * The decorator receives the resolved executor (either the default virtual-thread executor or
+     * the one provided via {@link #executor(Executor)}) and returns a wrapped replacement.
+     * This is useful for augmenting the executor with cross-cutting concerns such as context
+     * propagation (MDC, OpenTelemetry spans, etc.) or monitoring, without replacing the executor entirely.
+     *
+     * <p><b>Note:</b> The executor returned by the decorator must not <em>drop</em> tasks on rejection.
+     * Dropping tasks will cause the collector to wait for results that will never be produced,
+     * which can lead to deadlocks.
+     *
+     * @param decorator a function that wraps the resolved executor
+     *
+     * @return this configurer instance for fluent chaining
+     */
+    public CollectingConfigurer executorDecorator(UnaryOperator<Executor> decorator) {
+        Objects.requireNonNull(decorator, "executor decorator can't be null");
+
+        addOnce(new ConfigProcessor.Option.ExecutorDecorator(decorator));
+        return this;
+    }
+
+    /**
+     * Decorates each individual task before it is submitted to the executor.
+     * <p>
+     * The decorator receives the {@link Runnable} representing a single unit of work and returns a
+     * wrapped replacement that runs in its place. This is useful for propagating thread-local context
+     * (e.g. MDC entries, OpenTelemetry spans, {@code SecurityContext}) into worker threads, or for
+     * per-task instrumentation, without replacing the executor entirely.
+     *
+     * <p>Unlike {@link #executorDecorator(UnaryOperator)}, which wraps the executor as a whole,
+     * this decorator is applied to each task individually and runs on the worker thread.
+     *
+     * @param decorator a function that wraps each submitted task
+     *
+     * @return this configurer instance for fluent chaining
+     */
+    public CollectingConfigurer taskDecorator(UnaryOperator<Runnable> decorator) {
+        Objects.requireNonNull(decorator, "task decorator can't be null");
+
+        addOnce(new ConfigProcessor.Option.TaskDecorator(decorator));
+        return this;
+    }
+
     List<ConfigProcessor.Option> getConfig() {
         return Collections.unmodifiableList(modifiers);
     }

src/main/java/com/pivovarit/collectors/ConfigProcessor.java

@@ -20,6 +20,7 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.function.UnaryOperator;
 
 import static java.util.Objects.requireNonNull;
 
@@ -29,7 +30,7 @@ final class ConfigProcessor {
       .name("parallel-collectors-", 0)
       .factory());
 
-    record Config(boolean ordered, boolean batching, int parallelism, Executor executor) {
+    record Config(boolean ordered, boolean batching, int parallelism, Executor executor, UnaryOperator<Runnable> taskDecorator) {
         Config {
             Objects.requireNonNull(executor, "executor can't be null");
         }
@@ -42,21 +43,37 @@ static Config process(List<Option> options) {
         Boolean ordered = null;
         Integer parallelism = null;
         Executor executor = null;
+        UnaryOperator<Executor> decorator = null;
+        UnaryOperator<Runnable> taskDecorator = null;
 
         for (var option : options) {
             switch (option) {
                 case Option.Batched ignored -> batching = true;
                 case Option.Parallelism(var p) -> parallelism = p;
                 case Option.ThreadPool(var e) -> executor = e;
                 case Option.Ordered ignored -> ordered = true;
+                case Option.ExecutorDecorator(var d) -> decorator = d;
+                case Option.TaskDecorator(var d) -> taskDecorator = d;
             }
         }
 
+        var resolvedExecutor = Objects.requireNonNullElse(executor, DEFAULT_EXECUTOR);
+        if (decorator != null) {
+            resolvedExecutor = decorator.apply(resolvedExecutor);
+            Preconditions.requireValidExecutor(resolvedExecutor);
+        }
+        if (taskDecorator != null) {
+            var td = taskDecorator;
+            var delegate = resolvedExecutor;
+            resolvedExecutor = r -> delegate.execute(td.apply(r));
+        }
+
         return new Config(
           Objects.requireNonNullElse(ordered, false),
           Objects.requireNonNullElse(batching, false),
           Objects.requireNonNullElse(parallelism, 0),
-          Objects.requireNonNullElse(executor, DEFAULT_EXECUTOR));
+          resolvedExecutor,
+          taskDecorator);
     }
 
     static String toHumanReadableString(Option option) {
@@ -65,6 +82,8 @@ static String toHumanReadableString(Option option) {
             case Option.Parallelism ignored -> "parallelism";
             case Option.ThreadPool ignored -> "executor";
             case Option.Ordered ignored -> "ordered";
+            case Option.ExecutorDecorator ignored -> "executor decorator";
+            case Option.TaskDecorator ignored -> "task decorator";
         };
     }
 
@@ -89,5 +108,17 @@ record Parallelism(int parallelism) implements Option {
                 Preconditions.requireValidParallelism(parallelism);
             }
         }
+
+        record ExecutorDecorator(UnaryOperator<Executor> decorator) implements Option {
+            public ExecutorDecorator {
+                Objects.requireNonNull(decorator, "decorator can't be null");
+            }
+        }
+
+        record TaskDecorator(UnaryOperator<Runnable> decorator) implements Option {
+            public TaskDecorator {
+                Objects.requireNonNull(decorator, "decorator can't be null");
+            }
+        }
     }
 }

src/main/java/com/pivovarit/collectors/StreamingConfigurer.java

@@ -19,8 +19,10 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.Executor;
+import java.util.function.UnaryOperator;
 
 /**
  * Fluent configuration builder for collectors that <em>stream</em> results.
@@ -103,6 +105,51 @@ public StreamingConfigurer executor(Executor executor) {
         return this;
     }
 
+    /**
+     * Decorates the executor used for running tasks.
+     * <p>
+     * The decorator receives the resolved executor (either the default virtual-thread executor or
+     * the one provided via {@link #executor(Executor)}) and returns a wrapped replacement.
+     * This is useful for augmenting the executor with cross-cutting concerns such as context
+     * propagation (MDC, OpenTelemetry spans, etc.) or monitoring, without replacing the executor entirely.
+     *
+     * <p><b>Note:</b> The executor returned by the decorator must not <em>drop</em> tasks on rejection.
+     * Dropping tasks will cause the stream to wait for results that will never be produced,
+     * which can lead to deadlocks.
+     *
+     * @param decorator a function that wraps the resolved executor
+     *
+     * @return this configurer instance for fluent chaining
+     */
+    public StreamingConfigurer executorDecorator(UnaryOperator<Executor> decorator) {
+        Objects.requireNonNull(decorator, "executor decorator can't be null");
+
+        addOnce(new ConfigProcessor.Option.ExecutorDecorator(decorator));
+        return this;
+    }
+
+    /**
+     * Decorates each individual task before it is submitted to the executor.
+     * <p>
+     * The decorator receives the {@link Runnable} representing a single unit of work and returns a
+     * wrapped replacement that runs in its place. This is useful for propagating thread-local context
+     * (e.g. MDC entries, OpenTelemetry spans, {@code SecurityContext}) into worker threads, or for
+     * per-task instrumentation, without replacing the executor entirely.
+     *
+     * <p>Unlike {@link #executorDecorator(UnaryOperator)}, which wraps the executor as a whole,
+     * this decorator is applied to each task individually and runs on the worker thread.
+     *
+     * @param decorator a function that wraps each submitted task
+     *
+     * @return this configurer instance for fluent chaining
+     */
+    public StreamingConfigurer taskDecorator(UnaryOperator<Runnable> decorator) {
+        Objects.requireNonNull(decorator, "task decorator can't be null");
+
+        addOnce(new ConfigProcessor.Option.TaskDecorator(decorator));
+        return this;
+    }
+
     List<ConfigProcessor.Option> getConfig() {
         return Collections.unmodifiableList(modifiers);
     }

src/test/java/com/pivovarit/collectors/OptionTest.java

@@ -86,4 +86,34 @@ void shouldThrowOnDuplicateOrdered() {
           .hasMessageContaining("'ordered' can only be configured once");
     }
 
+    @Test
+    void shouldThrowOnNullExecutorDecorator() {
+        assertThatThrownBy(() -> new CollectingConfigurer().executorDecorator(null))
+          .isInstanceOf(NullPointerException.class);
+    }
+
+    @Test
+    void shouldThrowOnNullExecutorDecoratorStreaming() {
+        assertThatThrownBy(() -> new StreamingConfigurer().executorDecorator(null))
+          .isInstanceOf(NullPointerException.class);
+    }
+
+    @Test
+    void shouldThrowOnDuplicateExecutorDecorator() {
+        var configurer = new CollectingConfigurer();
+        configurer.executorDecorator(e -> e);
+        assertThatThrownBy(() -> configurer.executorDecorator(e -> e))
+          .isInstanceOf(IllegalArgumentException.class)
+          .hasMessageContaining("'executor decorator' can only be configured once");
+    }
+
+    @Test
+    void shouldThrowOnDuplicateExecutorDecoratorStreaming() {
+        var configurer = new StreamingConfigurer();
+        configurer.executorDecorator(e -> e);
+        assertThatThrownBy(() -> configurer.executorDecorator(e -> e))
+          .isInstanceOf(IllegalArgumentException.class)
+          .hasMessageContaining("'executor decorator' can only be configured once");
+    }
+
 }