"""A counting TCP proxy in front of a Postgres server: how many ROUND TRIPS a client really made.
A round trip is counted each time the client sends after the server last spoke (or first), which is
what a pg-wire client pays a full network latency for. Used by test_repository_pgwire.py; it needs a
real Postgres (ERP_TEST_PG_URL), so it is skipped everywhere else."""
import socket
import threading
class CountingProxy:
def __init__(self, host: str, port: int):
self.target = (host, port)
self.round_trips = 0
self._lock = threading.Lock()
self._srv = socket.socket()
self._srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
self._srv.bind(("127.0.0.1", 0))
self._srv.listen(16)
self.port = self._srv.getsockname()[1]
threading.Thread(target=self._accept, daemon=True).start()
def reset(self) -> None:
with self._lock:
self.round_trips = 0
def _accept(self) -> None:
while True:
try:
client, _ = self._srv.accept()
except OSError:
return
server = socket.create_connection(self.target)
state = {"last": "server"}
threading.Thread(target=self._pump, args=(client, server, state, "client"), daemon=True).start()
threading.Thread(target=self._pump, args=(server, client, state, "server"), daemon=True).start()
def _pump(self, src, dst, state, side) -> None:
try:
while True:
data = src.recv(65536)
if not data:
break
with self._lock:
if side == "client" and state["last"] == "server":
self.round_trips += 1
state["last"] = side
dst.sendall(data)
except OSError:
pass
finally:
for s in (src, dst):
try:
s.close()
except OSError:
pass
def close(self) -> None:
self._srv.close()