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_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
def __init__(
self,
conninfo: str = "",
*,
connection_class: Type[AsyncConnection[Any]] = AsyncConnection,
configure: Optional[Callable[[AsyncConnection[Any]],
Awaitable[None]]] = None,
reset: Optional[Callable[[AsyncConnection[Any]],
Awaitable[None]]] = None,
**kwargs: Any,
):
# https://bugs.python.org/issue42600
if sys.version_info < (3, 7):
raise e.NotSupportedError(
"async pool not supported before Python 3.7")
self.connection_class = connection_class
self._configure = configure
self._reset = reset
self._lock = asyncio.Lock()
self._waiting: Deque["AsyncClient"] = deque()
# to notify that the pool is full
self._pool_full_event: Optional[asyncio.Event] = None
self._sched = AsyncScheduler()
self._tasks: "asyncio.Queue[MaintenanceTask]" = asyncio.Queue()
self._workers: List[Task[None]] = []
super().__init__(conninfo, **kwargs)
self._sched_runner = create_task(self._sched.run(),
name=f"{self.name}-scheduler")
for i in range(self.num_workers):
t = create_task(
self.worker(self._tasks),
name=f"{self.name}-worker-{i}",
)
self._workers.append(t)
# populate the pool with initial min_size connections in background
for i in range(self._nconns):
self.run_task(AddConnection(self))
# Schedule a task to shrink the pool if connections over min_size have
# remained unused.
self.run_task(Schedule(self, ShrinkPool(self), self.max_idle))
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, retries):
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))
async for retry in retries:
with retry:
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_sched_error(caplog):
caplog.set_level(logging.WARNING, logger="psycopg")
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_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_stats_usage(dsn, retries):
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 for retry in retries:
with retry:
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 psycopg_pool.pool_async 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_queue_timeout(dsn, retries):
async def worker(n):
t0 = time()
try:
async with p.connection() as conn:
cur = await conn.execute(
"select pg_backend_pid() from pg_sleep(0.2)")
(pid, ) = await cur.fetchone() # type: ignore[misc]
except pool.PoolTimeout as e:
t1 = time()
errors.append((n, t1 - t0, e))
else:
t1 = time()
results.append((n, t1 - t0, pid))
async for retry in retries:
with retry:
results: List[Tuple[int, float, int]] = []
errors: List[Tuple[int, float, Exception]] = []
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) == 2
assert len(errors) == 2
for e in errors:
assert 0.1 < e[1] < 0.15
async def test_concurrency(aconn, what):
await aconn.set_autocommit(True)
e = [asyncio.Event() for i in range(3)]
async def worker(unlock, wait_on):
with pytest.raises(ProgrammingError) as ex:
async with aconn.transaction():
unlock.set()
await wait_on.wait()
await aconn.execute("select 1")
if what == "error":
1 / 0
elif what == "rollback":
raise Rollback()
else:
assert what == "commit"
if what == "error":
assert "transaction rollback" in str(ex.value)
assert isinstance(ex.value.__context__, ZeroDivisionError)
elif what == "rollback":
assert "transaction rollback" in str(ex.value)
assert isinstance(ex.value.__context__, Rollback)
else:
assert "transaction commit" in str(ex.value)
# Start a first transaction in a task
t1 = create_task(worker(unlock=e[0], wait_on=e[1]))
await e[0].wait()
# Start a nested transaction in a task
t2 = create_task(worker(unlock=e[1], wait_on=e[2]))
# Terminate the first transaction before the second does
await asyncio.gather(t1)
e[2].set()
await asyncio.gather(t2)
async def test_identify_closure(dsn, retries):
async def closer():
await asyncio.sleep(0.3)
await conn2.execute("select pg_terminate_backend(%s)",
[aconn.pgconn.backend_pid])
async for retry in retries:
with retry:
aconn = await psycopg.AsyncConnection.connect(dsn)
conn2 = await psycopg.AsyncConnection.connect(dsn)
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(psycopg.OperationalError):
await aconn.execute("select 1")
t1 = time.time()
assert 0.3 < t1 - t0 < 0.6
async def test_closed_queue(dsn):
async def w1():
async with p.connection() as conn:
e1.set() # Tell w0 that w1 got a connection
cur = await conn.execute("select 1")
assert await cur.fetchone() == (1, )
await e2.wait() # Wait until w0 has tested w2
success.append("w1")
async def w2():
try:
async with p.connection():
pass # unexpected
except pool.PoolClosed:
success.append("w2")
e1 = asyncio.Event()
e2 = asyncio.Event()
p = pool.AsyncConnectionPool(dsn, min_size=1)
await p.wait()
success: List[str] = []
t1 = create_task(w1())
# Wait until w1 has received a connection
await e1.wait()
t2 = create_task(w2())
# Wait until w2 is in the queue
while not p._waiting:
await asyncio.sleep(0)
await p.close()
# Wait for the workers to finish
e2.set()
await asyncio.gather(t1, t2)
assert len(success) == 2
async def test_closed_queue(dsn, retries):
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")
async for retry in retries:
with retry:
p = pool.AsyncConnectionPool(dsn, min_size=1)
success = []
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_identify_closure(dsn, retries):
async def closer():
await asyncio.sleep(0.2)
await conn2.execute("select pg_terminate_backend(%s)",
[aconn.pgconn.backend_pid])
async for retry in retries:
with retry:
aconn = await psycopg.AsyncConnection.connect(dsn)
conn2 = await psycopg.AsyncConnection.connect(dsn)
try:
t = create_task(closer())
t0 = time.time()
try:
with pytest.raises(psycopg.OperationalError):
await aconn.execute("select pg_sleep(1.0)")
t1 = time.time()
assert 0.2 < t1 - t0 < 0.4
finally:
await asyncio.gather(t)
finally:
await aconn.close()
await conn2.close()
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)
