feat: pre-warmed TCP connection pool to first hop
Add FirstHopPool that maintains a deque of pre-established TCP connections to chain[0]. Connections idle beyond pool_max_idle are evicted; a background task refills to pool_size. build_chain() tries the pool first, falls back to open_connection. Enabled with pool_size > 0 in config. Only pools the TCP handshake -- SOCKS/HTTP tunnels are consumed, not returned.
This commit is contained in:
105
tests/test_connpool.py
Normal file
105
tests/test_connpool.py
Normal file
@@ -0,0 +1,105 @@
|
||||
"""Tests for the first-hop connection pool."""
|
||||
|
||||
import asyncio
|
||||
|
||||
from s5p.config import ChainHop
|
||||
from s5p.connpool import FirstHopPool
|
||||
|
||||
|
||||
async def _echo_server(host="127.0.0.1", port=0):
|
||||
"""Start a TCP server that accepts connections and holds them open."""
|
||||
|
||||
async def handler(reader, writer):
|
||||
try:
|
||||
await reader.read(1) # block until client disconnects
|
||||
except (ConnectionError, asyncio.CancelledError):
|
||||
pass
|
||||
finally:
|
||||
writer.close()
|
||||
|
||||
server = await asyncio.start_server(handler, host, port)
|
||||
port = server.sockets[0].getsockname()[1]
|
||||
return server, port
|
||||
|
||||
|
||||
class TestFirstHopPool:
|
||||
"""Test connection pool lifecycle."""
|
||||
|
||||
def test_acquire_returns_connection(self):
|
||||
async def run():
|
||||
server, port = await _echo_server()
|
||||
async with server:
|
||||
hop = ChainHop(proto="socks5", host="127.0.0.1", port=port)
|
||||
pool = FirstHopPool(hop, size=2, max_idle=10.0)
|
||||
await pool.start()
|
||||
try:
|
||||
conn = await pool.acquire()
|
||||
assert conn is not None
|
||||
reader, writer = conn
|
||||
assert not writer.is_closing()
|
||||
writer.close()
|
||||
finally:
|
||||
await pool.stop()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_acquire_exhausts_pool(self):
|
||||
async def run():
|
||||
server, port = await _echo_server()
|
||||
async with server:
|
||||
hop = ChainHop(proto="socks5", host="127.0.0.1", port=port)
|
||||
pool = FirstHopPool(hop, size=2, max_idle=10.0)
|
||||
await pool.start()
|
||||
try:
|
||||
c1 = await pool.acquire()
|
||||
c2 = await pool.acquire()
|
||||
c3 = await pool.acquire()
|
||||
assert c1 is not None
|
||||
assert c2 is not None
|
||||
assert c3 is None # pool exhausted
|
||||
c1[1].close()
|
||||
c2[1].close()
|
||||
finally:
|
||||
await pool.stop()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_stale_connections_evicted(self):
|
||||
async def run():
|
||||
server, port = await _echo_server()
|
||||
async with server:
|
||||
hop = ChainHop(proto="socks5", host="127.0.0.1", port=port)
|
||||
pool = FirstHopPool(hop, size=2, max_idle=0.05)
|
||||
await pool._fill() # pre-warm without starting refill loop
|
||||
assert len(pool._pool) == 2
|
||||
await asyncio.sleep(0.1) # let connections go stale
|
||||
conn = await pool.acquire()
|
||||
assert conn is None # all stale, evicted
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_stop_drains_pool(self):
|
||||
async def run():
|
||||
server, port = await _echo_server()
|
||||
async with server:
|
||||
hop = ChainHop(proto="socks5", host="127.0.0.1", port=port)
|
||||
pool = FirstHopPool(hop, size=4, max_idle=30.0)
|
||||
await pool.start()
|
||||
await pool.stop()
|
||||
# Pool should be empty after stop
|
||||
conn = await pool.acquire()
|
||||
assert conn is None
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_unreachable_hop_graceful(self):
|
||||
"""Pool creation with an unreachable hop should not raise."""
|
||||
async def run():
|
||||
hop = ChainHop(proto="socks5", host="127.0.0.1", port=1) # nothing listening
|
||||
pool = FirstHopPool(hop, size=2, max_idle=10.0)
|
||||
await pool.start() # should not raise, just log warnings
|
||||
conn = await pool.acquire()
|
||||
assert conn is None # no connections could be established
|
||||
await pool.stop()
|
||||
|
||||
asyncio.run(run())
|
||||
Reference in New Issue
Block a user