jetstream v2 in zig stream.waow.tech
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363#!/usr/bin/env python3"""Offline process receipt for Jetstream V2 listener/client shutdown semantics."""
from __future__ import annotations
import atexitimport base64import hashlibimport http.serverimport osimport signalimport socketimport structimport subprocessimport tempfileimport threadingimport timeimport urllib.request
BIN = os.environ.get("STREAM_BIN", "./zig-out/bin/stream")
def unused_port() -> int: sock = socket.socket() sock.bind(("127.0.0.1", 0)) port = sock.getsockname()[1] sock.close() return port
def terminate(process: subprocess.Popen[str]) -> None: if process.poll() is not None: return process.kill() process.wait(timeout=3)
def start( root: str, port: int, *timeouts: str, upstream_port: int = 17997) -> subprocess.Popen[str]: process = subprocess.Popen( [ BIN, "serve", f"--addr=127.0.0.1:{port}", f"--data-dir={root}", f"--upstream=ws://127.0.0.1:{upstream_port}", f"--relay-http=http://127.0.0.1:{upstream_port}", f"--plc-url=http://127.0.0.1:{upstream_port}", "--compaction-interval=0", "--retry-interval=0", "--no-verify", *timeouts, ], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True, ) atexit.register(terminate, process) return process
class BootstrapFixture: def __init__(self) -> None: self.capture_connected = threading.Event() self.backfill_started = threading.Event() self.stop = threading.Event() self.requests: list[str] = [] fixture = self
class Handler(http.server.BaseHTTPRequestHandler): protocol_version = "HTTP/1.1"
def do_GET(self) -> None: fixture.requests.append(self.path) if self.path.startswith("/xrpc/com.atproto.sync.subscribeRepos"): key = self.headers["Sec-WebSocket-Key"] assert key is not None accept = base64.b64encode( hashlib.sha1( (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode() ).digest() ).decode() self.send_response(101) self.send_header("Upgrade", "websocket") self.send_header("Connection", "Upgrade") self.send_header("Sec-WebSocket-Accept", accept) self.end_headers() fixture.capture_connected.set() self.connection.settimeout(0.1) while not fixture.stop.is_set(): try: if not self.connection.recv(4096): break except TimeoutError: continue except OSError: break return if self.path.startswith("/xrpc/com.atproto.sync.listRepos"): fixture.backfill_started.set() while not fixture.stop.wait(0.1): pass return self.send_error(404)
def log_message(self, *_: object) -> None: pass
self.server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) self.server.daemon_threads = True self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) self.thread.start()
@property def port(self) -> int: return self.server.server_address[1]
def close(self) -> None: self.stop.set() self.server.shutdown() self.server.server_close() self.thread.join(timeout=2)
def wait_ready(process: subprocess.Popen[str], port: int, serving: bool = True) -> None: """Wait past the startup gates (listener-first startup: the listener answers from bind time, so a 200 alone no longer implies readiness). serving=True waits for "state live" (websocket upgrades accepted); serving=False only waits out storage init — a resumed/running bootstrap keeps serving false for its whole lifetime, and the bootstrap case needs exactly that state.""" deadline = time.monotonic() + 10 while time.monotonic() < deadline: if process.poll() is not None: stderr = process.stderr.read() if process.stderr else "" raise AssertionError( f"listener {port} exited with {process.returncode}:\n{stderr}" ) try: with urllib.request.urlopen(f"http://127.0.0.1:{port}/status", timeout=0.2) as response: body = response.read() if response.status == 200: if serving and b"state live" in body: return if not serving and b"state starting" not in body: return time.sleep(0.025) except OSError: time.sleep(0.025) raise AssertionError(f"listener {port} did not start")
def read_exact(sock: socket.socket, length: int) -> bytes: output = bytearray() while len(output) < length: chunk = sock.recv(length - len(output)) if not chunk: raise AssertionError("socket closed before complete frame") output.extend(chunk) return bytes(output)
def open_websocket(port: int) -> socket.socket: sock = socket.create_connection(("127.0.0.1", port), timeout=2) key = base64.b64encode(os.urandom(16)).decode() sock.sendall( ( "GET /subscribe HTTP/1.1\r\n" f"Host: 127.0.0.1:{port}\r\n" "Upgrade: websocket\r\n" "Connection: Upgrade\r\n" f"Sec-WebSocket-Key: {key}\r\n" "Sec-WebSocket-Version: 13\r\n\r\n" ).encode() ) response = bytearray() while b"\r\n\r\n" not in response: response.extend(read_exact(sock, 1)) if len(response) > 16 * 1024: raise AssertionError("oversized websocket handshake response") assert response.startswith(b"HTTP/1.1 101 Switching Protocols\r\n"), response return sock
def read_close(sock: socket.socket) -> tuple[int, bytes]: sock.settimeout(2) first, second = read_exact(sock, 2) assert first == 0x88, (first, second) assert second & 0x80 == 0, "server frames must not be masked" length = second & 0x7F if length == 126: length = struct.unpack("!H", read_exact(sock, 2))[0] elif length == 127: length = struct.unpack("!Q", read_exact(sock, 8))[0] payload = read_exact(sock, length) assert len(payload) >= 2 return struct.unpack("!H", payload[:2])[0], payload[2:]
def read_close_or_eof(sock: socket.socket) -> tuple[int, bytes] | None: """An already-expired upstream drain deadline makes close best-effort.""" sock.settimeout(2) first = sock.recv(1) if not first: return None second = read_exact(sock, 1) assert first[0] == 0x88, (first, second) assert second[0] & 0x80 == 0, "server frames must not be masked" length = second[0] & 0x7F if length == 126: length = struct.unpack("!H", read_exact(sock, 2))[0] elif length == 127: length = struct.unpack("!Q", read_exact(sock, 8))[0] payload = read_exact(sock, length) assert len(payload) >= 2 return struct.unpack("!H", payload[:2])[0], payload[2:]
def echo_close(sock: socket.socket, code: int) -> None: payload = struct.pack("!H", code) mask = os.urandom(4) masked = bytes(value ^ mask[index % 4] for index, value in enumerate(payload)) sock.sendall(bytes((0x88, 0x80 | len(payload))) + mask + masked)
def wait_clean(process: subprocess.Popen[str], timeout: float) -> None: process.wait(timeout=timeout) assert process.returncode == 0, process.stderr.read() if process.stderr else ""
def cooperative_case(root: str) -> float: port = unused_port() process = start( root, port, "--shutdown-timeout=2s", "--client-drain-timeout=5s", ) wait_ready(process, port) clients = [open_websocket(port) for _ in range(2)] started = time.monotonic() process.send_signal(signal.SIGTERM) try: for client in clients: code, reason = read_close(client) assert (code, reason) == (1001, b"server shutting down") echo_close(client, code) wait_clean(process, 3) finally: for client in clients: client.close() elapsed = time.monotonic() - started assert elapsed < 1.5, f"cooperative peers consumed their 5s budget: {elapsed:.3f}s" return elapsed
def silent_case(root: str) -> float: port = unused_port() process = start( root, port, "--shutdown-timeout=2s", "--client-drain-timeout=200ms", ) wait_ready(process, port) clients = [open_websocket(port) for _ in range(8)] started = time.monotonic() process.send_signal(signal.SIGTERM) try: for client in clients: code, reason = read_close(client) assert (code, reason) == (1001, b"server shutting down") wait_clean(process, 3) finally: for client in clients: client.close() elapsed = time.monotonic() - started assert elapsed >= 0.14, f"silent peers bypassed client drain: {elapsed:.3f}s" assert elapsed < 1.5, f"silent peers exceeded one shared 200ms budget: {elapsed:.3f}s" return elapsed
def nonpositive_case(root: str) -> float: port = unused_port() process = start( root, port, "--shutdown-timeout=0", "--client-drain-timeout=-1s", ) wait_ready(process, port) client = open_websocket(port) started = time.monotonic() process.send_signal(signal.SIGTERM) try: close = read_close_or_eof(client) if close is not None: assert close == (1001, b"server shutting down") wait_clean(process, 3) finally: client.close() elapsed = time.monotonic() - started assert elapsed < 1, f"non-positive client drain was not immediate: {elapsed:.3f}s" return elapsed
def bootstrap_case(root: str) -> tuple[float, float]: fixture = BootstrapFixture() try: elapsed: list[float] = [] for attempt in range(2): port = unused_port() process = start( root, port, "--shutdown-timeout=2s", "--client-drain-timeout=200ms", "--backfill", "--backfill-workers=1", "--backfill-async-flush-workers=0", upstream_port=fixture.port, ) wait_ready(process, port, serving=False) assert fixture.capture_connected.wait(3), "bootstrap capture did not connect" assert fixture.backfill_started.wait(3), ( f"bootstrap listRepos did not start; requests={fixture.requests!r}" ) started = time.monotonic() process.send_signal(signal.SIGTERM) wait_clean(process, 3) duration = time.monotonic() - started assert duration < 1.5, ( f"bootstrap attempt {attempt + 1} waited for the hanging " f"upstream request: {duration:.3f}s" ) elapsed.append(duration) fixture.capture_connected.clear() fixture.backfill_started.clear() return elapsed[0], elapsed[1] finally: fixture.close()
def main() -> None: with tempfile.TemporaryDirectory() as root: cooperative = cooperative_case(f"{root}/cooperative") silent = silent_case(f"{root}/silent") nonpositive = nonpositive_case(f"{root}/nonpositive") bootstrap, bootstrap_restart = bootstrap_case(f"{root}/bootstrap") print( "shutdown contract: " f"cooperative={cooperative:.3f}s silent={silent:.3f}s " f"nonpositive={nonpositive:.3f}s positive-close=1001 " f"bootstrap={bootstrap:.3f}s bootstrap-restart={bootstrap_restart:.3f}s " "nonpositive-close=best-effort shared-deadline=200ms durable-restart=ok" )
if __name__ == "__main__": main()