diff --git a/tools/station/server.py b/tools/station/server.py index e799492..59d05c9 100755 --- a/tools/station/server.py +++ b/tools/station/server.py @@ -17,6 +17,7 @@ import select import socket import sys import tempfile +import threading import time from datetime import datetime, timezone from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -59,29 +60,36 @@ class Station: self.since = now_iso() self.served = 0 self.recent: list[dict] = [] + # ThreadingHTTPServer runs one handler thread per connection, and several lamps + # can be downloading at once. Every mutation of served/recent and the snapshot + # that gets serialized to disk must happen under this same lock, or concurrent + # requests lose increments and stomp on each other's status.json writes. + self.lock = threading.RLock() self.write_state() def image_path(self, target: str) -> Path: return self.image_dir / f"ninlight_{target}_ota.img" def record(self, ip: str, target: str, version: str, nbytes: int, status: str) -> None: - if status == "ok": - self.served += 1 - self.recent.append({"ts": now_iso(), "ip": ip, "target": target, - "version": version, "bytes": nbytes, "status": status}) - del self.recent[:-RECENT_MAX] - self.write_state() + with self.lock: + if status == "ok": + self.served += 1 + self.recent.append({"ts": now_iso(), "ip": ip, "target": target, + "version": version, "bytes": nbytes, "status": status}) + del self.recent[:-RECENT_MAX] + self.write_state() def write_state(self) -> None: - payload = {"since": self.since, "served": self.served, "recent": self.recent} - try: - self.state_path.parent.mkdir(parents=True, exist_ok=True) - fd, tmp = tempfile.mkstemp(dir=self.state_path.parent, prefix=".status-") - with os.fdopen(fd, "w") as f: - json.dump(payload, f, indent=2) - os.replace(tmp, self.state_path) - except OSError as e: - print(f"warning: cannot write {self.state_path}: {e}", flush=True) + with self.lock: + payload = {"since": self.since, "served": self.served, "recent": list(self.recent)} + try: + self.state_path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(dir=self.state_path.parent, prefix=".status-") + with os.fdopen(fd, "w") as f: + json.dump(payload, f, indent=2) + os.replace(tmp, self.state_path) + except OSError as e: + print(f"warning: cannot write {self.state_path}: {e}", flush=True) class Handler(BaseHTTPRequestHandler): @@ -150,7 +158,6 @@ class Handler(BaseHTTPRequestHandler): sent = 0 status = "ok" - recorded = False remaining = size - start try: with path.open("rb") as f: @@ -180,20 +187,15 @@ class Handler(BaseHTTPRequestHandler): chunk = f.read(read_size) if not chunk: break - if sent + len(chunk) >= remaining: - # Record success before the last bytes go out: the lamp (and our - # tests) may observe the response completing and immediately check - # status.json, so the state must already be on disk by then. - self.station.record(self.client_address[0], target, version, - sent + len(chunk), "ok") - recorded = True self.wfile.write(chunk) sent += len(chunk) except (BrokenPipeError, ConnectionResetError, TimeoutError): status = "aborted" + # Record only now that the transfer has actually finished (or failed): "ok" must + # never be written before the last byte is confirmed sent, or status.json can + # claim success for a device that never received the end of the image. self.log(206 if start else 200, sent, started) - if not recorded: - self.station.record(self.client_address[0], target, version, sent, status) + self.station.record(self.client_address[0], target, version, sent, status) def main(argv: list[str] | None = None) -> int: diff --git a/tools/station/tests/test_server.py b/tools/station/tests/test_server.py index 052fca4..a027273 100644 --- a/tools/station/tests/test_server.py +++ b/tools/station/tests/test_server.py @@ -4,6 +4,7 @@ import json import socket import subprocess import sys +import threading import time from pathlib import Path @@ -11,6 +12,7 @@ import pytest SERVER = Path(__file__).resolve().parents[1] / "server.py" TRAILER_SIZE = 512 +RECENT_MAX = 8 def make_image(path: Path, device: str, version: str, payload: bytes) -> bytes: @@ -60,15 +62,29 @@ class Station: def conn(self): return http.client.HTTPConnection("127.0.0.1", self.port, timeout=5) - def read_state(self): - for _ in range(50): + def read_state(self, predicate=None): + """Reads status.json, optionally waiting for it to satisfy `predicate`. + + The server only writes the entry for a request after that request's response has + actually finished (or failed) on the wire - by design, so it never claims success + before the last byte is confirmed sent (see server.py). A client that has just + finished reading its own response can therefore still see a status.json that is a + few milliseconds stale. Rather than have the server promise something before it is + true, we wait here for the state to catch up. + """ + last = None + for _ in range(100): if self.state.exists(): try: - return json.loads(self.state.read_text()) + last = json.loads(self.state.read_text()) + if predicate is None or predicate(last): + return last except json.JSONDecodeError: pass time.sleep(0.05) - raise AssertionError("no state file written") + if last is None: + raise AssertionError("no state file written") + raise AssertionError(f"state never satisfied predicate, last seen: {last!r}") def stop(self): self.proc.terminate() @@ -137,7 +153,7 @@ def test_state_file_records_a_delivery(station): c = station.conn() c.request("GET", "/ninlight_evk_ota.img") c.getresponse().read() - state = station.read_state() + state = station.read_state(lambda s: s["served"] >= 1) assert state["served"] == 1 entry = state["recent"][-1] assert entry["target"] == "evk" @@ -151,7 +167,7 @@ def test_state_file_records_a_missing_image(station): c = station.conn() c.request("GET", "/ninlight_g95_cw_ota.img") c.getresponse().read() - state = station.read_state() + state = station.read_state(lambda s: len(s["recent"]) >= 1) assert state["served"] == 0 assert state["recent"][-1]["status"] == "notfound" assert state["recent"][-1]["target"] == "g95_cw" @@ -163,9 +179,9 @@ def test_recent_is_capped_at_eight(station): c = station.conn() c.request("GET", "/ninlight_evk_ota.img") c.getresponse().read() - state = station.read_state() + state = station.read_state(lambda s: s["served"] >= 12) assert state["served"] == 12 - assert len(state["recent"]) == 8 + assert len(state["recent"]) == RECENT_MAX def test_aborted_transfer_is_recorded(tmp_path): @@ -208,8 +224,61 @@ def test_stall_after_stops_sending_but_keeps_the_connection(tmp_path): got += chunk except socket.timeout: pass - # far less than the full image, and the peer never closed the connection + # far less than the full image assert len(got) < 20000 + # and the connection is genuinely still open: a server that (wrongly) closed it + # after the stall boundary would make this recv() return b"" (EOF) instead of + # timing out. Just counting received bytes above would not catch that - a server + # closing the connection right after --stall-after also sends far less than + # 20000 bytes. + sock.settimeout(1) + with pytest.raises(socket.timeout): + sock.recv(1) sock.close() finally: s.stop() + + +def test_concurrent_downloads_are_all_recorded(station): + """Several lamps hitting the station at once must not lose or corrupt status.json. + + Station.served is a shared counter and Station.recent a shared list, mutated from one + handler thread per connection under ThreadingHTTPServer. Without a lock around the + read-modify-write of `served` and the recent-list mutation + serialize-to-disk step, + concurrent requests can lose increments (two threads read the same old value) or have + one thread's status.json write clobber another's. + """ + make_image(station.image_dir / "ninlight_evk_ota.img", "evk", "1.0.0", b"c" * 5000) + expected_len = 5000 + TRAILER_SIZE + n_workers = 20 + errors = [] + results_lock = threading.Lock() + + def worker(): + try: + c = station.conn() + c.request("GET", "/ninlight_evk_ota.img") + r = c.getresponse() + body = r.read() + if r.status != 200 or len(body) != expected_len: + with results_lock: + errors.append((r.status, len(body))) + except Exception as e: # pragma: no cover - surfaced via the errors list below + with results_lock: + errors.append(e) + + threads = [threading.Thread(target=worker) for _ in range(n_workers)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + + assert not errors, f"some downloads failed or returned the wrong body: {errors}" + + state = station.read_state(lambda s: s["served"] >= n_workers) + assert state["served"] == n_workers + assert len(state["recent"]) == min(n_workers, RECENT_MAX) + for entry in state["recent"]: + assert entry["status"] == "ok" + assert entry["target"] == "evk" + assert entry["bytes"] == expected_len