Skip to content

fix(scheduling): make ScheduledAgentTask start/stop atomic state transitions (issue #215) #370

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
konlue wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(scheduling): make ScheduledAgentTask start/stop atomic state transitions (issue #215) #370

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
140 changes: 108 additions & 32 deletions ...core/src/main/java/io/github/agentic/ai/workflow/graph/scheduling/ScheduledAgentTask.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import java.util.Optional;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

import io.github.agentic.ai.workflow.graph.CompiledGraph;
import io.github.agentic.ai.workflow.graph.OverAllState;
Expand Down Expand Up @@ -51,9 +52,35 @@ public class ScheduledAgentTask {

private volatile ScheduledFuture<?> scheduledFuture;

private volatile boolean started = false;
/**
* Atomic lifecycle state (issue #215): NEW → STARTING → RUNNING → STOPPING →
* STOPPED, with FAILED as the start-failure terminal. A single CAS-able state
* replaces the previous independent started/stopped flags, so concurrent start
* and stop transitions cannot interleave.
*/
enum ScheduleState {

/** Created and registered, not yet scheduled. */
NEW,

/** A start is publishing the schedule; stop may still win the race. */
STARTING,

/** The schedule is live. */
RUNNING,

/** stop() is cancelling and unregistering. */
STOPPING,

/** Terminal after a successful stop; start is rejected. */
STOPPED,

private volatile boolean stopped = false;
/** Terminal after a failed start; registration has been released. */
FAILED

}

private final AtomicReference<ScheduleState> state = new AtomicReference<>(ScheduleState.NEW);

private final AtomicBoolean unregistered = new AtomicBoolean(false);

Expand All @@ -73,35 +100,51 @@ public ScheduledAgentTask(CompiledGraph graph, ScheduleConfig config) {
* Start the scheduled execution
*/
public ScheduledAgentTask start() {
if (started) {
throw new IllegalStateException("Schedule already started");
if (!state.compareAndSet(ScheduleState.NEW, ScheduleState.STARTING)) {
throw new IllegalStateException("Schedule already started or stopped: " + state.get());
}

try {
switch (config.getMode()) {
case CRON:
scheduledFuture = taskScheduler.schedule(this::executeGraph,
new CronTrigger(config.getCronExpression()));
break;
case FIXED_DELAY:
scheduledFuture = taskScheduler.scheduleWithFixedDelay(this::executeGraph,
Instant.now().plusMillis(config.getInitialDelay()), Duration.ofMillis(config.getFixedDelay()));
break;
case FIXED_RATE:
scheduledFuture = taskScheduler.scheduleAtFixedRate(this::executeGraph,
Instant.now().plusMillis(config.getInitialDelay()), Duration.ofMillis(config.getFixedRate()));
break;
case ONE_TIME:
scheduledFuture = taskScheduler.schedule(this::executeGraph,
Instant.now().plusMillis(config.getInitialDelay()));
break;
case TRIGGER:
scheduledFuture = taskScheduler.schedule(this::executeGraph, config.getTrigger());
break;
default:
throw new IllegalStateException("Unexpected value: " + config.getMode());
}
}
catch (RuntimeException e) {
// A failed start must not leave the task registered or scheduled (#215).
// CAS, not set: if stop() already won the STARTING window it owns the
// terminal state, and FAILED must not clobber its STOPPED.
state.compareAndSet(ScheduleState.STARTING, ScheduleState.FAILED);
unregisterTask();
throw e;
}

switch (config.getMode()) {
case CRON:
scheduledFuture = taskScheduler.schedule(this::executeGraph,
new CronTrigger(config.getCronExpression()));
break;
case FIXED_DELAY:
scheduledFuture = taskScheduler.scheduleWithFixedDelay(this::executeGraph,
Instant.now().plusMillis(config.getInitialDelay()), Duration.ofMillis(config.getFixedDelay()));
break;
case FIXED_RATE:
scheduledFuture = taskScheduler.scheduleAtFixedRate(this::executeGraph,
Instant.now().plusMillis(config.getInitialDelay()), Duration.ofMillis(config.getFixedRate()));
break;
case ONE_TIME:
scheduledFuture = taskScheduler.schedule(this::executeGraph,
Instant.now().plusMillis(config.getInitialDelay()));
break;
case TRIGGER:
scheduledFuture = taskScheduler.schedule(this::executeGraph, config.getTrigger());
break;
default:
throw new IllegalStateException("Unexpected value: " + config.getMode());
if (!state.compareAndSet(ScheduleState.STARTING, ScheduleState.RUNNING)) {
// stop() won the race while the future was being published: cancel the late
// publication so it cannot trigger new executions, and finish stopped.
cancelPublishedFuture();
throw new IllegalStateException("Schedule stopped while starting");
}

started = true;
notifyListeners(ScheduleLifecycleListener.ScheduleEvent.STARTED);
return this;
}
Expand All @@ -110,16 +153,44 @@ public ScheduledAgentTask start() {
* Stop the scheduled execution
*/
public void stop() {
if (scheduledFuture != null && !scheduledFuture.isCancelled()) {
scheduledFuture.cancel(false);
// One winner per stop: concurrent and repeated stops return idempotently.
if (!transitionToStopping()) {
return;
}
stopped = true;
cancelPublishedFuture();

// Unregister from the manager that owns this task.
unregisterTask();
log.debug("Stopped and unregistered ScheduledAgentTask with ID: {}", taskId);

notifyListeners(ScheduleLifecycleListener.ScheduleEvent.STOPPED);
state.set(ScheduleState.STOPPED);
}

/**
* Moves a non-terminal state to STOPPING under one CAS winner, so a concurrent or
* repeated stop is a no-op.
* @return true if this caller won the transition and must perform the stop work
*/
private boolean transitionToStopping() {
while (true) {
ScheduleState current = state.get();
if (current == ScheduleState.STOPPING || current == ScheduleState.STOPPED
|| current == ScheduleState.FAILED) {
return false;
}
if (state.compareAndSet(current, ScheduleState.STOPPING)) {
return true;
}
}
}

/** Cancels the published future if present; safe against a late publication race. */
private void cancelPublishedFuture() {
ScheduledFuture<?> active = scheduledFuture;
if (active != null && !active.isCancelled()) {
active.cancel(false);
}
}

/**
Expand Down Expand Up @@ -222,11 +293,16 @@ private void unregisterTask() {

// Getters for monitoring
public boolean isStarted() {
return started;
return state.get() != ScheduleState.NEW && state.get() != ScheduleState.FAILED;
}

public boolean isStopped() {
return stopped;
return state.get() == ScheduleState.STOPPING || state.get() == ScheduleState.STOPPED;
}

/** Test observation point for the lifecycle state machine. */
ScheduleState state() {
return state.get();
}

/**
Expand Down
Loading
Loading