Bugfix: fix process unknown after agent reconnects (#63)

* Make agent read timeout configurable for testing

* Create failing reproducer test for dropped messages on reconnect

* Check if websocket is closed before sending

* Add retry_exception utility

* Re-choose connection if the conn is dead

* Format with ruff

* Remove outdated comment

* Commit Cargo.lock version bump

* Fix clippy error

Author: gusostow

Cargo.lock

@@ -163,6 +163,16 @@ async def send(self, msg):
             raise WebSocketNotOpen(self.conn_id)
 
         data = serialize_message(msg)
+
+        # In rare cases the websocket may appear closed while it's still alive for tasks
+        # that run on the event loop before the websocket state can be updated. This
+        # forces it to check if it's closed to avoid quietly sending messages into a
+        # dead websocket.
+        try:
+            await asyncio.wait_for(self.ws.wait_closed(), timeout=0)
+        except asyncio.TimeoutError:
+            pass
+
         try:
             await self.ws.send(data)
         except ConnectionClosedError:

python/procstar/agent/conn.py

@@ -15,9 +15,10 @@
 from . import DEFAULT_PORT
 from .conn import Connection, Connections
 from .conn import choose_connection, wait_for_connection
-from .exc import NoConnectionError
+from .exc import NoConnectionError, WebSocketNotOpen
 from .proc import Processes, Process, Result
 from procstar import proto
+from procstar.lib.py import retry_exception
 from procstar.lib.time import now
 
 FROM_ENV = object()
@@ -199,7 +200,11 @@ async def _serve_connection(self, access_token, reconnect_timeout, ws):
         # Add or re-add the connection.
         try:
             conn = self.connections._add(
-                register_msg.conn, register_msg.proc, register_msg.shutdown_state, time, ws
+                register_msg.conn,
+                register_msg.proc,
+                register_msg.shutdown_state,
+                time,
+                ws,
             )
             conn.info.stats.num_received += 1  # the Register message
         except RuntimeError as exc:
@@ -311,13 +316,21 @@ async def request_start(
         except AttributeError:
             pass
 
-        conn = await choose_connection(
-            self.connections,
-            group_id,
-            timeout=conn_timeout,
-        )
-
-        await conn.send(proto.ProcStartRequest(specs={proc_id: spec}))
+        async def choose_and_start():
+            conn = await choose_connection(
+                self.connections,
+                group_id,
+                timeout=conn_timeout,
+            )
+            # raises a WebSocketNotOpen if the conn was closed after choosing the
+            # connection
+            await conn.send(proto.ProcStartRequest(specs={proc_id: spec}))
+            return conn
+
+        # We try twice sequentially to avoid an unfortunately common case where an agent
+        # disconnects directly after the connection is chosen. This is prone to
+        # happening after the event loop is unintentionally blocked for period of time.
+        conn = await retry_exception(choose_and_start, (WebSocketNotOpen,), retries=1, interval=0)
         return self.processes.create(conn, proc_id)
 
     async def start(self, *args, **kw_args) -> (Process, Result):

python/procstar/agent/server.py

@@ -1,3 +1,11 @@
+import asyncio
+import logging
+from typing import Any, Callable, Iterable, Type
+
+
+logger = logging.getLogger(__name__)
+
+
 class Interval:
     """
     Closed-open interval.
@@ -45,3 +53,23 @@ def format_call(__fn, *args, **kw_args) -> str:
 
 def format_ctor(obj, *args, **kw_args):
     return format_call(obj.__class__, *args, **kw_args)
+
+
+async def retry_exception(
+    fn: Callable[[], Any],
+    exc_types: Iterable[Type[BaseException]],
+    *,
+    retries,
+    interval,
+):
+    for retry in range(retries + 1):
+        try:
+            return await fn()
+        except tuple(exc_types) as exc:
+            if retry == retries:
+                logger.error(f"failed after {retries} retries")
+                raise
+            else:
+                logger.warning(str(exc))
+                logger.debug(f"retrying in {interval}s")
+                await asyncio.sleep(interval)

python/procstar/lib/py.py

@@ -208,7 +208,7 @@ async def wait_for_connect(conns):
                 for task in tasks:
                     task.cancel()
 
-    def start_procstar(self, *, group_id=proto.DEFAULT_GROUP, args=[]):
+    def start_procstar(self, *, group_id=proto.DEFAULT_GROUP, args=()):
         """
         Starts a single procstar instance.
         """
@@ -246,6 +246,7 @@ async def start(
         counts={"default": 1},
         access_token=DEFAULT,
         reconnect_timeout=None,
+        args=(),
     ):
         """
         Async context manager for a ready-to-go assembly.
@@ -255,7 +256,7 @@ async def start(
         """
         asm = cls(access_token=access_token, reconnect_timeout=reconnect_timeout)
         await asm.start_server()
-        await asm.start_procstars(counts)
+        await asm.start_procstars(counts, args=args)
         try:
             yield asm
         finally:

python/procstar/testing/agent.py

@@ -204,6 +204,8 @@ pub struct ConnectConfig {
     pub interval_max: Duration,
     /// Maximum consecutive failed connection attempts.
     pub count_max: u64,
+    /// Agent websocket read timeout.
+    pub read_timeout: Duration,
 }
 
 impl Default for ConnectConfig {
@@ -219,6 +221,7 @@ impl ConnectConfig {
             interval_mult: 2.0,
             interval_max: Duration::from_secs(60),
             count_max: u64::MAX,
+            read_timeout: READ_TIMEOUT,
         }
     }
 }
@@ -281,7 +284,7 @@ pub async fn run(
         let mut receiver = tokio_stream::StreamExt::timeout_repeating(
             receiver,
             // Avoid the instant first tick from a standard interval
-            time::interval_at(time::Instant::now() + READ_TIMEOUT, READ_TIMEOUT),
+            time::interval_at(time::Instant::now() + cfg.read_timeout, cfg.read_timeout),
         );
 
         // Simultaneously wait for an incoming websocket message or a
@@ -300,7 +303,7 @@ pub async fn run(
                 res = receiver.next() => {
                     let ws_msg = match res {
                         None => { warn!("msg stream end"); break },
-                        Some(Err(_)) => {warn!("timeout error after {:?}", READ_TIMEOUT); break; }
+                        Some(Err(_)) => {warn!("timeout error after {:?}", cfg.read_timeout); break; }
                         Some(Ok(ws_msg)) => ws_msg,
                     };
 

src/agent.rs

@@ -80,6 +80,9 @@ pub struct Args {
     /// Maximum number of consecutive agent connection attempts
     #[arg(long, value_name = "COUNT")]
     pub connect_count_max: Option<u64>,
+    /// Agent websocket read timeout in seconds; for testing only
+    #[arg(long, hide = true, value_name = "SECS")]
+    pub agent_read_timeout: Option<f64>,
 
     /// Process specification file, or "-" for stdin
     pub input: Option<String>,
@@ -117,5 +120,8 @@ pub fn get_connect_config(args: &Args) -> agent::ConnectConfig {
             .connect_interval_max
             .map_or(df.interval_max, Duration::from_secs_f64),
         count_max: args.connect_count_max.unwrap_or(df.count_max),
+        read_timeout: args
+            .agent_read_timeout
+            .map_or(df.read_timeout, Duration::from_secs_f64),
     }
 }

src/argv.rs

@@ -1,9 +1,11 @@
 import asyncio
 import shutil
+import time
 
 import pytest
 
 from procstar import spec
+from procstar.agent.exc import ProcessUnknownError
 from procstar.agent.proc import ConnectionTimeoutError
 from procstar.agent.server import maybe_set_reconnect_timeout
 from procstar.testing.agent import Assembly
@@ -45,3 +47,32 @@ async def mock_set_reconnect(*args, **kwargs):
 
         res = await asm.wait(proc)
         assert res.status.exit_code == 0
+
+
+@pytest.mark.asyncio
+async def test_dropped_procstart_on_reconnect():
+    """
+    Replicates the bug where ProcStartRequest messages are dropped when
+    the agent reconnects after a ping timeout due to blocked Python event loop.
+
+    This test demonstrates the specific timing issue where messages sent
+    during the brief window between disconnect and successful reconnect are lost.
+    """
+    # Use short read timeout to make test faster
+    async with Assembly.start(counts={"default": 1}, args=["--agent-read-timeout", "1"]) as asm:
+        # block the event loop for longer than agent read timeout
+        time.sleep(2)
+        try:
+            # request a proc start right away before the agent gets a chance to
+            # reconnect
+            proc1, _ = await asm.server.start(
+                "proc1",
+                spec.make_proc(["/usr/bin/true"]),
+                conn_timeout=5,
+            )
+            await anext(proc1.updates)
+        except ProcessUnknownError as e:
+            assert False, f"ProcStartRequest dropped: {e}"
+
+        res = await asm.wait(proc1)
+        assert res.status.exit_code == 0