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
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="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_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 psycopg.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, 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_identify_closure(aconn, dsn): conn2 = await psycopg.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(psycopg.OperationalError): await aconn.execute("select 1") t1 = time.time() assert 0.3 < t1 - t0 < 0.6
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)