From a5f4ecffb869664ef1b38dce48c85ee32467e14f Mon Sep 17 00:00:00 2001 From: konlue <2362586639@qq.com> Date: Sat, 14 Nov 2026 11:52:06 +0800 Subject: [PATCH] fix(scheduling): make task registration mutually exclusive with shutdown (issue #216) registerTask checked the shutdown flag before acquiring the write lock, so a registration that passed the check could be interleaved after shutdown() completed and stranded inside a shut-down manager's activeTasks forever. The rejection is now re-checked inside the write lock, which is the same boundary shutdown() uses to stop and clear tasks and set the flag, making check-and-register atomic with shutdown. The rejection contract is documented on the ScheduledAgentManager interface. A package-private constructor injects a test scheduler for deterministic interleaving tests: a registration passing the pre-lock check while shutdown holds the write lock in the task stop callback must be rejected by the in-lock re-check and leave no stranded record. Concurrent stop/shutdown stays idempotent without residue. --- .../DefaultScheduledAgentManager.java | 18 +- .../scheduling/ScheduledAgentManager.java | 3 + ...cheduledAgentManagerShutdownMutexTest.java | 197 ++++++++++++++++++ 3 files changed, 216 insertions(+), 2 deletions(-) create mode 100644 argi-orchestration/argi-workflow-graph-core/src/test/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentManagerShutdownMutexTest.java diff --git a/argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/DefaultScheduledAgentManager.java b/argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/DefaultScheduledAgentManager.java index b90d04cc9..a7a64928a 100644 --- a/argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/DefaultScheduledAgentManager.java +++ b/argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/DefaultScheduledAgentManager.java @@ -50,10 +50,18 @@ public class DefaultScheduledAgentManager implements ScheduledAgentManager { private volatile boolean shutdown = false; private DefaultScheduledAgentManager() { - this.taskScheduler = createTaskScheduler(); + this(createTaskScheduler()); log.info("Default Scheduled Agent Manager initialized with shared task scheduler"); } + /** + * Test seam: the eager singleton cannot be resurrected and its real scheduler cannot + * drive deterministic interleaving tests, so tests inject their own scheduler. + */ + DefaultScheduledAgentManager(TaskScheduler taskScheduler) { + this.taskScheduler = taskScheduler; + } + /** * Get the singleton instance of the default manager */ @@ -82,6 +90,12 @@ public String registerTask(ScheduledAgentTask task) { String taskId = String.format("agent-task-%s-%d", task.getName(), taskIdGenerator.getAndIncrement()); lock.writeLock().lock(); try { + // shutdown() sets the flag under the same write lock, so re-checking here + // makes the check-and-register atomic with shutdown: a registration that + // passed the fast-fail check cannot be interleaved into a shut-down manager. + if (shutdown) { + throw new IllegalStateException("Default Scheduled Agent Manager is shut down"); + } activeTasks.put(taskId, task); log.debug("Registered scheduled agent task: {}", taskId); return taskId; @@ -204,7 +218,7 @@ public boolean isShutdown() { return shutdown; } - private TaskScheduler createTaskScheduler() { + private static TaskScheduler createTaskScheduler() { ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler(); scheduler.setPoolSize(3); scheduler.setThreadNamePrefix("agent-scheduler-"); diff --git a/argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentManager.java b/argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentManager.java index 85c41440a..f8b375b54 100644 --- a/argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentManager.java +++ b/argi-orchestration/argi-workflow-graph-core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentManager.java @@ -34,6 +34,9 @@ public interface ScheduledAgentManager { * Register a scheduled agent task * @param task the task to register * @return the unique task ID assigned to this task + * @throws IllegalStateException if the manager is shut down; the rejection is decided + * atomically against shutdown, so a registration racing with an in-progress shutdown + * must also be rejected instead of leaving a stranded task */ String registerTask(ScheduledAgentTask task); diff --git a/argi-orchestration/argi-workflow-graph-core/src/test/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentManagerShutdownMutexTest.java b/argi-orchestration/argi-workflow-graph-core/src/test/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentManagerShutdownMutexTest.java new file mode 100644 index 000000000..274c2b9b5 --- /dev/null +++ b/argi-orchestration/argi-workflow-graph-core/src/test/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentManagerShutdownMutexTest.java @@ -0,0 +1,197 @@ +/* + * Copyright 2025-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.workflow.graph.scheduling; + +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import io.github.agentic.ai.workflow.graph.CompiledGraph; +import io.github.agentic.ai.workflow.graph.KeyStrategy; +import io.github.agentic.ai.workflow.graph.StateGraph; +import io.github.agentic.ai.workflow.graph.action.AsyncNodeAction; +import io.github.agentic.ai.workflow.graph.scheduling.ScheduledAgentTaskLifecycleTest.RecordingTaskScheduler; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import static io.github.agentic.ai.workflow.graph.StateGraph.END; +import static io.github.agentic.ai.workflow.graph.StateGraph.START; +import static io.github.agentic.ai.workflow.graph.scheduling.ScheduleLifecycleListener.ScheduleEvent.STOPPED; +import static io.github.agentic.ai.workflow.graph.scheduling.ScheduledAgentTaskLifecycleTest.graphFailingFirstAttempts; +import static org.junit.jupiter.api.Assertions.assertAll; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** 复现 #216：注册与 shutdown 必须原子互斥，通过锁外检查的注册不得滞留进已关闭的管理器。 */ +class ScheduledAgentManagerShutdownMutexTest { + + private final ScheduledAgentManagerFactory factory = ScheduledAgentManagerFactory.getInstance(); + + private final ScheduledAgentManager originalManager = factory.getManager(); + + /** 恢复测试前的管理器，避免全局 provider 影响其他测试。 */ + @AfterEach + void restoreOriginalManager() { + factory.registerProvider(() -> originalManager); + } + + /** 验证 shutdown 期间通过锁外检查的注册，在关闭完成后被锁内复查拒绝且不产生滞留记录。 */ + @Test + void aRegistrationPassingTheCheckDuringShutdownCannotEnterActiveTasks() throws Exception { + DefaultScheduledAgentManager manager = new DefaultScheduledAgentManager(new RecordingTaskScheduler()); + factory.registerProvider(() -> manager); + CountDownLatch stopCallbackEntered = new CountDownLatch(1); + CountDownLatch releaseStopCallback = new CountDownLatch(1); + new ScheduledAgentTask(graphFailingFirstAttempts(new AtomicInteger(), 0), + ScheduleConfig.builder().initialDelay(0).addListener(new ScheduleLifecycleListener() { + @Override + public void onEvent(ScheduleEvent event, Object data) { + if (event == STOPPED) { + stopCallbackEntered.countDown(); + try { + releaseStopCallback.await(10, TimeUnit.SECONDS); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + } + }).build()); + assertEquals(1, manager.getActiveTaskCount()); + + CountDownLatch nameRead = new CountDownLatch(1); + CountDownLatch releaseName = new CountDownLatch(1); + AtomicReference registrationFailure = new AtomicReference<>(); + Thread shutdownThread = new Thread(manager::shutdown, "shutdown-thread"); + try { + shutdownThread.start(); + // shutdown 持有写锁停在 stop 回调里，此刻 shutdown 标志仍为 false。 + assertTrue(stopCallbackEntered.await(5, TimeUnit.SECONDS), + "shutdown must hold the write lock inside the stop callback"); + // registerTask 在锁外检查与加锁之间调用 task.getName()，用它做可控缝隙。 + Thread registrationThread = new Thread(() -> { + try { + new ScheduledAgentTask(graphWithNameGatedBy(nameRead, releaseName), + ScheduleConfig.builder().initialDelay(0).build()); + } + catch (Throwable throwable) { + registrationFailure.set(throwable); + } + }, "registration-thread"); + registrationThread.start(); + assertTrue(nameRead.await(5, TimeUnit.SECONDS), + "registration must pass the pre-lock check while shutdown is in progress"); + + releaseStopCallback.countDown(); + shutdownThread.join(TimeUnit.SECONDS.toMillis(5)); + releaseName.countDown(); + registrationThread.join(TimeUnit.SECONDS.toMillis(5)); + + assertAll(() -> assertFalse(shutdownThread.isAlive(), "shutdown must finish"), + () -> assertFalse(registrationThread.isAlive(), "registration must finish"), + () -> assertTrue(registrationFailure.get() instanceof IllegalStateException, + "A registration racing with shutdown must be rejected, got: " + registrationFailure.get()), + () -> assertEquals(0, manager.getActiveTaskCount(), + "No task may be stranded in a shut-down manager's activeTasks"), + () -> assertTrue(manager.isShutdown())); + } + finally { + releaseStopCallback.countDown(); + releaseName.countDown(); + shutdownThread.join(TimeUnit.SECONDS.toMillis(5)); + } + } + + /** 验证关闭完成后的注册被明确拒绝，且不留任何任务记录。 */ + @Test + void aRegistrationAfterShutdownIsRejectedAndLeavesNoResidue() throws Exception { + DefaultScheduledAgentManager manager = new DefaultScheduledAgentManager(new RecordingTaskScheduler()); + factory.registerProvider(() -> manager); + new ScheduledAgentTask(graphFailingFirstAttempts(new AtomicInteger(), 0), + ScheduleConfig.builder().initialDelay(0).build()); + assertEquals(1, manager.getActiveTaskCount()); + + manager.shutdown(); + assertThrows(IllegalStateException.class, () -> new ScheduledAgentTask( + graphFailingFirstAttempts(new AtomicInteger(), 0), ScheduleConfig.builder().initialDelay(0).build())); + + assertAll(() -> assertTrue(manager.isShutdown()), () -> assertEquals(0, manager.getActiveTaskCount())); + } + + /** 验证并发 stop 与 shutdown 保持幂等：结束后无残留记录，所有任务都已停止。 */ + @Test + void concurrentStopsAndShutdownsRemainIdempotentWithoutResidue() throws Exception { + DefaultScheduledAgentManager manager = new DefaultScheduledAgentManager(new RecordingTaskScheduler()); + factory.registerProvider(() -> manager); + ScheduledAgentTask first = new ScheduledAgentTask(graphFailingFirstAttempts(new AtomicInteger(), 0), + ScheduleConfig.builder().initialDelay(0).build()); + ScheduledAgentTask second = new ScheduledAgentTask(graphFailingFirstAttempts(new AtomicInteger(), 0), + ScheduleConfig.builder().initialDelay(0).build()); + assertEquals(2, manager.getActiveTaskCount()); + + AtomicReference failure = new AtomicReference<>(); + Thread[] workers = new Thread[] { new Thread(manager::shutdown, "shutdown-1"), + new Thread(manager::shutdown, "shutdown-2"), new Thread(first::stop, "stop-1"), + new Thread(first::stop, "stop-1-again"), new Thread(second::stop, "stop-2"), + new Thread(second::stop, "stop-2-again") }; + for (Thread worker : workers) { + worker.setUncaughtExceptionHandler((thread, throwable) -> failure.compareAndSet(null, throwable)); + worker.start(); + } + for (Thread worker : workers) { + worker.join(TimeUnit.SECONDS.toMillis(5)); + } + + assertAll(() -> assertEquals(null, failure.get(), "concurrent stop and shutdown must not throw"), + () -> assertTrue(manager.isShutdown()), + () -> assertEquals(0, manager.getActiveTaskCount(), "no task may survive shutdown"), + () -> assertTrue(first.isStopped()), () -> assertTrue(second.isStopped())); + } + + /** + * Builds a graph whose {@code getName()} is gated on the given latches, hooking the + * gap between the shutdown check and the write lock inside {@code registerTask}. + */ + private static CompiledGraph graphWithNameGatedBy(CountDownLatch entered, CountDownLatch release) throws Exception { + Map strategies = new HashMap<>(); + strategies.put("result", (left, right) -> right); + StateGraph graph = new StateGraph("name-gated-registration", () -> strategies) { + + @Override + public String getName() { + entered.countDown(); + try { + release.await(10, TimeUnit.SECONDS); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return "name-gated-registration"; + } + + }; + graph.addNode("attempt", AsyncNodeAction.node_async(state -> Map.of("result", "ok"))) + .addEdge(START, "attempt") + .addEdge("attempt", END); + return graph.compile(); + } + +}