diff --git a/src/argv.rs b/src/argv.rs index 6fa3b58..7b9aee8 100644 --- a/src/argv.rs +++ b/src/argv.rs @@ -45,6 +45,9 @@ pub struct Args { /// When idle (all processes deleted), exit #[arg(short = 'w', long)] pub wait: bool, + /// Timeout in seconds for wait mode when no work is assigned + #[arg(long, value_name = "SECS", default_value_t = 600)] + pub wait_timeout: u64, /// Run an HTTP service #[arg(short, long)] diff --git a/src/main.rs b/src/main.rs index 6605d45..e88638e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,7 +6,7 @@ use log::*; // use procstar::fd::parse_fd; use procstar::agent; use procstar::http; -use procstar::procs::{restrict_exe, start_procs, SharedProcs}; +use procstar::procs::{restrict_exe, start_procs, Notification, NotificationSub, SharedProcs}; use procstar::proto; use procstar::res; use procstar::shutdown; @@ -15,6 +15,7 @@ use procstar::sig::{SIGINT, SIGQUIT, SIGTERM, SIGUSR1}; use procstar::spec; use procstar::systemd::api::{maybe_connect, SharedSystemdClient}; use std::rc::Rc; +use std::time::Duration; //------------------------------------------------------------------------------ @@ -99,17 +100,53 @@ async fn maybe_run_until_exit(args: &argv::Args, procs: &SharedProcs) { } } -async fn maybe_run_until_idle(args: &argv::Args, procs: &SharedProcs) { - if args.wait { - // Run until no processes are left, or until we receive a shutdown - // signal. +async fn wait_for_first_assignment(mut sub: NotificationSub) { + while let Some(notification) = sub.recv().await { + if let Notification::Start(_) = notification { + return; + } + } +} + +async fn wait_until_idle_then_shutdown(procs: &SharedProcs) { + tokio::select! { + _ = procs.wait_idle() => {}, + _ = procs.wait_for_shutdown() => {}, + }; + procs.set_shutdown(shutdown::State::Done); +} + +async fn run_agent_until_idle(args: &argv::Args, procs: &SharedProcs) { + if procs.is_empty() { tokio::select! { - _ = procs.wait_idle() => {}, - _ = procs.wait_for_shutdown() => {}, - }; + _ = procs.wait_for_shutdown() => return, + result = tokio::time::timeout( + Duration::from_secs(args.wait_timeout), + wait_for_first_assignment(procs.subscribe()), + ) => { + if result.is_err() { + warn!( + "agent timeout: no work assigned after {} seconds, shutting down", + args.wait_timeout + ); + procs.set_shutdown(shutdown::State::Done); + return; + } + } + } + } + wait_until_idle_then_shutdown(procs).await; +} - // Ready to shut down now. - procs.set_shutdown(shutdown::State::Done); +async fn maybe_run_until_idle(args: &argv::Args, procs: &SharedProcs) { + if !args.wait { + return; + } + + if args.agent { + run_agent_until_idle(args, procs).await; + } else { + wait_until_idle_then_shutdown(procs).await; } } diff --git a/tests/int/agent/test_wait_mode.py b/tests/int/agent/test_wait_mode.py new file mode 100644 index 0000000..f3693e8 --- /dev/null +++ b/tests/int/agent/test_wait_mode.py @@ -0,0 +1,115 @@ +import asyncio +import pytest +import signal + +from procstar import spec +from procstar.agent.proc import Result +from procstar.agent.exc import NoOpenConnectionInGroup +from procstar.testing.agent import Assembly + + +@pytest.mark.asyncio +async def test_agent_selection_logic(): + """ + Test that the agent selection logic works correctly. + """ + async with Assembly.start(args=["--wait"]) as asm: + available_conns = asm.server.connections._get_open_conns_in_group("default") + + assert len(available_conns) == 1, ( + f"Expected 1 available connection, got {len(available_conns)}" + ) + + conn = available_conns[0] + assert conn.shutdown_state.name == "active" + + +@pytest.mark.asyncio +async def test_agent_state_transitions(): + async with Assembly.start(args=["--wait"]) as asm: + conn = next(iter(asm.server.connections.values())) + assert len(asm.server.connections) >= 1 + assert conn.shutdown_state.name == "active", ( + f"Agent should be active when idle, got: {conn.shutdown_state.name}" + ) + + await asyncio.sleep(1) + # agent stays active until no work is assigned + assert conn.shutdown_state.name == "active", ( + f"Agent should remain active until some work is assigned, got: {conn.shutdown_state.name}" + ) + + # Start a process that runs for a short time + proc_spec = spec.Proc(["/bin/sleep", "0.2"]) + proc, result = await asm.server.start( + proc_id="test-proc", group_id="default", spec=proc_spec + ) + + async for update in proc.updates: + if isinstance(update, Result) and update.state != "running": + break + + conn_id = conn.info.conn.conn_id + conn = asm.server.connections.get(conn_id) + assert conn.shutdown_state.name == "active", ( + "Agent should stay active until all processes are deleted" + ) + + # Delete the completed process to trigger the shutdown sequence + await proc.delete() + + await asyncio.sleep(1) + + active_conns = asm.server.connections._get_open_conns_in_group("default") + assert len(active_conns) == 0, f"Expected no active connections, got {len(active_conns)}" + + # Attempt to start second process should fail due to no available connections + proc_spec2 = spec.Proc(["/bin/echo", "second"]) + + with pytest.raises(NoOpenConnectionInGroup): + await asm.server.start( + proc_id="test-proc-2", + group_id="default", + spec=proc_spec2, + conn_timeout=0.5, + ) + + +@pytest.mark.asyncio +async def test_agent_shutdown_before_receiving_work(): + """ + Test that agent responds to shutdown signals even before receiving work. + """ + async with Assembly.start(args=["--wait"]) as asm: + conn = next(iter(asm.server.connections.values())) + conn_id = conn.info.conn.conn_id + + # Send shutdown signal to the procstar process + procstar_proc = asm.conn_procs[conn_id] + procstar_proc.send_signal(signal.SIGUSR1) + + # Wait for graceful shutdown + await asyncio.sleep(0.5) + + # Connection should be cleaned up + assert conn_id not in asm.server.connections + + +@pytest.mark.asyncio +async def test_agent_timeout_no_work(): + """ + Test that agent times out and shuts down when no work is assigned within the timeout period. + """ + async with Assembly.start(args=["--wait", "--wait-timeout", "1"]) as asm: + # Verify agent starts with active connection + active_conns = asm.server.connections._get_open_conns_in_group("default") + assert len(active_conns) == 1, f"Expected 1 active connection, got {len(active_conns)}" + + # Wait for timeout + await asyncio.sleep(1) + + # Agent should have shut down due to timeout - no more active connections + active_conns = asm.server.connections._get_open_conns_in_group("default") + assert len(active_conns) == 0, ( + f"Expected no active connections after timeout, got {len(active_conns)}" + )