| 1 | from __future__ import annotations
|
| 2 |
|
| 3 | import socket
|
| 4 | import sys
|
| 5 | import threading
|
| 6 | import time
|
| 7 | import unittest
|
| 8 | from pathlib import Path
|
| 9 |
|
| 10 | sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts"))
|
| 11 |
|
| 12 | from oakbench.byteproxy import CountingProxy
|
| 13 |
|
| 14 |
|
| 15 | class EchoServer:
|
| 16 | def __init__(self) -> None:
|
| 17 | self._stop = threading.Event()
|
| 18 | self._listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
| 19 | self._listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
| 20 | self._listener.bind(("127.0.0.1", 0))
|
| 21 | self._listener.listen()
|
| 22 | self._listener.settimeout(0.2)
|
| 23 | self.address = self._listener.getsockname()[:2]
|
| 24 | self._thread = threading.Thread(target=self._serve, daemon=True)
|
| 25 |
|
| 26 | def start(self) -> "EchoServer":
|
| 27 | self._thread.start()
|
| 28 | return self
|
| 29 |
|
| 30 | def stop(self) -> None:
|
| 31 | self._stop.set()
|
| 32 | self._listener.close()
|
| 33 | self._thread.join(timeout=1.0)
|
| 34 |
|
| 35 | def __enter__(self) -> "EchoServer":
|
| 36 | return self.start()
|
| 37 |
|
| 38 | def __exit__(self, exc_type, exc, tb) -> None: # type: ignore[no-untyped-def]
|
| 39 | self.stop()
|
| 40 |
|
| 41 | def _serve(self) -> None:
|
| 42 | while not self._stop.is_set():
|
| 43 | try:
|
| 44 | conn, _addr = self._listener.accept()
|
| 45 | except socket.timeout:
|
| 46 | continue
|
| 47 | except OSError:
|
| 48 | break
|
| 49 | threading.Thread(target=self._handle, args=(conn,), daemon=True).start()
|
| 50 |
|
| 51 | def _handle(self, conn: socket.socket) -> None:
|
| 52 | with conn:
|
| 53 | chunks = []
|
| 54 | while True:
|
| 55 | chunk = conn.recv(4096)
|
| 56 | if not chunk:
|
| 57 | break
|
| 58 | chunks.append(chunk)
|
| 59 | conn.sendall(b"reply:" + b"".join(chunks))
|
| 60 |
|
| 61 |
|
| 62 | def roundtrip(address: tuple[str, int], payload: bytes) -> bytes:
|
| 63 | with socket.create_connection(address, timeout=2.0) as sock:
|
| 64 | sock.sendall(payload)
|
| 65 | sock.shutdown(socket.SHUT_WR)
|
| 66 | chunks = []
|
| 67 | while True:
|
| 68 | chunk = sock.recv(4096)
|
| 69 | if not chunk:
|
| 70 | break
|
| 71 | chunks.append(chunk)
|
| 72 | return b"".join(chunks)
|
| 73 |
|
| 74 |
|
| 75 | def wait_for(predicate, timeout_s: float = 2.0) -> None: # type: ignore[no-untyped-def]
|
| 76 | deadline = time.monotonic() + timeout_s
|
| 77 | while time.monotonic() < deadline:
|
| 78 | if predicate():
|
| 79 | return
|
| 80 | time.sleep(0.01)
|
| 81 | raise AssertionError("condition did not become true before timeout")
|
| 82 |
|
| 83 |
|
| 84 | class CountingProxyTest(unittest.TestCase):
|
| 85 | def test_proxy_forwards_and_counts_bytes_without_privileged_port(self) -> None:
|
| 86 | payload = b"hello through proxy"
|
| 87 | with EchoServer() as server:
|
| 88 | host, port = server.address
|
| 89 | with CountingProxy(host, port) as proxy:
|
| 90 | self.assertGreater(proxy.address[1], 0)
|
| 91 | self.assertEqual(roundtrip(proxy.address, payload), b"reply:" + payload)
|
| 92 | wait_for(lambda: proxy.snapshot().upstream_to_client_bytes == len(b"reply:" + payload))
|
| 93 | stats = proxy.snapshot()
|
| 94 |
|
| 95 | self.assertEqual(stats.connections, 1)
|
| 96 | self.assertEqual(stats.connection_errors, 0)
|
| 97 | self.assertEqual(stats.client_to_upstream_bytes, len(payload))
|
| 98 | self.assertEqual(stats.upstream_to_client_bytes, len(b"reply:" + payload))
|
| 99 | self.assertEqual(stats.total_bytes, len(payload) + len(b"reply:" + payload))
|
| 100 | self.assertIsNotNone(stats.first_upstream_to_client_ms)
|
| 101 | self.assertGreaterEqual(stats.first_upstream_to_client_ms, 0.0)
|
| 102 | self.assertEqual(stats.first_payload_bytes_server_to_client, len(b"reply:" + payload))
|
| 103 |
|
| 104 | def test_reset_counts_preserves_running_proxy(self) -> None:
|
| 105 | with EchoServer() as server:
|
| 106 | host, port = server.address
|
| 107 | with CountingProxy(host, port) as proxy:
|
| 108 | self.assertEqual(roundtrip(proxy.address, b"one"), b"reply:one")
|
| 109 | wait_for(lambda: proxy.snapshot().connections == 1)
|
| 110 | proxy.reset_counts()
|
| 111 | self.assertEqual(proxy.snapshot().total_bytes, 0)
|
| 112 | self.assertIsNone(proxy.snapshot().first_upstream_to_client_ms)
|
| 113 | self.assertIsNone(proxy.snapshot().first_payload_bytes_server_to_client)
|
| 114 |
|
| 115 | self.assertEqual(roundtrip(proxy.address, b"two"), b"reply:two")
|
| 116 | wait_for(lambda: proxy.snapshot().upstream_to_client_bytes == len(b"reply:two"))
|
| 117 | stats = proxy.snapshot()
|
| 118 |
|
| 119 | self.assertEqual(stats.connections, 1)
|
| 120 | self.assertEqual(stats.client_to_upstream_bytes, len(b"two"))
|
| 121 | self.assertEqual(stats.upstream_to_client_bytes, len(b"reply:two"))
|
| 122 | self.assertIsNotNone(stats.first_upstream_to_client_ms)
|
| 123 | self.assertEqual(stats.first_payload_bytes_server_to_client, len(b"reply:two"))
|
| 124 |
|
| 125 | def test_pump_eof_never_sets_first_payload_fields(self) -> None:
|
| 126 | """An empty recv is EOF, not payload: TTFD and first-payload size must
|
| 127 | stay null (ADR-0002 unmeasured) until a NON-EMPTY upstream-to-client
|
| 128 | chunk arrives."""
|
| 129 |
|
| 130 | class FakeSock:
|
| 131 | def __init__(self, chunks: list[bytes]) -> None:
|
| 132 | self.chunks = list(chunks)
|
| 133 | self.sent: list[bytes] = []
|
| 134 |
|
| 135 | def recv(self, _size: int) -> bytes:
|
| 136 | return self.chunks.pop(0) if self.chunks else b""
|
| 137 |
|
| 138 | def sendall(self, data: bytes) -> None:
|
| 139 | self.sent.append(data)
|
| 140 |
|
| 141 | def shutdown(self, _how: int) -> None:
|
| 142 | pass
|
| 143 |
|
| 144 | proxy = CountingProxy("127.0.0.1", 9) # never started; _pump driven directly
|
| 145 |
|
| 146 | # Upstream closes without sending a byte: EOF only, nothing measured.
|
| 147 | proxy._pump(FakeSock([]), FakeSock([]), "upstream_to_client")
|
| 148 | stats = proxy.snapshot()
|
| 149 | self.assertIsNone(stats.first_upstream_to_client_ms)
|
| 150 | self.assertIsNone(stats.first_payload_bytes_server_to_client)
|
| 151 | self.assertEqual(stats.upstream_to_client_bytes, 0)
|
| 152 |
|
| 153 | # First real payload chunk sets both fields; the size is the first
|
| 154 | # chunk's, not the total.
|
| 155 | proxy._pump(FakeSock([b"abc", b"defg"]), FakeSock([]), "upstream_to_client")
|
| 156 | stats = proxy.snapshot()
|
| 157 | self.assertIsNotNone(stats.first_upstream_to_client_ms)
|
| 158 | self.assertEqual(stats.first_payload_bytes_server_to_client, 3)
|
| 159 | self.assertEqual(stats.upstream_to_client_bytes, 7)
|
| 160 |
|
| 161 | # client_to_upstream traffic never touches the first-payload fields.
|
| 162 | fresh = CountingProxy("127.0.0.1", 9)
|
| 163 | fresh._pump(FakeSock([b"request"]), FakeSock([]), "client_to_upstream")
|
| 164 | stats = fresh.snapshot()
|
| 165 | self.assertIsNone(stats.first_upstream_to_client_ms)
|
| 166 | self.assertIsNone(stats.first_payload_bytes_server_to_client)
|
| 167 | self.assertEqual(stats.client_to_upstream_bytes, 7)
|
| 168 |
|
| 169 | def test_reset_counts_clears_stale_ttfd_between_measured_runs(self) -> None:
|
| 170 | """timed_row/timed_sequence_row reuse one proxy across runs and call
|
| 171 | reset_counts before each measured operation; run 2 must never report
|
| 172 | run 1's TTFD."""
|
| 173 |
|
| 174 | class FakeSock:
|
| 175 | def __init__(self, chunks: list[bytes]) -> None:
|
| 176 | self.chunks = list(chunks)
|
| 177 |
|
| 178 | def recv(self, _size: int) -> bytes:
|
| 179 | return self.chunks.pop(0) if self.chunks else b""
|
| 180 |
|
| 181 | def sendall(self, data: bytes) -> None:
|
| 182 | pass
|
| 183 |
|
| 184 | def shutdown(self, _how: int) -> None:
|
| 185 | pass
|
| 186 |
|
| 187 | proxy = CountingProxy("127.0.0.1", 9)
|
| 188 | proxy._pump(FakeSock([b"run-one"]), FakeSock([]), "upstream_to_client")
|
| 189 | self.assertIsNotNone(proxy.snapshot().first_upstream_to_client_ms)
|
| 190 |
|
| 191 | proxy.reset_counts()
|
| 192 | stats = proxy.snapshot()
|
| 193 | self.assertIsNone(stats.first_upstream_to_client_ms)
|
| 194 | self.assertIsNone(stats.first_payload_bytes_server_to_client)
|
| 195 |
|
| 196 | # A run whose upstream sends nothing stays null after the reset.
|
| 197 | proxy._pump(FakeSock([]), FakeSock([]), "upstream_to_client")
|
| 198 | self.assertIsNone(proxy.snapshot().first_upstream_to_client_ms)
|
| 199 |
|
| 200 | def test_connection_errors_are_counted(self) -> None:
|
| 201 | closed = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
| 202 | closed.bind(("127.0.0.1", 0))
|
| 203 | unused_port = closed.getsockname()[1]
|
| 204 | closed.close()
|
| 205 |
|
| 206 | with CountingProxy("127.0.0.1", unused_port, connect_timeout_s=0.1) as proxy:
|
| 207 | with socket.create_connection(proxy.address, timeout=2.0) as sock:
|
| 208 | sock.sendall(b"x")
|
| 209 | wait_for(lambda: proxy.snapshot().connection_errors == 1)
|
| 210 | stats = proxy.snapshot()
|
| 211 |
|
| 212 | self.assertEqual(stats.connections, 0)
|
| 213 | self.assertEqual(stats.connection_errors, 1)
|
| 214 |
|
| 215 |
|
| 216 | if __name__ == "__main__":
|
| 217 | unittest.main()
|