diff --git a/src/main/java/dev/shaaf/jgraphlet/CacheKey.java b/src/main/java/dev/shaaf/jgraphlet/CacheKey.java
index a8b6178..f1f8e10 100644
--- a/src/main/java/dev/shaaf/jgraphlet/CacheKey.java
+++ b/src/main/java/dev/shaaf/jgraphlet/CacheKey.java
@@ -1,5 +1,7 @@
package dev.shaaf.jgraphlet;
+import java.util.*;
+
/**
* Cache key used to uniquely identify the output of a task for a given input.
*
@@ -7,4 +9,19 @@
* @param input the input provided to that task
*/
public record CacheKey(String taskName, Object input) {
-}
\ No newline at end of file
+ public CacheKey {
+ Objects.requireNonNull(taskName, "taskName");
+ input = normalize(input); // snapshot for stable equals/hashCode
+ }
+ private static Object normalize(Object o) {
+ if (o == null) return null;
+ if (o instanceof Map m) {
+ // preserve iteration order to keep equals/hash stable
+ var copy = new LinkedHashMap<>(m);
+ return Map.copyOf(copy);
+ }
+ if (o instanceof List l) return List.copyOf(l);
+ if (o instanceof Set s) return Set.copyOf(s);
+ return o; // assume immutable or stable
+ }
+}
diff --git a/src/main/java/dev/shaaf/jgraphlet/TaskPipeline.java b/src/main/java/dev/shaaf/jgraphlet/TaskPipeline.java
index 102c16b..f69a821 100644
--- a/src/main/java/dev/shaaf/jgraphlet/TaskPipeline.java
+++ b/src/main/java/dev/shaaf/jgraphlet/TaskPipeline.java
@@ -1,10 +1,12 @@
package dev.shaaf.jgraphlet;
import java.util.*;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Collectors;
@@ -12,10 +14,10 @@
/**
* Async task pipeline that can be used to execute a chain of tasks in a defined order.
* It supports parallel execution of independent tasks, caching, and complex dependencies.
- *
- * This class implements {@link AutoCloseable} to support try-with-resources for
+ *
+ *
This class implements {@link AutoCloseable} to support try-with-resources for
* automatic cleanup of internal executor resources.
- *
+ *
* Example usage:
* {@code
* try (TaskPipeline pipeline = new TaskPipeline()) {
@@ -28,23 +30,31 @@
public class TaskPipeline implements AutoCloseable {
private static final Logger logger = Logger.getLogger(TaskPipeline.class.getName());
+ private static final long DEFAULT_SHUTDOWN_TIMEOUT_SECONDS = 30;
+ private static final long DEFAULT_SHUTDOWN_NOW_TIMEOUT_SECONDS = 10;
private final Map> tasks = new ConcurrentHashMap<>();
private final Map> graph = new ConcurrentHashMap<>();
+ private final Map> reverseGraph = new ConcurrentHashMap<>(); // For O(1) predecessor lookups
private final Map cache = new ConcurrentHashMap<>();
private final Map> futureCache = new ConcurrentHashMap<>();
private final ExecutorService executor;
private final boolean ownedExecutor;
+ private final long shutdownTimeoutSeconds;
+ private final long shutdownNowTimeoutSeconds;
- private volatile String lastAddedTaskName;
+ private final AtomicReference lastAddedTaskName = new AtomicReference<>();
+
+ private final ReadWriteLock pipelineLock = new ReentrantReadWriteLock();
+ private final Lock readLock = pipelineLock.readLock();
+ private final Lock writeLock = pipelineLock.writeLock();
/**
* Creates a TaskPipeline with an internally managed work-stealing executor.
* Close the pipeline or call shutdown() when finished to release resources.
*/
public TaskPipeline() {
- this.executor = Executors.newWorkStealingPool();
- this.ownedExecutor = true;
+ this(Executors.newWorkStealingPool(), true, DEFAULT_SHUTDOWN_TIMEOUT_SECONDS, DEFAULT_SHUTDOWN_NOW_TIMEOUT_SECONDS);
}
/**
@@ -54,8 +64,28 @@ public TaskPipeline() {
* @param executor the executor service to run tasks on
*/
public TaskPipeline(ExecutorService executor) {
+ this(executor, false, DEFAULT_SHUTDOWN_TIMEOUT_SECONDS, DEFAULT_SHUTDOWN_NOW_TIMEOUT_SECONDS);
+ }
+
+ /**
+ * Creates a TaskPipeline with custom shutdown timeout settings.
+ *
+ * @param shutdownTimeoutSeconds timeout in seconds for graceful shutdown
+ * @param shutdownNowTimeoutSeconds timeout in seconds for forced shutdown
+ */
+ public TaskPipeline(long shutdownTimeoutSeconds, long shutdownNowTimeoutSeconds) {
+ this(Executors.newWorkStealingPool(), true, shutdownTimeoutSeconds, shutdownNowTimeoutSeconds);
+ }
+
+ /**
+ * Private constructor for internal initialization.
+ */
+ private TaskPipeline(ExecutorService executor, boolean ownedExecutor,
+ long shutdownTimeoutSeconds, long shutdownNowTimeoutSeconds) {
this.executor = executor;
- this.ownedExecutor = false;
+ this.ownedExecutor = ownedExecutor;
+ this.shutdownTimeoutSeconds = shutdownTimeoutSeconds;
+ this.shutdownNowTimeoutSeconds = shutdownNowTimeoutSeconds;
}
/**
@@ -67,14 +97,19 @@ public TaskPipeline(ExecutorService executor) {
* @throws IllegalArgumentException if a task with the same name has already been added
*/
public TaskPipeline add(String taskName, Task task) {
- logger.log(Level.FINE, "Adding task {0} to the pipeline.", taskName);
- if (tasks.containsKey(taskName)) {
- throw new IllegalArgumentException("Task '" + taskName + "' has already been added.");
+ writeLock.lock();
+ try {
+ logger.log(Level.FINE, "Adding task {0} to the pipeline.", taskName);
+ if (tasks.putIfAbsent(taskName, task) != null) {
+ throw new IllegalArgumentException("Task '" + taskName + "' has already been added.");
+ }
+ graph.computeIfAbsent(taskName, k -> new CopyOnWriteArrayList<>());
+ reverseGraph.put(taskName, new CopyOnWriteArrayList<>());
+ lastAddedTaskName.set(taskName);
+ return this;
+ } finally {
+ writeLock.unlock();
}
- tasks.put(taskName, task);
- graph.put(taskName, new ArrayList<>());
- lastAddedTaskName = taskName;
- return this;
}
/**
@@ -87,12 +122,17 @@ public TaskPipeline add(String taskName, Task task) {
* @throws IllegalArgumentException if a task with the same name has already been added
*/
public TaskPipeline addTask(String taskName, Task task) {
- if (tasks.containsKey(taskName)) {
- throw new IllegalArgumentException("Task '" + taskName + "' has already been added.");
+ writeLock.lock();
+ try {
+ if (tasks.putIfAbsent(taskName, task) != null) {
+ throw new IllegalArgumentException("Task '" + taskName + "' has already been added.");
+ }
+ graph.computeIfAbsent(taskName, k -> new CopyOnWriteArrayList<>());
+ reverseGraph.computeIfAbsent(taskName, k -> new CopyOnWriteArrayList<>());
+ return this;
+ } finally {
+ writeLock.unlock();
}
- tasks.put(taskName, task);
- graph.put(taskName, new ArrayList<>());
- return this;
}
/**
@@ -104,14 +144,29 @@ public TaskPipeline addTask(String taskName, Task task) {
* @throws IllegalStateException if called before {@link #add(String, Task)}
*/
public TaskPipeline then(String nextTaskName, Task nextTask) {
- if (lastAddedTaskName == null) {
- throw new IllegalStateException("You must call 'add()' before calling 'then()'.");
+ writeLock.lock();
+ try {
+ String prev = lastAddedTaskName.get();
+ if (prev == null) {
+ throw new IllegalStateException("You must call 'add()' before calling 'then()'.");
+ }
+
+ // Inline addTask logic to avoid nested locking
+ if (tasks.putIfAbsent(nextTaskName, nextTask) != null) {
+ throw new IllegalArgumentException("Task '" + nextTaskName + "' has already been added.");
+ }
+ graph.computeIfAbsent(nextTaskName, k -> new CopyOnWriteArrayList<>());
+ reverseGraph.computeIfAbsent(nextTaskName, k -> new CopyOnWriteArrayList<>());
+
+ // Inline connect logic
+ graph.computeIfAbsent(prev, k -> new CopyOnWriteArrayList<>()).add(nextTaskName);
+ reverseGraph.computeIfAbsent(nextTaskName, k -> new CopyOnWriteArrayList<>()).add(prev);
+
+ this.lastAddedTaskName.set(nextTaskName);
+ return this;
+ } finally {
+ writeLock.unlock();
}
- addTask(nextTaskName, nextTask);
- String fromTaskName = this.lastAddedTaskName;
- connect(fromTaskName, nextTaskName);
- this.lastAddedTaskName = nextTaskName;
- return this;
}
/**
@@ -121,28 +176,33 @@ public TaskPipeline then(String nextTaskName, Task nextTask) {
* @param toTaskName the child task name that depends on the parent
* @return The pipeline instance for fluent chaining.
* @throws IllegalArgumentException if task names are null or identical
- * @throws IllegalStateException if either task has not been added yet
+ * @throws IllegalStateException if either task has not been added yet
*/
public TaskPipeline connect(String fromTaskName, String toTaskName) {
+ writeLock.lock();
+ try {
+ if (toTaskName == null || fromTaskName == null) {
+ throw new IllegalArgumentException("Task names cannot be null.");
+ }
- if (toTaskName == null || fromTaskName == null) {
- throw new IllegalArgumentException("Task names cannot be null.");
- }
+ if (fromTaskName.equals(toTaskName)) {
+ throw new IllegalArgumentException("Cannot connect a task to itself: " + fromTaskName);
+ }
- if (fromTaskName.equals(toTaskName)) {
- throw new IllegalArgumentException("Cannot connect a task to itself: " + fromTaskName);
- }
+ if (!tasks.containsKey(fromTaskName)) {
+ throw new IllegalStateException("The 'from' task '" + fromTaskName + "' must be added to the pipeline before connecting from it.");
+ }
- if (!tasks.containsKey(fromTaskName)) {
- throw new IllegalStateException("The 'from' task '" + fromTaskName + "' must be added to the pipeline before connecting from it.");
- }
+ if (!tasks.containsKey(toTaskName)) {
+ throw new IllegalStateException("The 'to' task '" + toTaskName + "' must be added to the pipeline before connecting to it.");
+ }
- if (!tasks.containsKey(toTaskName)) {
- throw new IllegalStateException("The 'to' task '" + toTaskName + "' must be added to the pipeline before connecting to it.");
+ graph.computeIfAbsent(fromTaskName, k -> new CopyOnWriteArrayList<>()).add(toTaskName);
+ reverseGraph.computeIfAbsent(toTaskName, k -> new CopyOnWriteArrayList<>()).add(fromTaskName); // Maintain reverse adjacency for O(1) lookups
+ return this;
+ } finally {
+ writeLock.unlock();
}
-
- graph.get(fromTaskName).add(toTaskName);
- return this;
}
/**
@@ -153,56 +213,61 @@ public TaskPipeline connect(String fromTaskName, String toTaskName) {
*/
@SuppressWarnings("unchecked")
public CompletableFuture run(Object initialInput) {
- PipelineContext context = new PipelineContext();
- List executionOrder = topologicalSort();
- Map> results = new HashMap<>();
+ readLock.lock();
+ try {
+ PipelineContext context = new PipelineContext();
+ List executionOrder = topologicalSort();
+ Map> results = new HashMap<>();
- for (String taskName : executionOrder) {
- Task currentTask = (Task) tasks.get(taskName);
- List predecessors = findPredecessorsFor(taskName);
+ for (String taskName : executionOrder) {
+ Task currentTask = (Task) tasks.get(taskName);
+ List predecessors = findPredecessorsFor(taskName);
- CompletableFuture[] parentFutures = predecessors.stream()
+ CompletableFuture[] parentFutures = predecessors.stream()
.map(results::get)
.toArray(CompletableFuture[]::new);
- CompletableFuture allParentsDone = CompletableFuture.allOf(parentFutures);
+ CompletableFuture allParentsDone = CompletableFuture.allOf(parentFutures);
- CompletableFuture currentFuture = allParentsDone.thenComposeAsync(v -> {
- Object input = gatherInputsFromCompletedParents(predecessors, results, initialInput);
- if (currentTask == null) {
- throw new TaskRunException("Task '" + taskName + "' was not found in the pipeline.");
- }
+ CompletableFuture currentFuture = allParentsDone.thenComposeAsync(v -> {
+ Object input = gatherInputsFromCompletedParents(predecessors, results, initialInput);
+ if (currentTask == null) {
+ throw new TaskRunException("Task '" + taskName + "' was not found in the pipeline.");
+ }
- if (currentTask.isCacheable()) {
- CacheKey cacheKey = new CacheKey(taskName, input);
-
- // Use computeIfAbsent to atomically check cache and compute if needed
- return futureCache.computeIfAbsent(cacheKey, k -> {
- logger.log(Level.FINE, "Executing task {0}, (cache miss).", taskName);
- CompletableFuture taskResultFuture = currentTask.execute(input, context);
-
- // Populate the object cache when the future completes
- return taskResultFuture.thenApply(result -> {
- cache.put(cacheKey, result);
- return result;
+ if (currentTask.isCacheable()) {
+ CacheKey cacheKey = new CacheKey(taskName, input);
+
+ // Use computeIfAbsent to atomically check cache and compute if needed
+ return futureCache.computeIfAbsent(cacheKey, k -> {
+ logger.log(Level.FINE, "Executing task {0}, (cache miss).", taskName);
+ CompletableFuture taskResultFuture = currentTask.execute(input, context);
+
+ // Populate the object cache when the future completes
+ return taskResultFuture.thenApply(result -> {
+ cache.put(cacheKey, result);
+ return result;
+ });
});
- });
- } else {
- logger.log(Level.FINE, "Executing task {0}, (cache miss).", taskName);
- return currentTask.execute(input, context);
- }
+ } else {
+ logger.log(Level.FINE, "Executing task {0}, (cache miss).", taskName);
+ return currentTask.execute(input, context);
+ }
- }, executor);
+ }, executor);
- results.put(taskName, currentFuture);
- }
+ results.put(taskName, currentFuture);
+ }
- if (executionOrder.isEmpty()) {
- return CompletableFuture.completedFuture(initialInput);
- }
+ if (executionOrder.isEmpty()) {
+ return CompletableFuture.completedFuture(initialInput);
+ }
- String lastTaskName = executionOrder.get(executionOrder.size() - 1);
- return results.get(lastTaskName);
+ String lastTaskName = executionOrder.get(executionOrder.size() - 1);
+ return results.get(lastTaskName);
+ } finally {
+ readLock.unlock();
+ }
}
/**
@@ -248,18 +313,14 @@ private List topologicalSort() {
/**
* Finds all direct parent tasks for a given task in the graph.
+ * Now runs in O(1) time using the reverse adjacency map.
*
* @param taskName the task for which to find direct predecessors
* @return a list of task names that directly precede the given task
*/
private List findPredecessorsFor(String taskName) {
- List predecessors = new ArrayList<>();
- for (Map.Entry> entry : graph.entrySet()) {
- if (entry.getValue().contains(taskName)) {
- predecessors.add(entry.getKey());
- }
- }
- return predecessors;
+ // O(1) lookup using reverse graph instead of O(n*m) iteration
+ return reverseGraph.getOrDefault(taskName, List.of());
}
/**
@@ -282,32 +343,66 @@ private Object gatherInputsFromCompletedParents(List predecessors, Map name,
- name -> results.get(name).join()
- ));
+ .collect(Collectors.toMap(
+ name -> name,
+ name -> results.get(name).join()
+ ));
}
/**
* Shuts down the internal executor if it was created by this TaskPipeline.
* Call this method when you're done with the pipeline to prevent resource leaks.
- *
- * Note: If you provided a custom executor via constructor, you are responsible
+ * This method will wait for currently executing tasks to complete, up to the configured timeout.
+ *
+ * Note: If you provided a custom executor via constructor, you are responsible
* for shutting it down yourself.
*/
public void shutdown() {
if (ownedExecutor && !executor.isShutdown()) {
+ logger.log(Level.FINE, "Initiating graceful shutdown of TaskPipeline executor");
executor.shutdown();
+ try {
+ if (!executor.awaitTermination(shutdownTimeoutSeconds, TimeUnit.SECONDS)) {
+ logger.log(Level.WARNING,
+ "Executor did not terminate within {0} seconds, forcing shutdown",
+ shutdownTimeoutSeconds);
+ shutdownNow();
+ } else {
+ logger.log(Level.FINE, "TaskPipeline executor shutdown completed successfully");
+ }
+ } catch (InterruptedException e) {
+ logger.log(Level.WARNING, "Shutdown interrupted, forcing immediate termination", e);
+ shutdownNow();
+ Thread.currentThread().interrupt();
+ }
}
}
/**
* Forcibly shuts down the internal executor if it was created by this TaskPipeline.
- * This may interrupt running tasks.
+ * This may interrupt running tasks. Waits for tasks to respond to interruption.
*/
public void shutdownNow() {
if (ownedExecutor && !executor.isShutdown()) {
- executor.shutdownNow();
+ logger.log(Level.FINE, "Forcing immediate shutdown of TaskPipeline executor");
+ List pendingTasks = executor.shutdownNow();
+ if (!pendingTasks.isEmpty()) {
+ logger.log(Level.INFO, "Cancelled {0} pending tasks during forced shutdown",
+ pendingTasks.size());
+ }
+ try {
+ if (!executor.awaitTermination(shutdownNowTimeoutSeconds, TimeUnit.SECONDS)) {
+ logger.log(Level.SEVERE,
+ "Executor did not terminate after shutdownNow within {0} seconds. " +
+ "Some threads may still be running.",
+ shutdownNowTimeoutSeconds);
+ } else {
+ logger.log(Level.FINE, "Forced shutdown completed");
+ }
+ } catch (InterruptedException e) {
+ logger.log(Level.SEVERE, "Forced shutdown interrupted", e);
+ Thread.currentThread().interrupt();
+ }
}
}
@@ -315,7 +410,7 @@ public void shutdownNow() {
* Closes the TaskPipeline by gracefully shutting down the internal executor
* if it was created by this instance. This method is called automatically
* when using try-with-resources.
- *
+ *
* This is equivalent to calling {@link #shutdown()}.
*/
@Override
diff --git a/src/test/java/dev/shaaf/jgraphlet/TaskPipelinePerformanceTest.java b/src/test/java/dev/shaaf/jgraphlet/TaskPipelinePerformanceTest.java
new file mode 100644
index 0000000..5c03043
--- /dev/null
+++ b/src/test/java/dev/shaaf/jgraphlet/TaskPipelinePerformanceTest.java
@@ -0,0 +1,180 @@
+package dev.shaaf.jgraphlet;
+
+import org.junit.jupiter.api.Test;
+
+import java.time.Duration;
+import java.time.Instant;
+import java.util.concurrent.CompletableFuture;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+/**
+ * Performance tests for TaskPipeline optimizations.
+ * Demonstrates the performance improvements from using reverse adjacency map.
+ */
+class TaskPipelinePerformanceTest {
+
+ @Test
+ void shouldHandleLargeFanInGraphEfficiently() {
+ // This test creates a large fan-in graph to test the O(1) predecessor lookup
+ // Before optimization: O(n*m) where n = number of tasks, m = average connections per task
+ // After optimization: O(1) lookup time
+
+ TaskPipeline pipeline = new TaskPipeline();
+
+ // Create a simple task for testing
+ Task simpleTask = (input, context) ->
+ CompletableFuture.completedFuture(input + "-processed");
+
+ Task, String> aggregatorTask = (inputs, context) ->
+ CompletableFuture.completedFuture("Aggregated " + inputs.size() + " inputs");
+
+ // Create a large fan-in scenario: many tasks feeding into one
+ int numFeederTasks = 100;
+
+ // Add the aggregator task
+ pipeline.addTask("aggregator", aggregatorTask);
+
+ // Add many feeder tasks and connect them all to the aggregator
+ Instant startSetup = Instant.now();
+ for (int i = 0; i < numFeederTasks; i++) {
+ String taskName = "feeder" + i;
+ pipeline.addTask(taskName, simpleTask);
+ pipeline.connect(taskName, "aggregator");
+ }
+ Duration setupTime = Duration.between(startSetup, Instant.now());
+
+ // Run the pipeline - this will call findPredecessorsFor many times
+ Instant startRun = Instant.now();
+ CompletableFuture result = pipeline.run("test-input");
+ Object finalResult = result.join();
+ Duration runTime = Duration.between(startRun, Instant.now());
+
+ // Assert
+ assertNotNull(finalResult);
+ assertEquals("Aggregated " + numFeederTasks + " inputs", finalResult);
+
+ // Performance assertions
+ assertTrue(setupTime.toMillis() < 100,
+ "Setup should be fast even with " + numFeederTasks + " connections (took: " + setupTime.toMillis() + "ms)");
+ assertTrue(runTime.toMillis() < 500,
+ "Execution should be fast with O(1) predecessor lookups (took: " + runTime.toMillis() + "ms)");
+
+ System.out.println("✓ Large fan-in graph (" + numFeederTasks + " tasks) setup: " +
+ setupTime.toMillis() + "ms, execution: " + runTime.toMillis() + "ms");
+
+ pipeline.shutdown();
+ }
+
+ @Test
+ void shouldHandleComplexDiamondGraphEfficiently() {
+ // Test a diamond-shaped graph with multiple paths
+ // This tests that the reverse graph correctly maintains multiple predecessors
+
+ TaskPipeline pipeline = new TaskPipeline();
+
+ Task transformTask = (input, context) ->
+ CompletableFuture.completedFuture(input + "-transformed");
+
+ Task, String> mergeTask = (inputs, context) -> {
+ String path1 = (String) inputs.get("path1");
+ String path2 = (String) inputs.get("path2");
+ return CompletableFuture.completedFuture(path1 + " & " + path2);
+ };
+
+ // Create diamond shape: start -> path1 & path2 -> end
+ pipeline.addTask("start", transformTask)
+ .addTask("path1", transformTask)
+ .addTask("path2", transformTask)
+ .addTask("end", mergeTask);
+
+ pipeline.connect("start", "path1")
+ .connect("start", "path2")
+ .connect("path1", "end")
+ .connect("path2", "end");
+
+ // Run and verify
+ String result = (String) pipeline.run("input").join();
+
+ assertEquals("input-transformed-transformed & input-transformed-transformed", result);
+ System.out.println("✓ Diamond graph executed correctly with optimized predecessor lookups");
+
+ pipeline.shutdown();
+ }
+
+ @Test
+ void shouldScaleWellWithManyTasks() {
+ // Test scalability with a large number of tasks in a chain
+ TaskPipeline pipeline = new TaskPipeline();
+
+ Task incrementTask = (input, context) ->
+ CompletableFuture.completedFuture(input + 1);
+
+ int chainLength = 50;
+
+ // Build a long chain of tasks
+ Instant startSetup = Instant.now();
+ pipeline.add("task0", incrementTask);
+ for (int i = 1; i < chainLength; i++) {
+ pipeline.then("task" + i, incrementTask);
+ }
+ Duration setupTime = Duration.between(startSetup, Instant.now());
+
+ // Execute the chain
+ Instant startRun = Instant.now();
+ Integer result = (Integer) pipeline.run(0).join();
+ Duration runTime = Duration.between(startRun, Instant.now());
+
+ // Verify result
+ assertEquals(chainLength, result, "Should increment " + chainLength + " times");
+
+ // Performance checks
+ assertTrue(setupTime.toMillis() < 50,
+ "Chain setup should be fast (took: " + setupTime.toMillis() + "ms)");
+ assertTrue(runTime.toMillis() < 200,
+ "Chain execution should be efficient (took: " + runTime.toMillis() + "ms)");
+
+ System.out.println("✓ Chain of " + chainLength + " tasks - setup: " +
+ setupTime.toMillis() + "ms, execution: " + runTime.toMillis() + "ms");
+
+ pipeline.shutdown();
+ }
+
+ @Test
+ void shouldMaintainCorrectReverseGraphAfterMultipleConnections() {
+ // Test that reverse graph is correctly maintained with complex connections
+ TaskPipeline pipeline = new TaskPipeline();
+
+ Task task = (input, context) ->
+ CompletableFuture.completedFuture(input);
+
+ Task, String> collectTask = (inputs, context) ->
+ CompletableFuture.completedFuture("Collected: " + inputs.keySet());
+
+ // Create a complex graph structure
+ pipeline.addTask("A", task)
+ .addTask("B", task)
+ .addTask("C", task)
+ .addTask("D", collectTask)
+ .addTask("E", collectTask);
+
+ // Multiple connections
+ pipeline.connect("A", "D") // A -> D
+ .connect("B", "D") // B -> D
+ .connect("A", "E") // A -> E
+ .connect("C", "E"); // C -> E
+
+ // D should have predecessors: A, B
+ // E should have predecessors: A, C
+
+ // Run pipeline - this will use findPredecessorsFor internally
+ CompletableFuture futureD = pipeline.run("test");
+
+ // The execution should work correctly with the optimized reverse graph
+ assertDoesNotThrow(() -> futureD.join());
+
+ System.out.println("✓ Complex graph with multiple connections handled correctly");
+
+ pipeline.shutdown();
+ }
+}
\ No newline at end of file
diff --git a/src/test/java/dev/shaaf/jgraphlet/TaskPipelineShutdownTest.java b/src/test/java/dev/shaaf/jgraphlet/TaskPipelineShutdownTest.java
new file mode 100644
index 0000000..48567b7
--- /dev/null
+++ b/src/test/java/dev/shaaf/jgraphlet/TaskPipelineShutdownTest.java
@@ -0,0 +1,202 @@
+package dev.shaaf.jgraphlet;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class TaskPipelineShutdownTest {
+
+ @Test
+ void shouldShutdownGracefullyWithTimeout() throws InterruptedException {
+ // Arrange: Create pipeline with reasonable timeout
+ TaskPipeline pipeline = new TaskPipeline(5, 2);
+
+ CountDownLatch taskStarted = new CountDownLatch(1);
+ CountDownLatch allowCompletion = new CountDownLatch(1);
+ AtomicBoolean taskCompleted = new AtomicBoolean(false);
+
+ // Create a task that waits for signal before completing
+ Task controllableTask = (input, context) -> {
+ taskStarted.countDown();
+ try {
+ // Wait for signal to complete
+ if (allowCompletion.await(10, TimeUnit.SECONDS)) {
+ taskCompleted.set(true);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ return CompletableFuture.completedFuture("completed");
+ };
+
+ // Act
+ pipeline.add("controllableTask", controllableTask);
+ CompletableFuture future = pipeline.run("input");
+
+ // Wait for task to start
+ assertTrue(taskStarted.await(2, TimeUnit.SECONDS), "Task should have started");
+
+ // Signal task to complete
+ allowCompletion.countDown();
+
+ // Wait for completion
+ future.join();
+
+ // Shutdown should complete quickly since task is done
+ long startTime = System.currentTimeMillis();
+ pipeline.shutdown();
+ long shutdownTime = System.currentTimeMillis() - startTime;
+
+ // Assert
+ assertTrue(taskCompleted.get(), "Task should have completed");
+ assertTrue(shutdownTime < 1000, "Shutdown should be quick when no tasks running");
+ }
+
+ @Test
+ void shouldNotShutdownExternalExecutor() throws InterruptedException {
+ // Arrange: Create external executor
+ var externalExecutor = java.util.concurrent.Executors.newFixedThreadPool(2);
+
+ try {
+ TaskPipeline pipeline = new TaskPipeline(externalExecutor);
+
+ Task task = (input, context) ->
+ CompletableFuture.completedFuture("result");
+
+ // Act
+ pipeline.add("task", task);
+ pipeline.run("input").join();
+ pipeline.shutdown();
+
+ // Assert - external executor should still be running
+ assertFalse(externalExecutor.isShutdown(),
+ "External executor should not be shut down by pipeline");
+
+ } finally {
+ // Clean up
+ externalExecutor.shutdown();
+ externalExecutor.awaitTermination(1, TimeUnit.SECONDS);
+ }
+ }
+
+ @Test
+ void shouldHandleMultipleShutdownCallsGracefully() {
+ // Arrange
+ TaskPipeline pipeline = new TaskPipeline();
+
+ Task simpleTask = (input, context) ->
+ CompletableFuture.completedFuture("done");
+
+ // Act
+ pipeline.add("simpleTask", simpleTask);
+ Object result = pipeline.run("input").join();
+
+ // Multiple shutdown calls should not cause issues
+ pipeline.shutdown();
+ pipeline.shutdown();
+ pipeline.close();
+
+ // Assert
+ assertEquals("done", result);
+ // No exceptions should be thrown
+ }
+
+ @Test
+ void shouldUseCustomTimeoutValues() throws InterruptedException {
+ // This test verifies that custom timeout values are properly used
+ // We create a pipeline with short timeouts and verify the shutdown behavior
+
+ // Arrange: Create pipeline with custom timeouts
+ long gracefulTimeout = 2; // 2 seconds for graceful shutdown
+ long forcedTimeout = 1; // 1 second for forced shutdown
+ TaskPipeline pipeline = new TaskPipeline(gracefulTimeout, forcedTimeout);
+
+ // Create a simple task that completes quickly
+ Task quickTask = (input, context) ->
+ CompletableFuture.completedFuture("done quickly");
+
+ // Act
+ pipeline.add("quickTask", quickTask);
+ Object result = pipeline.run("input").join();
+
+ // Measure shutdown time - should be very quick since no tasks are running
+ long startTime = System.currentTimeMillis();
+ pipeline.shutdown();
+ long shutdownTime = System.currentTimeMillis() - startTime;
+
+ // Assert
+ assertEquals("done quickly", result);
+ // Shutdown should be nearly instant when no tasks are running
+ assertTrue(shutdownTime < 500, "Shutdown should be quick when no tasks are running");
+
+ // Verify we can create another pipeline with different timeouts
+ TaskPipeline pipeline2 = new TaskPipeline(10, 5);
+ pipeline2.add("task", quickTask);
+ pipeline2.run("test").join();
+ pipeline2.shutdown();
+ // No exception should be thrown
+ }
+
+ @Test
+ void shouldCleanupResourcesWithTryWithResources() throws Exception {
+ String result;
+
+ // Use try-with-resources to ensure automatic cleanup
+ try (TaskPipeline pipeline = new TaskPipeline()) {
+ Task task = (input, context) ->
+ CompletableFuture.completedFuture(input.toUpperCase());
+
+ pipeline.add("upperCase", task);
+ result = (String) pipeline.run("hello").join();
+ } // Pipeline.close() called automatically here
+
+ // Assert
+ assertEquals("HELLO", result);
+ // Pipeline resources should be cleaned up automatically
+ }
+
+ @Test
+ void shouldHandleShutdownNowCorrectly() {
+ // Arrange
+ TaskPipeline pipeline = new TaskPipeline();
+
+ CountDownLatch taskStarted = new CountDownLatch(1);
+ AtomicBoolean taskInterrupted = new AtomicBoolean(false);
+
+ Task longTask = (input, context) ->
+ CompletableFuture.supplyAsync(() -> {
+ taskStarted.countDown();
+ try {
+ Thread.sleep(5000);
+ } catch (InterruptedException e) {
+ taskInterrupted.set(true);
+ Thread.currentThread().interrupt();
+ }
+ return "should be interrupted";
+ });
+
+ // Act
+ pipeline.add("longTask", longTask);
+ pipeline.run("input"); // Start but don't wait
+
+ try {
+ taskStarted.await(1, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ fail("Test interrupted");
+ }
+
+ // Force immediate shutdown
+ pipeline.shutdownNow();
+
+ // Assert - task should be interrupted
+ // Note: The interruption happens in the ForkJoinPool threads
+ // We can't guarantee the task sees the interruption immediately
+ // but shutdownNow() should have been called
+ assertNotNull(pipeline); // Pipeline should still be valid after shutdownNow
+ }
+}
\ No newline at end of file
diff --git a/src/test/java/dev/shaaf/jgraphlet/TaskPipelineTest.java b/src/test/java/dev/shaaf/jgraphlet/TaskPipelineTest.java
index 1c448db..5386c7f 100644
--- a/src/test/java/dev/shaaf/jgraphlet/TaskPipelineTest.java
+++ b/src/test/java/dev/shaaf/jgraphlet/TaskPipelineTest.java
@@ -86,7 +86,11 @@ void shouldSetLastAddedTaskNameCorrectlyAfterAdd() {
try {
var field = TaskPipeline.class.getDeclaredField("lastAddedTaskName");
field.setAccessible(true);
- String lastAddedTaskName = (String) field.get(pipeline);
+ // Now lastAddedTaskName is an AtomicReference
+ @SuppressWarnings("unchecked")
+ java.util.concurrent.atomic.AtomicReference atomicRef =
+ (java.util.concurrent.atomic.AtomicReference) field.get(pipeline);
+ String lastAddedTaskName = atomicRef.get();
// Assert
assertEquals(taskName, lastAddedTaskName, "lastAddedTaskName should match the name of the last added task.");
diff --git a/src/test/java/dev/shaaf/jgraphlet/TaskPipelineVirtualThreadsTest.java b/src/test/java/dev/shaaf/jgraphlet/TaskPipelineVirtualThreadsTest.java
new file mode 100644
index 0000000..ff3fe5a
--- /dev/null
+++ b/src/test/java/dev/shaaf/jgraphlet/TaskPipelineVirtualThreadsTest.java
@@ -0,0 +1,373 @@
+package dev.shaaf.jgraphlet;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledForJreRange;
+import org.junit.jupiter.api.condition.JRE;
+
+import java.lang.reflect.Method;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.IntStream;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+/**
+ * Tests for TaskPipeline compatibility with virtual threads.
+ * These tests only run on Java 21+ where virtual threads are available.
+ * The project still targets Java 17 for compatibility.
+ */
+@EnabledForJreRange(min = JRE.JAVA_21)
+class TaskPipelineVirtualThreadsTest {
+
+ /**
+ * Creates a virtual thread executor using reflection to maintain Java 17 compatibility.
+ * This method will only be called when running on Java 21+.
+ */
+ private ExecutorService createVirtualThreadExecutor() {
+ try {
+ // Use reflection to call Executors.newVirtualThreadPerTaskExecutor()
+ // This allows the code to compile on Java 17 but use virtual threads on Java 21+
+ Class executorsClass = Executors.class;
+ Method method = executorsClass.getMethod("newVirtualThreadPerTaskExecutor");
+ return (ExecutorService) method.invoke(null);
+ } catch (Exception e) {
+ // Should not happen as this test only runs on Java 21+
+ throw new RuntimeException("Failed to create virtual thread executor", e);
+ }
+ }
+
+ @Test
+ void shouldWorkWithVirtualThreadExecutor() throws Exception {
+ // Arrange
+ ExecutorService virtualExecutor = createVirtualThreadExecutor();
+ TaskPipeline pipeline = new TaskPipeline(virtualExecutor);
+
+ try {
+ // Create tasks that would benefit from virtual threads (I/O simulation)
+ Task fetchTask = (input, context) ->
+ CompletableFuture.supplyAsync(() -> {
+ simulateIOOperation(100); // Simulate I/O
+ return "fetched: " + input;
+ });
+
+ Task processTask = (input, context) ->
+ CompletableFuture.supplyAsync(() -> {
+ simulateIOOperation(50); // Simulate processing
+ return input.toUpperCase();
+ });
+
+ // Act
+ pipeline.add("fetch", fetchTask)
+ .then("process", processTask);
+
+ String result = (String) pipeline.run("data").get(5, TimeUnit.SECONDS);
+
+ // Assert
+ assertEquals("FETCHED: DATA", result);
+ System.out.println("✓ Virtual thread executor works with TaskPipeline");
+
+ } finally {
+ pipeline.shutdown();
+ virtualExecutor.shutdown();
+ assertTrue(virtualExecutor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ void shouldHandleHighConcurrencyWithVirtualThreads() throws Exception {
+ // Virtual threads excel at high concurrency scenarios
+ ExecutorService virtualExecutor = createVirtualThreadExecutor();
+ TaskPipeline pipeline = new TaskPipeline(virtualExecutor);
+
+ try {
+ AtomicInteger completedTasks = new AtomicInteger(0);
+
+ // Create a simple I/O-bound task
+ Task ioTask = (input, context) ->
+ CompletableFuture.supplyAsync(() -> {
+ simulateIOOperation(100); // Simulate blocking I/O
+ completedTasks.incrementAndGet();
+ return "Task " + input + " completed";
+ });
+
+ pipeline.add("ioTask", ioTask);
+
+ // Act: Run many tasks concurrently
+ int numTasks = 100;
+ List> futures = new ArrayList<>();
+
+ Instant start = Instant.now();
+ for (int i = 0; i < numTasks; i++) {
+ futures.add(pipeline.run(i));
+ }
+
+ // Wait for all to complete
+ CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
+ .get(30, TimeUnit.SECONDS);
+
+ Duration duration = Duration.between(start, Instant.now());
+
+ // Assert
+ assertEquals(numTasks, completedTasks.get(),
+ "All tasks should complete");
+
+ // With virtual threads, 100 tasks with 100ms I/O should complete much faster than sequential
+ // Sequential would take 10 seconds, we should complete in much less
+ assertTrue(duration.getSeconds() < 5,
+ "High concurrency tasks should complete quickly with virtual threads (took: " + duration.getSeconds() + "s)");
+
+ // Virtual threads allow many concurrent operations
+ System.out.println("✓ Completed " + numTasks + " tasks in " + duration.toMillis() + "ms with virtual threads");
+
+ } finally {
+ pipeline.shutdown();
+ virtualExecutor.shutdown();
+ assertTrue(virtualExecutor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ void shouldHandleFanInPatternWithVirtualThreads() throws Exception {
+ ExecutorService virtualExecutor = createVirtualThreadExecutor();
+ TaskPipeline pipeline = new TaskPipeline(virtualExecutor);
+
+ try {
+ // Create multiple I/O-bound tasks for fan-in
+ Task dbTask = (input, context) ->
+ CompletableFuture.supplyAsync(() -> {
+ simulateIOOperation(150); // Simulate DB query
+ return "db-data";
+ });
+
+ Task apiTask = (input, context) ->
+ CompletableFuture.supplyAsync(() -> {
+ simulateIOOperation(200); // Simulate API call
+ return "api-data";
+ });
+
+ Task cacheTask = (input, context) ->
+ CompletableFuture.supplyAsync(() -> {
+ simulateIOOperation(50); // Simulate cache lookup
+ return "cache-data";
+ });
+
+ Task, String> aggregateTask = (inputs, context) ->
+ CompletableFuture.supplyAsync(() -> {
+ String db = (String) inputs.get("db");
+ String api = (String) inputs.get("api");
+ String cache = (String) inputs.get("cache");
+ return String.format("Aggregated: [%s, %s, %s]", db, api, cache);
+ });
+
+ // Build fan-in pipeline
+ pipeline.addTask("db", dbTask)
+ .addTask("api", apiTask)
+ .addTask("cache", cacheTask)
+ .addTask("aggregate", aggregateTask);
+
+ pipeline.connect("db", "aggregate")
+ .connect("api", "aggregate")
+ .connect("cache", "aggregate");
+
+ // Act
+ Instant start = Instant.now();
+ String result = (String) pipeline.run("input").get(5, TimeUnit.SECONDS);
+ Duration duration = Duration.between(start, Instant.now());
+
+ // Assert
+ assertEquals("Aggregated: [db-data, api-data, cache-data]", result);
+
+ // Should complete in roughly the time of the slowest task (200ms), not the sum
+ assertTrue(duration.toMillis() < 500,
+ "Fan-in should execute in parallel with virtual threads");
+
+ System.out.println("✓ Fan-in pattern completed in " + duration.toMillis() + "ms");
+
+ } finally {
+ pipeline.shutdown();
+ virtualExecutor.shutdown();
+ assertTrue(virtualExecutor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ void shouldNotSufferFromThreadPoolStarvation() throws Exception {
+ // Virtual threads should handle blocking operations without thread pool starvation
+ ExecutorService virtualExecutor = createVirtualThreadExecutor();
+ TaskPipeline pipeline = new TaskPipeline(virtualExecutor);
+
+ try {
+ CountDownLatch allTasksStarted = new CountDownLatch(50);
+ CountDownLatch proceedSignal = new CountDownLatch(1);
+ AtomicInteger startedCount = new AtomicInteger(0);
+
+ // Create a task that blocks until signaled
+ Task blockingTask = (input, context) ->
+ CompletableFuture.completedFuture(input)
+ .thenApplyAsync(i -> {
+ startedCount.incrementAndGet();
+ allTasksStarted.countDown();
+ try {
+ // Block until signaled - this would cause thread pool starvation with platform threads
+ proceedSignal.await(10, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ return "Task " + i + " completed";
+ }, virtualExecutor); // Use the virtual thread executor explicitly
+
+ pipeline.add("blocking", blockingTask);
+
+ // Act: Start many blocking tasks
+ int taskCount = 50;
+ List> futures = IntStream.range(0, taskCount)
+ .mapToObj(i -> pipeline.run(i))
+ .toList();
+
+ // Wait a bit for tasks to start (more lenient timeout)
+ boolean allStarted = allTasksStarted.await(5, TimeUnit.SECONDS);
+
+ // Signal all tasks to proceed
+ proceedSignal.countDown();
+
+ // Wait for completion
+ CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
+ .get(10, TimeUnit.SECONDS);
+
+ // Assert - check that most tasks started (virtual threads should handle many concurrent blocks)
+ int started = startedCount.get();
+ System.out.println("✓ Started " + started + " out of " + taskCount + " blocking tasks with virtual threads");
+
+ assertTrue(started >= taskCount * 0.8,
+ "Most tasks should start with virtual threads (started: " + started + " out of " + taskCount + ")");
+
+ } finally {
+ pipeline.shutdown();
+ virtualExecutor.shutdown();
+ assertTrue(virtualExecutor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ void shouldHandleMixedCPUAndIOTasksEfficiently() throws Exception {
+ ExecutorService virtualExecutor = createVirtualThreadExecutor();
+ TaskPipeline pipeline = new TaskPipeline(virtualExecutor);
+
+ try {
+ // I/O-bound task (benefits from virtual threads)
+ Task ioTask = (input, context) ->
+ CompletableFuture.supplyAsync(() -> {
+ simulateIOOperation(100);
+ context.put("ioResult", "io-done");
+ return input + "-io";
+ });
+
+ // CPU-bound task (doesn't benefit as much from virtual threads)
+ Task cpuTask = (input, context) ->
+ CompletableFuture.supplyAsync(() -> {
+ // Simulate CPU-intensive work
+ long sum = 0;
+ for (int i = 0; i < 1_000_000; i++) {
+ sum += i;
+ }
+ context.put("cpuResult", sum);
+ return input.length() + (int) (sum % 100);
+ });
+
+ // Another I/O task
+ Task finalIoTask = (input, context) ->
+ CompletableFuture.supplyAsync(() -> {
+ simulateIOOperation(50);
+ String ioResult = context.get("ioResult", String.class).orElse("none");
+ Long cpuResult = context.get("cpuResult", Long.class).orElse(0L);
+ return String.format("Final: io=%s, cpu=%d, input=%d", ioResult, cpuResult, input);
+ });
+
+ // Build pipeline
+ pipeline.add("io1", ioTask)
+ .then("cpu", cpuTask)
+ .then("io2", finalIoTask);
+
+ // Act
+ String result = (String) pipeline.run("test").get(5, TimeUnit.SECONDS);
+
+ // Assert
+ assertNotNull(result);
+ assertTrue(result.contains("io=io-done"));
+ assertTrue(result.contains("cpu="));
+
+ System.out.println("✓ Mixed CPU and I/O tasks completed successfully");
+
+ } finally {
+ pipeline.shutdown();
+ virtualExecutor.shutdown();
+ assertTrue(virtualExecutor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ @Test
+ void shouldMaintainThreadLocalSafety() throws Exception {
+ // Virtual threads have their own thread locals
+ ExecutorService virtualExecutor = createVirtualThreadExecutor();
+ TaskPipeline pipeline = new TaskPipeline(virtualExecutor);
+
+ try {
+ ThreadLocal threadLocal = new ThreadLocal<>();
+ ConcurrentHashMap threadValues = new ConcurrentHashMap<>();
+
+ Task taskWithThreadLocal = (input, context) ->
+ CompletableFuture.supplyAsync(() -> {
+ String value = "thread-" + input;
+ threadLocal.set(value);
+
+ // Simulate some work
+ simulateIOOperation(10);
+
+ // Verify thread local is maintained
+ String retrieved = threadLocal.get();
+ threadValues.put(value, retrieved);
+
+ return retrieved;
+ });
+
+ pipeline.add("threadLocalTask", taskWithThreadLocal);
+
+ // Run multiple tasks concurrently
+ List> futures = IntStream.range(0, 50)
+ .mapToObj(pipeline::run)
+ .toList();
+
+ CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
+ .get(5, TimeUnit.SECONDS);
+
+ // Assert - each virtual thread maintained its own thread local
+ assertEquals(50, threadValues.size());
+ threadValues.forEach((key, value) ->
+ assertEquals(key, value, "Thread local should be maintained per virtual thread"));
+
+ System.out.println("✓ Thread locals maintained correctly across " + threadValues.size() + " virtual threads");
+
+ } finally {
+ pipeline.shutdown();
+ virtualExecutor.shutdown();
+ assertTrue(virtualExecutor.awaitTermination(5, TimeUnit.SECONDS));
+ }
+ }
+
+ /**
+ * Simulates an I/O operation that blocks the thread.
+ * Virtual threads handle blocking operations efficiently.
+ */
+ private void simulateIOOperation(int milliseconds) {
+ try {
+ Thread.sleep(milliseconds);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ }
+}
\ No newline at end of file