Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions src/main/java/com/openjiuwen/core/graph/pregel/NodeTask.java
Original file line number Diff line number Diff line change
Expand Up @@ -26,10 +26,14 @@
*/
public class NodeTask implements Callable<Object> {
private static final LoggerProtocol logger = Loggers.GRAPH;
private static final Runnable NO_OP_INVOCATION_COMPLETED_HANDLER = () -> {
// Standalone NodeTask execution has no executor lifecycle to notify.
};

private final PregelNode node;
private final PregelConfig config;
private final int version;
private final Runnable invocationCompletedHandler;

/**
* NodeTask.
Expand All @@ -40,9 +44,22 @@ public class NodeTask implements Callable<Object> {
* @since 0.1.7
*/
public NodeTask(PregelNode node, PregelConfig config, int version) {
this(node, config, version, NO_OP_INVOCATION_COMPLETED_HANDLER);
}

/**
* Creates a node task with an invocation lifecycle callback.
*
* @param node node to execute
* @param config Pregel configuration
* @param version node version
* @param invocationCompletedHandler callback invoked after the node function returns normally
*/
NodeTask(PregelNode node, PregelConfig config, int version, Runnable invocationCompletedHandler) {
this.node = node;
this.config = config;
this.version = version;
this.invocationCompletedHandler = invocationCompletedHandler;
}

/**
Expand Down Expand Up @@ -79,6 +96,7 @@ public Object call() throws Exception {

// Invoke the node function
invokeFunc(func, kwargs);
invocationCompletedHandler.run();
throwIfInterrupted();

// Route messages
Expand Down
127 changes: 117 additions & 10 deletions src/main/java/com/openjiuwen/core/graph/pregel/TaskExecutorPool.java
Original file line number Diff line number Diff line change
Expand Up @@ -81,10 +81,11 @@ public void submit(PregelNode node, int version) {
RunningTask runningTask = new RunningTask(node);
TaskFuture execution = new TaskFuture(() -> {
if (!runningTask.tryStart()) {
runningTask.complete();
throw new CancellationException("Pregel node task cancelled before execution");
}
try {
return new NodeTask(node, config, version).call();
return new NodeTask(node, config, version, runningTask::markInvocationCompleted).call();
} finally {
runningTask.complete();
}
Expand Down Expand Up @@ -127,6 +128,7 @@ public void waitAll() throws Exception {
} catch (InterruptedException e) {
cancelPendingFutures(futures);
awaitActualCompletion(futures);
Thread.currentThread().interrupt();
throw new CancellationException("Pregel task execution cancelled");
} catch (ExecutionException | CancellationException ignored) {
// Individual errors will be handled below
Expand All @@ -139,8 +141,14 @@ public void waitAll() throws Exception {
}
}
if (!pendingFutures.isEmpty() && !allDone.isDone()) {
cancelPendingFutures(pendingFutures);
awaitActualCompletion(pendingFutures);
try {
settlePendingTasksAfterFailure(pendingFutures);
} catch (InterruptedException exception) {
cancelPendingFutures(pendingFutures);
awaitActualCompletion(pendingFutures);
Thread.currentThread().interrupt();
throw new CancellationException("Pregel task settlement interrupted");
}
}

// Check if the first failure is only a GraphInterrupt (not a real exception).
Expand Down Expand Up @@ -317,6 +325,38 @@ private void cancelPendingFutures(List<CompletableFuture<Object>> pendingFutures
}
}

/**
* Requests cooperative cancellation after a sibling failure and waits until every affected task settles.
*
* @param pendingFutures unfinished sibling futures
* @throws InterruptedException if the caller is interrupted while waiting
* @since 0.1.7
*/
private void settlePendingTasksAfterFailure(List<CompletableFuture<Object>> pendingFutures)
throws InterruptedException {
List<CompletableFuture<Void>> settlements = new ArrayList<>();
for (CompletableFuture<Object> future : pendingFutures) {
if (future.isDone()) {
continue;
}
RunningTask runningTask = runningTasks.get(future);
if (runningTask == null) {
future.cancel(false);
settlements.add(future.handle((value, throwable) -> null));
continue;
}
settlements.add(runningTask.cancelForSiblingFailure(future));
}
if (settlements.isEmpty()) {
return;
}
try {
CompletableFuture.allOf(settlements.toArray(new CompletableFuture[0])).get();
} catch (ExecutionException exception) {
throw new IllegalStateException("Failed to settle graph tasks after a sibling failure", exception);
}
}

/**
* Request cancellation of an executor-managed task.
*
Expand Down Expand Up @@ -384,8 +424,9 @@ private static final class RunningTask {
private final ReentrantLock lifecycleLock = new ReentrantLock();

private FutureTask<Object> execution;
private boolean hasStarted;
private boolean isCancellationRequested;
private Thread executionThread;
private TaskPhase phase = TaskPhase.NOT_STARTED;
private TaskCancellation cancellation = TaskCancellation.NONE;

private RunningTask(PregelNode node) {
this.node = node;
Expand All @@ -395,8 +436,12 @@ private void attach(FutureTask<Object> taskExecution) {
lifecycleLock.lock();
try {
execution = taskExecution;
if (isCancellationRequested) {
if (cancellation == TaskCancellation.FORCED) {
execution.cancel(true);
} else if (cancellation == TaskCancellation.SIBLING_FAILURE
&& phase == TaskPhase.NOT_STARTED
&& execution.cancel(false)) {
completion.complete(null);
}
} finally {
lifecycleLock.unlock();
Expand All @@ -406,10 +451,11 @@ private void attach(FutureTask<Object> taskExecution) {
private boolean tryStart() {
lifecycleLock.lock();
try {
if (isCancellationRequested) {
if (cancellation != TaskCancellation.NONE) {
return false;
}
hasStarted = true;
phase = TaskPhase.INVOKING;
executionThread = Thread.currentThread();
return true;
} finally {
lifecycleLock.unlock();
Expand All @@ -419,15 +465,56 @@ private boolean tryStart() {
private void cancel() {
lifecycleLock.lock();
try {
isCancellationRequested = true;
if (execution != null && execution.cancel(true) && !hasStarted) {
cancellation = TaskCancellation.FORCED;
if (execution != null && execution.cancel(true) && phase == TaskPhase.NOT_STARTED) {
completion.complete(null);
}
} finally {
lifecycleLock.unlock();
}
}

private CompletableFuture<Void> cancelForSiblingFailure(CompletableFuture<Object> result) {
lifecycleLock.lock();
try {
if (phase == TaskPhase.NOT_STARTED && cancellation != TaskCancellation.FORCED) {
cancellation = TaskCancellation.SIBLING_FAILURE;
if (execution == null) {
result.cancel(false);
completion.complete(null);
} else if (execution.cancel(false)) {
completion.complete(null);
}
} else if (phase == TaskPhase.INVOKING && cancellation != TaskCancellation.FORCED) {
cancellation = TaskCancellation.SIBLING_FAILURE;
if (executionThread != null) {
executionThread.interrupt();
}
}
CompletableFuture<Void> resultCompletion = result.handle((value, throwable) -> null);
return CompletableFuture.allOf(completion, resultCompletion);
} finally {
lifecycleLock.unlock();
}
}

private void markInvocationCompleted() {
lifecycleLock.lock();
try {
if (phase != TaskPhase.INVOKING) {
return;
}
phase = TaskPhase.ROUTING;
if (cancellation == TaskCancellation.SIBLING_FAILURE
&& executionThread == Thread.currentThread()) {
// Clear only the cooperative sibling-failure signal before routing starts.
Thread.interrupted();
}
} finally {
lifecycleLock.unlock();
}
}

private PregelNode node() {
return node;
}
Expand All @@ -437,10 +524,30 @@ private CompletableFuture<Void> completion() {
}

private void complete() {
lifecycleLock.lock();
try {
phase = TaskPhase.COMPLETED;
executionThread = null;
} finally {
lifecycleLock.unlock();
}
completion.complete(null);
}
}

private enum TaskPhase {
NOT_STARTED,
INVOKING,
ROUTING,
COMPLETED
}

private enum TaskCancellation {
NONE,
SIBLING_FAILURE,
FORCED
}

/**
* Bridges an executor-managed task to the result future used by the Pregel pool.
*
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,12 @@

package com.openjiuwen.core.graph.pregel;

import static org.junit.jupiter.api.Assertions.*;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Nested;
Expand All @@ -14,6 +19,7 @@
import java.util.concurrent.Callable;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

/**
Expand Down Expand Up @@ -71,13 +77,75 @@ void testPoolRuntimeException() throws Exception {
assertTrue(pool.getFailed().containsKey("B"));
assertEquals("__error__", pool.getFailed().get("B").getStatus());

// A propagates interruption and remains pending for recovery.
assertTrue(pool.getFailed().containsKey("A"));

// C succeeds
assertFalse(pool.getFailed().containsKey("C"));
// C's messages should be collected
boolean cMessageFound = pool.getSucceedMessages().stream().anyMatch(m -> "C".equals(m.getSender()));
assertTrue(cMessageFound, "C's success messages should be collected");
}

@Test
@DisplayName("Sibling failure preserves a node invocation that returns normally")
void testSiblingFailurePreservesNormallyReturnedInvocation() throws Exception {
CountDownLatch invocationStarted = new CountDownLatch(1);
CountDownLatch invocationBlocker = new CountDownLatch(1);
AtomicBoolean invocationInterrupted = new AtomicBoolean(false);
AtomicBoolean routerSawInterruption = new AtomicBoolean(false);
AtomicInteger invocationCount = new AtomicInteger(0);
AtomicInteger routingCount = new AtomicInteger(0);
Callable<Object> successfulTask = () -> {
int currentInvocation = invocationCount.incrementAndGet();
if (currentInvocation == 1) {
invocationStarted.countDown();
try {
if (invocationBlocker.await(5L, TimeUnit.SECONDS)) {
throw new IllegalStateException("Invocation blocker was released unexpectedly");
}
} catch (InterruptedException exception) {
invocationInterrupted.set(true);
Thread.currentThread().interrupt();
}
}
return null;
};
Callable<Object> failingTask = () -> {
if (!invocationStarted.await(1L, TimeUnit.SECONDS)) {
throw new IllegalStateException("Successful sibling did not start");
}
throw new IllegalStateException("Sibling failed");
};
IRouter router = sourceNode -> {
routingCount.incrementAndGet();
routerSawInterruption.set(Thread.currentThread().isInterrupted());
return List.of(new TriggerMessage(sourceNode, "target"));
};
PregelNode successfulNode = new PregelNode(
"successful",
successfulTask,
List.of(router));
PregelNode failingNode = new PregelNode("failing", failingTask, List.of());
TaskExecutorPool pool = new TaskExecutorPool(new PregelConfig());
pool.submit(successfulNode, 1);
pool.submit(failingNode, 1);

IllegalStateException exception = assertThrows(IllegalStateException.class, pool::waitAll);
if (pool.getFailed().containsKey("successful")) {
new NodeTask(successfulNode, new PregelConfig(), 1).call();
}

assertEquals("Sibling failed", exception.getMessage());
assertTrue(invocationInterrupted.get());
assertFalse(routerSawInterruption.get(), "A sibling-failure interrupt must not leak into routing");
assertEquals(1, invocationCount.get(), "A normally returned invocation must not execute again");
assertEquals(1, routingCount.get());
assertFalse(pool.getFailed().containsKey("successful"));
assertTrue(pool.getSucceedMessages().stream()
.anyMatch(message -> "successful".equals(message.getSender())));
}

@Test
@DisplayName("Interrupt exception: B interrupts, A may complete, C succeeds")
void testPoolInterruptException() throws Exception {
Expand Down Expand Up @@ -189,6 +257,38 @@ void testCancelAllInterruptsRunningTask() throws Exception {

assertTrue(interrupted.await(1, TimeUnit.SECONDS));
}

@Test
@DisplayName("cancelAll interrupts routing after the node invocation returns")
void testCancelAllInterruptsRouting() throws Exception {
CountDownLatch routingStarted = new CountDownLatch(1);
CountDownLatch routingBlocker = new CountDownLatch(1);
CountDownLatch routingInterrupted = new CountDownLatch(1);
AtomicInteger invocationCount = new AtomicInteger(0);
IRouter blockingRouter = sourceNode -> {
routingStarted.countDown();
try {
routingBlocker.await();
} catch (InterruptedException exception) {
routingInterrupted.countDown();
Thread.currentThread().interrupt();
}
return List.of(new TriggerMessage(sourceNode, "target"));
};
PregelNode node = new PregelNode(
"routing",
(Runnable) invocationCount::incrementAndGet,
List.of(blockingRouter));
TaskExecutorPool pool = new TaskExecutorPool(new PregelConfig());
pool.submit(node, 1);

assertTrue(routingStarted.await(1L, TimeUnit.SECONDS));
pool.cancelAll();

assertEquals(1, invocationCount.get());
assertTrue(routingInterrupted.await(1L, TimeUnit.SECONDS));
assertTrue(pool.getFailed().containsKey("routing"));
}
}

// ---------- Node function invocation tests ----------
Expand Down