sdk/python: keep a handle's relay connection across calls

docs/sdk-python.md

@@ -170,6 +170,16 @@ automatically with the same transparent wake. A claim with a connection live
 when the sweep checks it (a relay stream, a buffered exec, a preview dial, an
 egress request) is not swept; the idle clock restarts when that connection ends.
 
+Data-plane calls share a handle's relay connection: after a call the SDK
+keeps the connection for 30 seconds (`Client(..., keep_alive=...)` tunes the
+window; 0 dials per call) and the next call on that handle sends its request
+on it, so a busy handle pays the dial, upgrade and TLS handshake once. A kept
+connection counts as live for `idle_hibernate_seconds` until it closes, so
+keep the window below that setting; `close` and `hibernate` drop it at once.
+Streams (`watch`, `open_pty`, `dial_port`, an LSP session) take a connection
+of their own, and a guest whose silkd predates the back-to-back protocol
+gets one connection per call as before.
+
 If that deployment also enables `archive_after_seconds`, archiving replaces
 the original claim deadline with the archive-retention deadline (or no
 deadline when archives are kept forever). Waking an archive starts a fresh

sdk/python/cocoonsandbox/client.py

@@ -28,8 +28,15 @@ class Client:
     """Talks to one sandboxd node (and, transparently, its cluster)."""
 
     def __init__(
-        self, addr: str, api_token: str = "", timeout: float = 120.0, *, ssl_context: ssl.SSLContext | None = None
+        self,
+        addr: str,
+        api_token: str = "",
+        timeout: float = 120.0,
+        *,
+        ssl_context: ssl.SSLContext | None = None,
+        keep_alive: float = 30.0,
     ) -> None:
+        """keep_alive keeps a handle's idle relay connection, which holds the sandbox's idle clock; 0 dials per call."""
         endpoint = _endpoint_url(addr.split(",")[0].strip())
         self.addr = endpoint.geturl().removeprefix("http://")
         self._scheme = endpoint.scheme
@@ -39,6 +46,7 @@ def __init__(
         self._opener: urllib.request.OpenerDirector | None = None
         self.api_token = api_token
         self.timeout = timeout
+        self.keep_alive = keep_alive
 
     def new(
         self,

sdk/python/cocoonsandbox/conn.py

@@ -3,8 +3,10 @@
 from __future__ import annotations
 
 import contextlib
+import select
 import socket
 import ssl
+import threading
 import time
 import urllib.parse
 from collections.abc import Iterator
@@ -13,6 +15,8 @@
 from .errors import APIError, ProtocolError, SilkdError
 from .frames import MAX_FRAME, decode_response, encode_request
 
+KEEP_ALIVE_CONNS = 8
+
 _CloseableT = TypeVar("_CloseableT", bound="_Closeable")
 
 
@@ -69,6 +73,14 @@ def recv_until(self, *terminal: str) -> Iterator[dict[str, Any]]:
             if frame["type"] in terminal:
                 return
 
+    def quiet(self) -> bool:
+        """Reports whether the peer has neither hung up nor spoken since the last frame."""
+        try:
+            readable, _, _ = select.select([self._sock], [], [], 0)
+        except (OSError, ValueError):
+            return False
+        return not readable
+
     def close(self) -> None:
         self.abort()
         try:
@@ -77,6 +89,60 @@ def close(self) -> None:
             self._sock.close()
 
 
+class ConnPool:
+    """Parks a handle's idle relay connections between calls; one sweeper timer closes them as they expire."""
+
+    def __init__(self, idle: float) -> None:
+        self._idle = idle
+        self._lock = threading.Lock()
+        self._parked: list[tuple[float, Conn]] = []
+        self._sweep: threading.Timer | None = None
+
+    def take(self) -> Conn | None:
+        while True:
+            with self._lock:
+                if not self._parked:
+                    return None
+                expires, conn = self._parked.pop()
+            if time.monotonic() < expires and conn.quiet():
+                return conn
+            conn.close()
+
+    def park(self, conn: Conn) -> None:
+        with self._lock:
+            if self._idle > 0 and len(self._parked) < KEEP_ALIVE_CONNS:
+                self._parked.append((time.monotonic() + self._idle, conn))
+                if self._sweep is None:
+                    self._arm(self._idle)
+                return
+        conn.close()
+
+    def drain(self) -> None:
+        with self._lock:
+            parked, self._parked = self._parked, []
+            if self._sweep is not None:
+                self._sweep.cancel()
+                self._sweep = None
+        for _, conn in parked:
+            conn.close()
+
+    def _arm(self, delay: float) -> None:
+        self._sweep = threading.Timer(delay, self._evict)
+        self._sweep.daemon = True
+        self._sweep.start()
+
+    def _evict(self) -> None:
+        now = time.monotonic()
+        with self._lock:
+            expired = [conn for expires, conn in self._parked if expires <= now]
+            self._parked = [entry for entry in self._parked if entry[0] > now]
+            self._sweep = None
+            if self._parked:
+                self._arm(self._parked[0][0] - now)
+        for conn in expired:
+            conn.close()
+
+
 def dial_agent(
     addr: str,
     sandbox_id: str,

sdk/python/cocoonsandbox/frames.py

@@ -8,6 +8,8 @@
 from typing import Any
 
 PROTO_VERSION = 1
+# the info proto from which silkd serves RPCs back to back on one connection
+KEEP_ALIVE_PROTO = 2
 MAX_FRAME = 8 * 1024 * 1024
 FS_CHUNK = 256 * 1024
 # bulk streams chunk larger than silkd's FS_CHUNK: fewer frames per byte, still under MAX_FRAME after base64.

sdk/python/cocoonsandbox/sandbox.py

@@ -10,9 +10,9 @@
 from typing import TYPE_CHECKING, Any, cast
 
 from .checkpoint import Checkpoint
-from .conn import Conn, _Closeable
-from .errors import APIError, ExitError, ProtocolError, SandboxError
-from .frames import BULK_CHUNK, FS_CHUNK
+from .conn import Conn, ConnPool, _Closeable
+from .errors import APIError, ExitError, ProtocolError, SandboxError, SilkdError
+from .frames import BULK_CHUNK, FS_CHUNK, KEEP_ALIVE_PROTO
 from .template import Template
 
 if TYPE_CHECKING:
@@ -41,6 +41,8 @@ def __init__(
         self.from_checkpoint = from_checkpoint
         self.template_digest = template_digest
         self.volumes = [dict(volume) for volume in volumes or []]
+        self._pool = ConnPool(client.keep_alive)
+        self._proto = 0
 
     def __enter__(self) -> Sandbox:
         return self
@@ -97,13 +99,14 @@ def run(
         deadline = None if timeout is None else time.monotonic() + timeout
         expired = threading.Event()
         try:
-            conn = self._dial(deadline)
+            conn = self._connect(deadline)
         except (ProtocolError, TimeoutError):
             if deadline is not None and time.monotonic() >= deadline:
                 raise TimeoutError(f"command did not finish within {timeout}s") from None
             raise
-        with conn:
+        with self._lease(conn):
             watchdog = _arm_watchdog(conn, deadline, expired)
+            pump = threading.Thread(target=_feed_stdin, args=(conn, stdin), daemon=True) if stdin else None
             try:
                 conn.send(
                     "exec",
@@ -114,17 +117,25 @@ def run(
                     detach=False,
                     session=session or None,
                 )
-                pump = threading.Thread(target=_feed_stdin, args=(conn, stdin), daemon=True)
-                pump.start()
+                if pump is not None:
+                    pump.start()
+                else:
+                    with contextlib.suppress(OSError):
+                        conn.send("stdin_close")
                 code = _pump_stdio(conn, on_stdout, on_stderr)
             except (ProtocolError, OSError):
                 if expired.is_set():
                     raise TimeoutError(f"command did not finish within {timeout}s") from None
                 raise
+            except SilkdError:
+                if pump is not None:
+                    pump.join()
+                raise
             finally:
                 if watchdog is not None:
                     watchdog.cancel()
-        pump.join()  # the closed conn fails a stalled send, so this cannot hang
+            if pump is not None:
+                pump.join()  # the pump's frames must not land on the next RPC
         if code is None:
             raise ProtocolError("exec stream ended without an exit frame")
         return code
@@ -151,7 +162,7 @@ def logs(
         on_stderr: Callable[[bytes], object] | None = None,
     ) -> int | None:
         """Replays buffered output and returns the exit code, or None if the process still runs."""
-        return self._drain_proc("logs", pid, on_stdout, on_stderr)
+        return self._drain_proc("logs", pid, on_stdout, on_stderr, trailing_done=True)
 
     def attach(
         self,
@@ -164,19 +175,19 @@ def attach(
 
     def write_file(self, path: str, data: bytes, mode: int | None = None) -> None:
         """Writes data to path atomically (temp + rename on the guest)."""
-        with self._dial() as conn:
+        with self._lease() as conn:
             conn.send("fs_write", path=path, mode=mode)
             _send_chunks(conn, data)
             conn.send("data_end")
             _expect(conn, "done")
 
     def read_file(self, path: str) -> bytes:
-        with self._dial() as conn:
+        with self._lease() as conn:
             conn.send("fs_read", path=path)
             return _drain_data(conn)
 
     def list_dir(self, path: str) -> list[dict[str, Any]]:
-        with self._dial() as conn:
+        with self._lease() as conn:
             conn.send("fs_list", path=path)
             entries: list[dict[str, Any]] = []
             for frame in conn.recv_until("done"):
@@ -198,15 +209,15 @@ def rename(self, src: str, dst: str) -> None:
 
     def push(self, dest: str, tar_stream: bytes) -> None:
         """Extracts a tar stream into dest; a truncated stream leaves dest untouched."""
-        with self._dial() as conn:
+        with self._lease() as conn:
             conn.send("fs_push", dest=dest)
             _send_chunks(conn, tar_stream, chunk=BULK_CHUNK)
             conn.send("data_end")
             _expect(conn, "done")
 
     def pull(self, path: str) -> bytes:
         """Returns path (file or tree) as a tar archive."""
-        with self._dial() as conn:
+        with self._lease() as conn:
             conn.send("fs_pull", path=path)
             return _drain_data(conn)
 
@@ -215,14 +226,14 @@ def find(self, path: str, pattern: str, glob: str = "") -> list[dict[str, Any]]:
 
     def find_iter(self, path: str, pattern: str, glob: str = "") -> Iterator[dict[str, Any]]:
         """Yields matches as they stream."""
-        with self._dial() as conn:
+        with self._lease() as conn:
             conn.send("fs_find", path=path, pattern=pattern, glob=glob or None)
             for f in conn.recv_until("done"):
                 if f["type"] == "match":
                     yield f
 
     def replace(self, files: list[str], pattern: str, replacement: str) -> list[dict[str, Any]]:
-        with self._dial() as conn:
+        with self._lease() as conn:
             conn.send("fs_replace", files=files, pattern=pattern, replacement=replacement)
             return [f for f in conn.recv_until("done") if f["type"] == "replaced"]
 
@@ -284,6 +295,7 @@ def fork(self, count: int, ttl_seconds: int = 0) -> list[Sandbox]:
 
     def hibernate(self) -> None:
         """Snapshots and stops the VM; the next guest call restores its state."""
+        self._pool.drain()
         self._client._request(
             self.owner, "POST", f"/v1/sandboxes/{self.id}/hibernate", None, "hibernate", bearer=self.token
         )
@@ -347,6 +359,7 @@ def dial_port(self, port: int) -> PortConn:
 
     def close(self) -> None:
         """Releases the sandbox; its VM is destroyed."""
+        self._pool.drain()
         try:
             self._client._request(
                 self.owner, "POST", f"/v1/sandboxes/{self.id}/release", None, "release", bearer=self.token
@@ -358,8 +371,49 @@ def close(self) -> None:
     def _dial(self, deadline: float | None = None) -> Conn:
         return self._client._dial(self.owner, self.id, self.token, deadline)
 
+    def _connect(self, deadline: float | None = None) -> Conn:
+        """Takes a parked connection or dials one, asking the daemon's proto on a handle's first kept dial."""
+        conn = self._pool.take()
+        if conn is not None:
+            return conn
+        conn = self._dial(deadline)
+        if self._proto or self._client.keep_alive <= 0:
+            return conn
+        try:
+            conn.send("info")
+            proto = int(_expect(conn, "info").get("proto") or 1)
+        except Exception:
+            conn.close()
+            raise
+        self._proto = proto
+        if proto >= KEEP_ALIVE_PROTO:
+            return conn
+        conn.close()
+        return self._dial(deadline)
+
+    @contextlib.contextmanager
+    def _lease(self, conn: Conn | None = None) -> Iterator[Conn]:
+        """Runs one RPC on conn, dialed when absent; a terminal frame parks it and anything else drops it."""
+        if conn is None:
+            conn = self._connect()
+        try:
+            yield conn
+        except SilkdError:
+            self._park(conn)
+            raise
+        except BaseException:
+            conn.close()
+            raise
+        self._park(conn)
+
+    def _park(self, conn: Conn) -> None:
+        if self._proto >= KEEP_ALIVE_PROTO:
+            self._pool.park(conn)
+        else:
+            conn.close()
+
     def _open_stream(self, op: str, expect: str = "ready", **fields: object) -> tuple[Conn, dict[str, Any]]:
-        conn = self._dial()
+        conn = self._connect()
         try:
             conn.send(op, **fields)
             frame = _expect(conn, expect)
@@ -409,7 +463,7 @@ def pump_out() -> None:
                 local.close()
 
     def _call(self, op: str, expect: str, **fields: object) -> dict[str, Any]:
-        with self._dial() as conn:
+        with self._lease() as conn:
             conn.send(op, **fields)
             return _expect(conn, expect)
 
@@ -422,10 +476,14 @@ def _drain_proc(
         pid: int,
         on_stdout: Callable[[bytes], object] | None,
         on_stderr: Callable[[bytes], object] | None,
+        trailing_done: bool = False,
     ) -> int | None:
-        with self._dial() as conn:
+        with self._lease() as conn:
             conn.send(op, pid=pid)
-            return _pump_stdio(conn, on_stdout, on_stderr)
+            code = _pump_stdio(conn, on_stdout, on_stderr)
+            if trailing_done and code is not None:
+                _expect(conn, "done")
+            return code
 
 
 class Session(_Closeable):