async def test_stats_measures(dsn):
async def worker(n):
async with p.connection() as conn:
await conn.execute("select pg_sleep(0.2)")
async with pool.AsyncConnectionPool(dsn, min_size=2, max_size=4) as p:
await p.wait(2.0)
stats = p.get_stats()
assert stats["pool_min"] == 2
assert stats["pool_max"] == 4
assert stats["pool_size"] == 2
assert stats["pool_available"] == 2
assert stats["requests_waiting"] == 0
ts = [create_task(worker(i)) for i in range(3)]
await asyncio.sleep(0.1)
stats = p.get_stats()
await asyncio.gather(*ts)
assert stats["pool_min"] == 2
assert stats["pool_max"] == 4
assert stats["pool_size"] == 3
assert stats["pool_available"] == 0
assert stats["requests_waiting"] == 0
await p.wait(2.0)
ts = [create_task(worker(i)) for i in range(7)]
await asyncio.sleep(0.1)
stats = p.get_stats()
await asyncio.gather(*ts)
assert stats["pool_min"] == 2
assert stats["pool_max"] == 4
assert stats["pool_size"] == 4
assert stats["pool_available"] == 0
assert stats["requests_waiting"] == 3
async def test_queue_size(dsn):
async def worker(t, ev=None):
try:
async with p.connection():
if ev:
ev.set()
await asyncio.sleep(t)
except pool.TooManyRequests as e:
errors.append(e)
else:
success.append(True)
errors = []
success = []
async with pool.AsyncConnectionPool(dsn, min_size=1, max_waiting=3) as p:
await p.wait()
ev = asyncio.Event()
create_task(worker(0.3, ev))
await ev.wait()
ts = [create_task(worker(0.1)) for i in range(4)]
await asyncio.gather(*ts)
assert len(success) == 4
assert len(errors) == 1
assert isinstance(errors[0], pool.TooManyRequests)
assert p.name in str(errors[0])
assert str(p.max_waiting) in str(errors[0])
assert p.get_stats()["requests_errors"] == 1
async def test_stats_usage(dsn):
async def worker(n):
try:
async with p.connection(timeout=0.3) as conn:
await conn.execute("select pg_sleep(0.2)")
except pool.PoolTimeout:
pass
async with pool.AsyncConnectionPool(dsn, min_size=3) as p:
await p.wait(2.0)
ts = [create_task(worker(i)) for i in range(7)]
await asyncio.gather(*ts)
stats = p.get_stats()
assert stats["requests_num"] == 7
assert stats["requests_queued"] == 4
assert 850 <= stats["requests_wait_ms"] <= 950
assert stats["requests_errors"] == 1
assert 1150 <= stats["usage_ms"] <= 1250
assert stats.get("returns_bad", 0) == 0
async with p.connection() as conn:
await conn.close()
await p.wait()
stats = p.pop_stats()
assert stats["requests_num"] == 8
assert stats["returns_bad"] == 1
async with p.connection():
pass
assert p.get_stats()["requests_num"] == 1
async def test_shrink(dsn, monkeypatch):
from psycopg3.pool.async_pool import ShrinkPool
results = []
async def run_hacked(self, pool):
n0 = pool._nconns
await orig_run(self, pool)
n1 = pool._nconns
results.append((n0, n1))
orig_run = ShrinkPool._run
monkeypatch.setattr(ShrinkPool, "_run", run_hacked)
async def worker(n):
async with p.connection() as conn:
await conn.execute("select pg_sleep(0.1)")
async with pool.AsyncConnectionPool(dsn,
min_size=2,
max_size=4,
max_idle=0.2) as p:
await p.wait(5.0)
assert p.max_idle == 0.2
ts = [create_task(worker(i)) for i in range(4)]
await asyncio.gather(*ts)
await asyncio.sleep(1)
assert results == [(4, 4), (4, 3), (3, 2), (2, 2), (2, 2)]
async def test_grow(dsn, monkeypatch, retries):
delay_connection(monkeypatch, 0.1)
async def worker(n):
t0 = time()
async with p.connection() as conn:
await conn.execute("select 1 from pg_sleep(0.2)")
t1 = time()
results.append((n, t1 - t0))
async for retry in retries:
with retry:
async with pool.AsyncConnectionPool(dsn,
min_size=2,
max_size=4,
num_workers=3) as p:
await p.wait(1.0)
ts = []
results = []
ts = [create_task(worker(i)) for i in range(6)]
await asyncio.gather(*ts)
want_times = [0.2, 0.2, 0.3, 0.4, 0.4, 0.4]
times = [item[1] for item in results]
for got, want in zip(times, want_times):
assert got == pytest.approx(want, 0.1), times
async def test_queue_timeout_override(dsn):
async def worker(n):
t0 = time()
timeout = 0.25 if n == 3 else None
try:
async with p.connection(timeout=timeout) as conn:
cur = await conn.execute(
"select pg_backend_pid() from pg_sleep(0.2)")
(pid, ) = await cur.fetchone()
except pool.PoolTimeout as e:
t1 = time()
errors.append((n, t1 - t0, e))
else:
t1 = time()
results.append((n, t1 - t0, pid))
results = []
errors = []
async with pool.AsyncConnectionPool(dsn, min_size=2, timeout=0.1) as p:
ts = [create_task(worker(i)) for i in range(4)]
await asyncio.gather(*ts)
assert len(results) == 3
assert len(errors) == 1
for e in errors:
assert 0.1 < e[1] < 0.15
async def test_empty_queue_timeout():
s = AsyncScheduler()
t0 = time()
times = []
wait_orig = s._event.wait
async def wait_logging():
try:
rv = await wait_orig()
finally:
times.append(time() - t0)
return rv
s._event.wait = wait_logging
s.EMPTY_QUEUE_TIMEOUT = 0.2
t = create_task(s.run())
await asyncio.sleep(0.5)
await s.enter(0.5, None)
await asyncio.gather(t)
times.append(time() - t0)
for got, want in zip(times, [0.2, 0.4, 0.5, 1.0]):
assert got == pytest.approx(want, 0.1), times
async def test_sched_error(caplog):
caplog.set_level(logging.WARNING, logger="psycopg3")
s = AsyncScheduler()
t = create_task(s.run())
results = []
async def worker(i):
results.append((i, time()))
async def error():
1 / 0
t0 = time()
await s.enter(0.1, partial(worker, 1))
await s.enter(0.4, None)
await s.enter(0.3, partial(worker, 2))
await s.enter(0.2, error)
await asyncio.gather(t)
t1 = time()
assert t1 - t0 == pytest.approx(0.4, 0.1)
assert len(results) == 2
assert results[0][0] == 1
assert results[0][1] - t0 == pytest.approx(0.1, 0.1)
assert results[1][0] == 2
assert results[1][1] - t0 == pytest.approx(0.3, 0.1)
assert len(caplog.records) == 1
assert "ZeroDivisionError" in caplog.records[0].message
async def test_first_task_rescheduling():
s = AsyncScheduler()
t0 = time()
times = []
wait_orig = s._event.wait
async def wait_logging():
try:
rv = await wait_orig()
finally:
times.append(time() - t0)
return rv
s._event.wait = wait_logging
s.EMPTY_QUEUE_TIMEOUT = 0.1
async def noop():
pass
await s.enter(0.4, noop)
t = create_task(s.run())
await s.enter(0.6, None) # this task doesn't trigger a reschedule
await asyncio.sleep(0.1)
await s.enter(0.1, noop) # this triggers a reschedule
await asyncio.gather(t)
times.append(time() - t0)
for got, want in zip(times, [0.1, 0.2, 0.4, 0.6, 0.6]):
assert got == pytest.approx(want, 0.1), times
async def test_identify_closure(aconn, dsn):
conn2 = await psycopg3.AsyncConnection.connect(dsn)
async def closer():
await asyncio.sleep(0.3)
await conn2.execute("select pg_terminate_backend(%s)",
[aconn.pgconn.backend_pid])
t0 = time.time()
ev = asyncio.Event()
loop = asyncio.get_event_loop()
loop.add_reader(aconn.fileno(), ev.set)
create_task(closer())
await asyncio.wait_for(ev.wait(), 1.0)
with pytest.raises(psycopg3.OperationalError):
await aconn.execute("select 1")
t1 = time.time()
assert 0.3 < t1 - t0 < 0.5
async def test_closed_queue(dsn):
p = pool.AsyncConnectionPool(dsn, min_size=1)
success = []
async def w1():
async with p.connection() as conn:
res = await conn.execute("select 1 from pg_sleep(0.2)")
assert await res.fetchone() == (1, )
success.append("w1")
async def w2():
with pytest.raises(pool.PoolClosed):
async with p.connection():
pass
success.append("w2")
t1 = create_task(w1())
await asyncio.sleep(0.1)
t2 = create_task(w2())
await p.close()
await asyncio.gather(t1, t2)
assert len(success) == 2
async def test_resize(dsn):
async def sampler():
await asyncio.sleep(0.05) # ensure sampling happens after shrink check
while True:
await asyncio.sleep(0.2)
if p.closed:
break
size.append(len(p._pool))
async def client(t):
async with p.connection() as conn:
await conn.execute("select pg_sleep(%s)", [t])
size = []
async with pool.AsyncConnectionPool(dsn, min_size=2, max_idle=0.2) as p:
s = create_task(sampler())
await asyncio.sleep(0.3)
c = create_task(client(0.4))
await asyncio.sleep(0.2)
await p.resize(4)
assert p.min_size == 4
assert p.max_size == 4
await asyncio.sleep(0.4)
await p.resize(2)
assert p.min_size == 2
assert p.max_size == 2
await asyncio.sleep(0.6)
await asyncio.gather(s, c)
assert size == [2, 1, 3, 4, 3, 2, 2]
async def test_spike(dsn, monkeypatch):
# Inspired to https://github.com/brettwooldridge/HikariCP/blob/dev/
# documents/Welcome-To-The-Jungle.md
delay_connection(monkeypatch, 0.15)
async def worker():
async with p.connection():
await asyncio.sleep(0.002)
async with pool.AsyncConnectionPool(dsn, min_size=5, max_size=10) as p:
await p.wait()
ts = [create_task(worker()) for i in range(50)]
await asyncio.gather(*ts)
await p.wait()
assert len(p._pool) < 7
async def test_dead_client(dsn):
async def worker(i, timeout):
try:
async with p.connection(timeout=timeout) as conn:
await conn.execute("select pg_sleep(0.3)")
results.append(i)
except pool.PoolTimeout:
if timeout > 0.2:
raise
async with pool.AsyncConnectionPool(dsn, min_size=2) as p:
results = []
ts = [
create_task(worker(i, timeout))
for i, timeout in enumerate([0.4, 0.4, 0.1, 0.4, 0.4])
]
await asyncio.gather(*ts)
await asyncio.sleep(0.2)
assert set(results) == set([0, 1, 3, 4])
assert len(p._pool) == 2 # no connection was lost
async def test_queue(dsn, retries):
async def worker(n):
t0 = time()
async with p.connection() as conn:
cur = await conn.execute(
"select pg_backend_pid() from pg_sleep(0.2)")
(pid, ) = await cur.fetchone()
t1 = time()
results.append((n, t1 - t0, pid))
async for retry in retries:
with retry:
results = []
async with pool.AsyncConnectionPool(dsn, min_size=2) as p:
ts = [create_task(worker(i)) for i in range(6)]
await asyncio.gather(*ts)
times = [item[1] for item in results]
want_times = [0.2, 0.2, 0.4, 0.4, 0.6, 0.6]
for got, want in zip(times, want_times):
assert got == pytest.approx(want, 0.1), times
assert len(set(r[2] for r in results)) == 2, results
async def test_sched_task():
s = AsyncScheduler()
t = create_task(s.run())
results = []
async def worker(i):
results.append((i, time()))
t0 = time()
await s.enter(0.1, partial(worker, 1))
await s.enter(0.4, partial(worker, 3))
await s.enter(0.3, None)
await s.enter(0.2, partial(worker, 2))
await asyncio.gather(t)
t1 = time()
assert t1 - t0 == pytest.approx(0.3, 0.1)
assert len(results) == 2
assert results[0][0] == 1
assert results[0][1] - t0 == pytest.approx(0.1, 0.1)
assert results[1][0] == 2
assert results[1][1] - t0 == pytest.approx(0.2, 0.1)
