sdk/python: a handle's first call sends info and its request together

docs/sandboxd-api.md

@@ -820,11 +820,11 @@ the connection is a byte-for-byte relay to the guest's silkd, carrying RPCs
 back to back (see [silkd](silkd.md)). An open relay holds the sandbox's idle
 clock, so a client that keeps a connection warm must close it when idle for
 `idle_hibernate_seconds` to apply. Frames sent right behind the request,
-before the 101, reach silkd in order; the Go SDK sends its first frames that
-way over plain HTTP, so a proxy in front of a plain-HTTP endpoint must forward
-them too (Go's `httputil.ReverseProxy` drops them; Caddy forwards them). 426
-without the upgrade header, 404 unknown sandbox, 409 a failed sandbox, 502
-guest unreachable.
+before the 101, reach silkd in order; the Go and Python SDKs send their first
+frames that way over plain HTTP, so a proxy in front of a plain-HTTP endpoint
+must forward them too (Go's `httputil.ReverseProxy` drops them; Caddy forwards
+them). 426 without the upgrade header, 404 unknown sandbox, 409 a failed
+sandbox, 502 guest unreachable.
 
 ## GET /v1/sandboxes/{id}/ports/{port}
 

sdk/python/cocoonsandbox/sandbox.py

@@ -104,7 +104,16 @@ def run(
         expired = threading.Event()
         started: deque[int] = deque(maxlen=1)
         try:
-            conn = self._connect(deadline)
+            conn = self._connect(
+                "exec",
+                deadline,
+                argv=argv,
+                cwd=cwd or None,
+                env=env,
+                user=user or None,
+                detach=False,
+                session=session or None,
+            )
         except (ProtocolError, TimeoutError):
             if deadline is not None and time.monotonic() >= deadline:
                 raise SandboxTimeout(f"command did not finish within {timeout}s") from None
@@ -113,15 +122,6 @@ def run(
             watchdog = None if deadline is None else _arm_watchdog(conn, deadline, expired)
             pump = threading.Thread(target=_feed_stdin, args=(conn, stdin), daemon=True) if stdin else None
             try:
-                conn.send(
-                    "exec",
-                    argv=argv,
-                    cwd=cwd or None,
-                    env=env,
-                    user=user or None,
-                    detach=False,
-                    session=session or None,
-                )
                 if pump is not None:
                     pump.start()
                 else:
@@ -182,18 +182,15 @@ 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._lease() as conn:
-            conn.send("fs_write", path=path, mode=mode)
+        with self._lease(self._connect("fs_write", path=path, mode=mode)) as conn:
             _upload(conn, data, FS_CHUNK)
 
     def read_file(self, path: str) -> bytes:
-        with self._lease() as conn:
-            conn.send("fs_read", path=path)
+        with self._lease(self._connect("fs_read", path=path)) as conn:
             return _drain_data(conn)
 
     def list_dir(self, path: str) -> list[dict[str, Any]]:
-        with self._lease() as conn:
-            conn.send("fs_list", path=path)
+        with self._lease(self._connect("fs_list", path=path)) as conn:
             entries: list[dict[str, Any]] = []
             for frame in conn.recv_until("done"):
                 entries.extend(frame.get("entries") or [])
@@ -214,30 +211,26 @@ 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._lease() as conn:
-            conn.send("fs_push", dest=dest)
+        with self._lease(self._connect("fs_push", dest=dest)) as conn:
             _upload(conn, tar_stream, BULK_CHUNK)
 
     def pull(self, path: str) -> bytes:
         """Returns path (file or tree) as a tar archive."""
-        with self._lease() as conn:
-            conn.send("fs_pull", path=path)
+        with self._lease(self._connect("fs_pull", path=path)) as conn:
             return _drain_data(conn)
 
     def find(self, path: str, pattern: str, glob: str = "") -> list[dict[str, Any]]:
         return list(self.find_iter(path, pattern, glob))
 
     def find_iter(self, path: str, pattern: str, glob: str = "") -> Iterator[dict[str, Any]]:
         """Yields matches as they stream."""
-        with self._lease() as conn:
-            conn.send("fs_find", path=path, pattern=pattern, glob=glob or None)
+        with self._lease(self._connect("fs_find", path=path, pattern=pattern, glob=glob or None)) as conn:
             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._lease() as conn:
-            conn.send("fs_replace", files=files, pattern=pattern, replacement=replacement)
+        with self._lease(self._connect("fs_replace", files=files, pattern=pattern, replacement=replacement)) as conn:
             return [f for f in conn.recv_until("done") if f["type"] == "replaced"]
 
     def git_clone(self, url: str, path: str, branch: str = "", depth: int = 0, auth: str = "") -> None:
@@ -385,33 +378,34 @@ 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."""
+    def _connect(self, op: str, deadline: float | None = None, /, **fields: object) -> Conn:
+        """Sends the request on a parked or new connection, behind an info probe 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
+        if conn is None:
+            conn = self._dial(deadline)
+        probe = not self._proto and self._client.keep_alive > 0
+        bounded = probe or deadline is not None
         try:
-            conn.settimeout(remaining_timeout(self._client.timeout, deadline, "agent probe"))
-            conn.send("info")
-            proto = int(_expect(conn, "info").get("proto") or 1)
-            conn.settimeout(None)
+            if bounded:
+                conn.settimeout(remaining_timeout(self._client.timeout, deadline, "agent request"))
+            if probe:
+                conn.send("info")
+            conn.send(op, **fields)
+            if probe:
+                self._proto = int(_expect(conn, "info").get("proto") or 1)
+            if bounded:
+                conn.settimeout(None)
         except BaseException:
             conn.close()
             raise
-        self._proto = proto
-        if proto >= KEEP_ALIVE_PROTO:
+        if not probe or self._proto >= KEEP_ALIVE_PROTO:
             return conn
         conn.close()
-        return self._dial(deadline)
+        return self._connect(op, deadline, **fields)
 
     @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()
+    def _lease(self, conn: Conn) -> Iterator[Conn]:
+        """Runs one RPC on conn; a terminal frame parks it and anything else drops it."""
         try:
             yield conn
         except SilkdError:
@@ -429,9 +423,8 @@ def _park(self, conn: Conn) -> None:
             conn.close()
 
     def _open_stream(self, op: str, expect: str = "ready", **fields: object) -> tuple[Conn, dict[str, Any]]:
-        conn = self._connect()
+        conn = self._connect(op, **fields)
         try:
-            conn.send(op, **fields)
             frame = _expect(conn, expect)
         except SilkdError:
             self._park(conn)
@@ -487,8 +480,7 @@ def pump_out() -> None:
                 local.close()
 
     def _call(self, op: str, expect: str, **fields: object) -> dict[str, Any]:
-        with self._lease() as conn:
-            conn.send(op, **fields)
+        with self._lease(self._connect(op, **fields)) as conn:
             return _expect(conn, expect)
 
     def _done_rpc(self, op: str, **fields: object) -> None:
@@ -502,19 +494,17 @@ def _drain_proc(
         on_stderr: Callable[[bytes], object] | None,
         trailing_done: bool = False,
     ) -> int | None:
-        with self._lease() as conn:
-            conn.send(op, pid=pid)
+        with self._lease(self._connect(op, pid=pid)) as conn:
             code = _pump_stdio(conn, on_stdout, on_stderr)
             if trailing_done and code is not None:
                 _expect(conn, "done")
             return code
 
     def _kill_after_cut(self, pid: int) -> None:
         deadline = time.monotonic() + _KILL_WAIT_SECONDS
-        with contextlib.suppress(SandboxError, OSError), self._lease(self._connect(deadline)) as conn:
+        with contextlib.suppress(SandboxError, OSError), self._lease(self._connect("kill", deadline, pid=pid)) as conn:
             watchdog = _arm_watchdog(conn, deadline, threading.Event())
             try:
-                conn.send("kill", pid=pid)
                 _expect(conn, "done")
             finally:
                 watchdog.cancel()

sdk/python/tests/test_keepalive.py

@@ -54,6 +54,17 @@ def test_old_daemon_dials_per_call():
     assert agent.upgrades == 4, "the proto probe plus one dial per call"
 
 
+def test_old_daemon_that_closes_after_info_gets_the_request_on_a_redial():
+    agent = FakeAgent(proto=1, slam=True)
+    sb = sandbox_at(agent.addr)
+    try:
+        for _ in range(2):
+            assert sb.stat("/")["kind"] == "dir"
+    finally:
+        agent.stop()
+    assert agent.upgrades == 3, "the proto probe plus one dial per call"
+
+
 def test_keep_alive_off_dials_per_call():
     agent = FakeAgent()
     sb = sandbox_at(agent.addr, keep_alive=0)

sdk/python/tests/test_stream.py

@@ -17,27 +17,6 @@
 TIMEOUT = 0.2
 
 
-class BlockedSendConn:
-    def __init__(self) -> None:
-        self.aborted = threading.Event()
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, *exc) -> None:
-        pass
-
-    def send(self, op: str, **fields) -> None:
-        assert self.aborted.wait(5 * TIMEOUT), "exec send was not aborted"
-        raise OSError("connection cut")
-
-    def abort(self) -> None:
-        self.aborted.set()
-
-    def close(self) -> None:
-        pass
-
-
 def legacy_sandbox(addr: str) -> Sandbox:
     return sandbox_at(addr, timeout=TIMEOUT, keep_alive=0)
 
@@ -135,13 +114,12 @@ def test_run_timeout_bounds_the_kill_it_sends(monkeypatch):
     assert kills == [{"v": 1, "op": "kill", "pid": 7}], kills
 
 
-def test_run_timeout_cuts_a_blocked_exec_send(monkeypatch):
-    sb = legacy_sandbox("127.0.0.1:1")
-    conn = BlockedSendConn()
-    monkeypatch.setattr(sb, "_dial", lambda deadline=None: conn)
+def test_run_timeout_cuts_a_blocked_exec_send(black_hole):
+    sb = sandbox_at(black_hole, timeout=30, keep_alive=0)
+    started = time.monotonic()
     with pytest.raises(TimeoutError):
-        sb.run(["echo", "hello"], timeout=TIMEOUT)
-    assert conn.aborted.is_set()
+        sb.run(["echo", "x" * (7 << 20)], timeout=TIMEOUT)
+    assert time.monotonic() - started < 5 * TIMEOUT
 
 
 def test_run_rejects_a_non_positive_timeout():

sdk/python/tests/test_upgrade.py

@@ -21,10 +21,11 @@
 
 
 class EarlyFrameAgent:
-    """Answers the 101 only after a frame has arrived behind the request, then serves frames until the peer leaves."""
+    """Answers the 101 only once its frames arrive behind the request, then serves frames until the peer leaves."""
 
-    def __init__(self, reply: bytes = b"HTTP/1.1 101 Switching Protocols\r\n\r\n") -> None:
+    def __init__(self, reply: bytes = b"HTTP/1.1 101 Switching Protocols\r\n\r\n", frames: int = 1) -> None:
         self.reply = reply
+        self.frames = frames
         self.early: list[str] = []
         self._server = socket.create_server(("127.0.0.1", 0))
         self.addr = f"127.0.0.1:{self._server.getsockname()[1]}"
@@ -40,19 +41,19 @@ def _serve(self) -> None:
             while reader.readline() not in (b"\r\n", b""):
                 pass
             conn.settimeout(1)
-            op = json.loads(reader.readline())["op"]
-            self.early.append(op)
+            for _ in range(self.frames):
+                self.early.append(json.loads(reader.readline())["op"])
             conn.settimeout(None)
-            conn.sendall(self.reply)
-            while op:
-                conn.sendall(json.dumps(REPLIES[op]).encode() + b"\n")
-                line = reader.readline()
-                op = json.loads(line)["op"] if line else ""
+            conn.sendall(self.reply + b"".join(json.dumps(REPLIES[op]).encode() + b"\n" for op in self.early))
+            for line in iter(reader.readline, b""):
+                conn.sendall(json.dumps(REPLIES[json.loads(line)["op"]]).encode() + b"\n")
 
 
-@pytest.mark.parametrize(("keep_alive", "early"), [(0, ["fs_mkdir"]), (30, ["info"])], ids=["no probe", "proto probe"])
-def test_first_frame_rides_ahead_of_the_101(keep_alive, early):
-    agent = EarlyFrameAgent()
+@pytest.mark.parametrize(
+    ("keep_alive", "early"), [(0, ["fs_mkdir"]), (30, ["info", "fs_mkdir"])], ids=["no probe", "proto probe"]
+)
+def test_first_frames_ride_ahead_of_the_101(keep_alive, early):
+    agent = EarlyFrameAgent(frames=len(early))
     sb = sandbox_at(agent.addr, timeout=5, keep_alive=keep_alive)
     try:
         sb.mkdir("/w")