248f5c3306
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.
106 lines
3.7 KiB
Python
106 lines
3.7 KiB
Python
"""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())
|