diff --git a/test/conftest.py b/test/conftest.py index 811cad2..c3c40cd 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,3 +1,5 @@ +import time + import pytest from sqlalchemy import create_engine from sqlalchemy.orm import Session @@ -41,3 +43,24 @@ def autocommit_sess(engine): session = Session(ac_engine) yield session + + +@pytest.fixture(scope="function") +def wait_until(): + """Poll `fetch()` until `predicate(result)` is true, instead of a fixed sleep. + + Returns the last fetched value on success, so callers can use it directly. + Raises AssertionError with the last observed value if `timeout` is exceeded. + """ + def _wait_until(fetch, predicate, timeout=10, interval=0.1, description="condition"): + deadline = time.time() + timeout + result = None + while time.time() < deadline: + result = fetch() + if predicate(result): + return result + time.sleep(interval) + raise AssertionError( + f"Timed out after {timeout}s waiting for {description} (last value: {result!r})" + ) + return _wait_until diff --git a/test/test_http_requests_deleted_after_ttl.py b/test/test_http_requests_deleted_after_ttl.py index 071d795..906696e 100644 --- a/test/test_http_requests_deleted_after_ttl.py +++ b/test/test_http_requests_deleted_after_ttl.py @@ -1,65 +1,73 @@ -import time - import pytest from sqlalchemy import text -def test_http_responses_deleted_after_ttl(sess, autocommit_sess): +def test_http_responses_deleted_after_ttl(sess, autocommit_sess, wait_until): """Check that http responses will be deleted when they reach their ttl, not immediately but when the worker wakes again""" autocommit_sess.execute(text("alter system set pg_net.ttl to '1 second'")) autocommit_sess.execute(text("select net.worker_restart()")) autocommit_sess.execute(text("select net.wait_until_running()")) - # Create a request - (request_id,) = sess.execute(text( - """ - select net.http_get( - 'http://localhost:8080/anything' - ); - """ - )).fetchone() - - # Commit so background worker can start - sess.commit() - - # Confirm that the request was retrievable - response = sess.execute( - text( + try: + # Create a request + (request_id,) = sess.execute(text( """ - select * from net._http_collect_response(:request_id, async:=false); - """ - ), - {"request_id": request_id}, - ).fetchone() - assert response[0] == "SUCCESS" + select net.http_get( + 'http://localhost:8080/anything' + ); + """ + )).fetchone() - # Sleep until after request should have been deleted - time.sleep(1.1) + # Commit so background worker can start + sess.commit() - # Wake the worker manually, under normal operation this will happen when new requests are received - sess.execute(text("select net.wake()")) + # Confirm that the request was retrievable + response = sess.execute( + text( + """ + select * from net._http_collect_response(:request_id, async:=false); + """ + ), + {"request_id": request_id}, + ).fetchone() + assert response[0] == "SUCCESS" + + # Wait until the request's age exceeds the ttl + wait_until( + fetch=lambda: sess.execute( + text("select extract(epoch from clock_timestamp() - created) from net._http_response where id = :request_id"), + {"request_id": request_id}, + ).scalar(), + predicate=lambda age: age is None or age >= 1.0, + timeout=5, + description=f"request {request_id} to exceed the ttl", + ) - sess.commit() # commit so worker wakes + # Wake the worker manually, under normal operation this will happen when new requests are received + sess.execute(text("select net.wake()")) - time.sleep(0.1) # wait for deletion + sess.commit() # commit so worker wakes - # Ensure the response is now empty - (count,) = sess.execute( - text( - """ - select count(*) from net._http_response where id = :request_id; - """ - ), - {"request_id": request_id}, - ).fetchone() - assert count == 0 - - autocommit_sess.execute(text("alter system reset pg_net.ttl")) - autocommit_sess.execute(text("select net.worker_restart()")) - autocommit_sess.execute(text("select net.wait_until_running()")) + # Ensure the response is now empty + (count,) = wait_until( + fetch=lambda: sess.execute( + text("select count(*) from net._http_response where id = :request_id"), + {"request_id": request_id}, + ).fetchone(), + predicate=lambda r: r[0] == 0, + timeout=5, + description=f"deletion of expired response {request_id}", + ) + assert count == 0 + finally: + # Always restore ttl, even on assertion failure - it's set via ALTER + # SYSTEM so it would otherwise leak into and destabilize later tests. + autocommit_sess.execute(text("alter system reset pg_net.ttl")) + autocommit_sess.execute(text("select net.worker_restart()")) + autocommit_sess.execute(text("select net.wait_until_running()")) -def test_http_responses_will_complete_deletion(sess, autocommit_sess): +def test_http_responses_will_complete_deletion(sess, autocommit_sess, wait_until): """Check that http responses will keep being deleted until completion despite no new requests coming""" (request_id,) = sess.execute(text( @@ -95,42 +103,48 @@ def test_http_responses_will_complete_deletion(sess, autocommit_sess): autocommit_sess.execute(text("alter system set pg_net.batch_size to 2;")) autocommit_sess.execute(text("select pg_reload_conf();")) - # wait for ttl - time.sleep(1) - - # Wake the worker manually, under normal operation this will happen when new requests are received - sess.execute(text("select net.wake()")) - sess.commit() # commit so worker wakes - - time.sleep(0.1) - - (count,) = sess.execute( - text( - """ - select count(*) from net._http_response - """ + try: + # wait for ttl + wait_until( + fetch=lambda: sess.execute( + text("select extract(epoch from clock_timestamp() - created) from net._http_response where id = :request_id"), + {"request_id": request_id}, + ).scalar(), + predicate=lambda age: age is None or age >= 1.0, + timeout=5, + description=f"request {request_id} to exceed the ttl", ) - ).fetchone() - assert count == 2 - # wait for another batch - time.sleep(1.1) + # Wake the worker manually, under normal operation this will happen when new requests are received + sess.execute(text("select net.wake()")) + sess.commit() # commit so worker wakes - (count,) = sess.execute( - text( - """ - select count(*) from net._http_response - """ + # only one batch (batch_size=2) should be deleted by this wake + (count,) = wait_until( + fetch=lambda: sess.execute(text("select count(*) from net._http_response")).fetchone(), + predicate=lambda r: r[0] <= 2, + timeout=2, + description="first batch of expired responses to be deleted", ) - ).fetchone() - assert count == 0 - - autocommit_sess.execute(text("alter system reset pg_net.ttl")) - autocommit_sess.execute(text("alter system reset pg_net.batch_size")) - autocommit_sess.execute(text("select pg_reload_conf();")) + assert count == 2 + + # wait for another batch + (count,) = wait_until( + fetch=lambda: sess.execute(text("select count(*) from net._http_response")).fetchone(), + predicate=lambda r: r[0] == 0, + timeout=5, + description="remaining batch of expired responses to be deleted", + ) + assert count == 0 + finally: + # Always restore ttl/batch_size, even on assertion failure - they're set + # via ALTER SYSTEM so they would otherwise leak into later tests. + autocommit_sess.execute(text("alter system reset pg_net.ttl")) + autocommit_sess.execute(text("alter system reset pg_net.batch_size")) + autocommit_sess.execute(text("select pg_reload_conf();")) -def test_http_responses_will_delete_despite_restart(sess, autocommit_sess): +def test_http_responses_will_delete_despite_restart(sess, autocommit_sess, wait_until): """Check that http responses will keep being despite no new requests coming" and despite restart""" (request_id,) = sess.execute(text( @@ -168,20 +182,20 @@ def test_http_responses_will_delete_despite_restart(sess, autocommit_sess): autocommit_sess.execute(text("select net.worker_restart()")) autocommit_sess.execute(text("select net.wait_until_running()")) - # wait for ttl - time.sleep(1.1) - - (count,) = sess.execute( - text( - """ - select count(*) from net._http_response - """ + try: + # wait for ttl to elapse and the worker to clean up all expired responses + # (multiple worker cycles may be needed since batch_size=2 but there are 4 rows) + (count,) = wait_until( + fetch=lambda: sess.execute(text("select count(*) from net._http_response")).fetchone(), + predicate=lambda r: r[0] == 0, + timeout=20, + description="all expired responses to be deleted after restart", ) - ).fetchone() - assert count == 0 - - # reset - autocommit_sess.execute(text("alter system reset pg_net.ttl")) - autocommit_sess.execute(text("alter system reset pg_net.batch_size")) - autocommit_sess.execute(text("select net.worker_restart()")) - autocommit_sess.execute(text("select net.wait_until_running()")) + assert count == 0 + finally: + # Always restore ttl/batch_size, even on assertion failure - they're set + # via ALTER SYSTEM so they would otherwise leak into later tests. + autocommit_sess.execute(text("alter system reset pg_net.ttl")) + autocommit_sess.execute(text("alter system reset pg_net.batch_size")) + autocommit_sess.execute(text("select net.worker_restart()")) + autocommit_sess.execute(text("select net.wait_until_running()")) diff --git a/test/test_http_timeout.py b/test/test_http_timeout.py index bdf08ec..2fee4f2 100644 --- a/test/test_http_timeout.py +++ b/test/test_http_timeout.py @@ -1,10 +1,8 @@ -import time - import pytest import re from sqlalchemy import text -def test_http_get_timeout_reached(sess): +def test_http_get_timeout_reached(sess, wait_until): """net.http_get with timeout errs on a slow reply""" (request_id,) = sess.execute(text( @@ -16,16 +14,15 @@ def test_http_get_timeout_reached(sess): sess.commit() # wait for timeout - time.sleep(7) - - (content_type, content, response,timed_out) = sess.execute( - text( - """ - select content_type, content, error_msg, timed_out from net._http_response where id = :request_id; - """ - ), - {"request_id": request_id}, - ).fetchone() + (content_type, content, response, timed_out) = wait_until( + fetch=lambda: sess.execute( + text("select content_type, content, error_msg, timed_out from net._http_response where id = :request_id"), + {"request_id": request_id}, + ).fetchone(), + predicate=lambda r: r is not None, + timeout=10, + description=f"timeout response for request {request_id}", + ) assert content_type == None assert content == None @@ -33,7 +30,7 @@ def test_http_get_timeout_reached(sess): assert response.startswith("Timeout of 5000 ms reached") -def test_http_detailed_timeout(sess): +def test_http_detailed_timeout(sess, wait_until): """the timeout shows a detailed error msg""" pattern = r""" @@ -67,16 +64,15 @@ def test_http_detailed_timeout(sess): sess.commit() # wait for timeout - time.sleep(2.1) - - (content_type, content, response,timed_out) = sess.execute( - text( - """ - select content_type, content, error_msg, timed_out from net._http_response where id = :request_id; - """ - ), - {"request_id": request_id}, - ).fetchone() + (content_type, content, response, timed_out) = wait_until( + fetch=lambda: sess.execute( + text("select content_type, content, error_msg, timed_out from net._http_response where id = :request_id"), + {"request_id": request_id}, + ).fetchone(), + predicate=lambda r: r is not None, + timeout=10, + description=f"timeout response for request {request_id}", + ) match = regex.search(response) @@ -93,7 +89,7 @@ def test_http_detailed_timeout(sess): assert tcp_ssl_time > 0 assert http_time > 0 -def test_http_get_succeed_with_gt_timeout(sess): +def test_http_get_succeed_with_gt_timeout(sess, wait_until): """net.http_get with timeout succeeds when the timeout is greater than the slow reply response time""" (request_id,) = sess.execute(text( @@ -104,20 +100,19 @@ def test_http_get_succeed_with_gt_timeout(sess): sess.commit() - time.sleep(4.5) - - (status_code,) = sess.execute( - text( - """ - select status_code from net._http_response where id = :request_id; - """ - ), - {"request_id": request_id}, - ).fetchone() + (status_code,) = wait_until( + fetch=lambda: sess.execute( + text("select status_code from net._http_response where id = :request_id"), + {"request_id": request_id}, + ).fetchone(), + predicate=lambda r: r is not None, + timeout=10, + description=f"response for request {request_id}", + ) assert status_code == 200 -def test_many_slow_mixed_with_fast(sess): +def test_many_slow_mixed_with_fast(sess, wait_until): """many fast responses finish despite being mixed with slow responses, the fast responses will wait the timeout duration""" sess.execute(text( @@ -133,17 +128,20 @@ def test_many_slow_mixed_with_fast(sess): sess.commit() - # wait for timeouts - time.sleep(3) - - (request_successes, request_timeouts) = sess.execute(text( - """ - select - count(*) filter (where error_msg is null and status_code = 200) as request_successes, - count(*) filter (where error_msg is not null and error_msg like 'Timeout of 1000 ms reached%') as request_timeouts - from net._http_response; - """ - )).fetchone() + # wait for all 100 requests (successes + timeouts) to land + (request_successes, request_timeouts) = wait_until( + fetch=lambda: sess.execute(text( + """ + select + count(*) filter (where error_msg is null and status_code = 200) as request_successes, + count(*) filter (where error_msg is not null and error_msg like 'Timeout of 1000 ms reached%') as request_timeouts + from net._http_response; + """ + )).fetchone(), + predicate=lambda r: r[0] + r[1] == 100, + timeout=10, + description="100 mixed fast/slow responses in net._http_response", + ) assert request_successes == 50 assert request_timeouts == 50 diff --git a/test/test_privileges.py b/test/test_privileges.py index aca3b4d..64818fb 100644 --- a/test/test_privileges.py +++ b/test/test_privileges.py @@ -75,48 +75,53 @@ def test_net_on_new_role(sess): create role another; """)) - # Create a request - (request_id, current_user) = sess.execute(text( - """ - set local role to another; - select net.http_get( - 'http://localhost:8080/anything' - ), current_user; - """ - )).fetchone() - assert request_id == 1 - assert current_user == 'another' - - # Commit so background worker can start - sess.commit() - - # Confirm that the request was retrievable - response = sess.execute( - text( + try: + # Create a request + (request_id, current_user) = sess.execute(text( """ - set local role to another; - select *, current_user from net._http_collect_response(:request_id, async:=false); - """ - ), - {"request_id": request_id}, - ).fetchone() - assert response[0] == "SUCCESS" - assert response[3] == 'another' # current-user - - ## can use the net.worker_restart function - (res, current_user) = sess.execute( - text( - """ - set local role to another; - select net.worker_restart(), current_user; - """ - ) - ).fetchone() - assert res == True - assert current_user == 'another' - - sess.execute(text(""" - select net.wait_until_running(); - set local role postgres; - drop role another; - """)) + set local role to another; + select net.http_get( + 'http://localhost:8080/anything' + ), current_user; + """ + )).fetchone() + assert request_id == 1 + assert current_user == 'another' + + # Commit so background worker can start + sess.commit() + + # Confirm that the request was retrievable + response = sess.execute( + text( + """ + set local role to another; + select *, current_user from net._http_collect_response(:request_id, async:=false); + """ + ), + {"request_id": request_id}, + ).fetchone() + assert response[0] == "SUCCESS" + assert response[3] == 'another' # current-user + + ## can use the net.worker_restart function + (res, current_user) = sess.execute( + text( + """ + set local role to another; + select net.worker_restart(), current_user; + """ + ) + ).fetchone() + assert res == True + assert current_user == 'another' + + sess.execute(text("select net.wait_until_running();")) + finally: + # Always drop the role, even on assertion failure - otherwise a rerun + # against the same cluster fails with a duplicate-object error on + # `create role another;` instead of surfacing the real regression. + sess.execute(text(""" + set local role postgres; + drop role if exists another; + """)) diff --git a/test/test_stat_statements.py b/test/test_stat_statements.py index 6117c8c..d04614b 100644 --- a/test/test_stat_statements.py +++ b/test/test_stat_statements.py @@ -17,47 +17,51 @@ def test_query_stat_statements(sess): sess.execute(text( """ - create extension pg_stat_statements; + create extension if not exists pg_stat_statements; """ )) sess.commit() - time.sleep(1) + try: + time.sleep(1) - (old_calls,) = sess.execute(text( + (old_calls,) = sess.execute(text( + """ + select coalesce(sum(calls), 0) + from pg_stat_statements + where + query ilike '%DELETE FROM net._http_response r %' or + query ilike '%DELETE FROM net.http_request_queue%'; """ - select coalesce(sum(calls), 0) - from pg_stat_statements - where - query ilike '%DELETE FROM net._http_response r %' or - query ilike '%DELETE FROM net.http_request_queue%'; - """ - )).fetchone() - - # sleep for some time to see if new queries arrive - time.sleep(3) - - (new_calls,) = sess.execute(text( + )).fetchone() + + # sleep for some time to see if new queries arrive + time.sleep(3) + + (new_calls,) = sess.execute(text( + """ + select coalesce(sum(calls), 0) + from pg_stat_statements + where + query ilike '%DELETE FROM net._http_response r %' or + query ilike '%DELETE FROM net.http_request_queue%'; """ - select coalesce(sum(calls), 0) - from pg_stat_statements - where - query ilike '%DELETE FROM net._http_response r %' or - query ilike '%DELETE FROM net.http_request_queue%'; - """ - )).fetchone() - - assert new_calls == old_calls - - sess.execute(text( + )).fetchone() + + assert new_calls == old_calls + finally: + # Always drop the extension, even on assertion failure - otherwise the + # next test's unconditional `create extension pg_stat_statements;` + # fails outright with a duplicate-object error. + sess.execute(text( + """ + select pg_stat_statements_reset(); + drop extension if exists pg_stat_statements; """ - select pg_stat_statements_reset(); - drop extension pg_stat_statements; - """ - )) + )) - sess.commit() + sess.commit() def test_wakes_at_commit_time(sess): @@ -74,80 +78,83 @@ def test_wakes_at_commit_time(sess): sess.execute(text( """ - create extension pg_stat_statements; + create extension if not exists pg_stat_statements; """ )) sess.commit() - # wait for initial queries - time.sleep(1) - - (initial_calls,) = sess.execute(text( + try: + # wait for initial queries + time.sleep(1) + + (initial_calls,) = sess.execute(text( + """ + select coalesce(sum(calls), 0) + from pg_stat_statements + where + query ilike '%DELETE FROM net._http_response r %' or + query ilike '%DELETE FROM net.http_request_queue%'; """ - select coalesce(sum(calls), 0) - from pg_stat_statements - where - query ilike '%DELETE FROM net._http_response r %' or - query ilike '%DELETE FROM net.http_request_queue%'; - """ - )).fetchone() + )).fetchone() - assert initial_calls >= 0 + assert initial_calls >= 0 - sess.execute(text( + sess.execute(text( + """ + select net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,100); """ - select net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,100); - """ - )) + )) - sess.commit() + sess.commit() - # wait for reqs - time.sleep(2) + # wait for reqs + time.sleep(2) - (commit_calls,) = sess.execute(text( + (commit_calls,) = sess.execute(text( + """ + select coalesce(sum(calls), 0) + from pg_stat_statements + where + query ilike '%DELETE FROM net._http_response r %' or + query ilike '%DELETE FROM net.http_request_queue%'; """ - select coalesce(sum(calls), 0) - from pg_stat_statements - where - query ilike '%DELETE FROM net._http_response r %' or - query ilike '%DELETE FROM net.http_request_queue%'; - """ - )).fetchone() + )).fetchone() - assert commit_calls == initial_calls + 4 # only 4 queries should be made for the above requests - # 2 queries at wake, 2 extra to check if there are more rows to be processed + assert commit_calls == initial_calls + 4 # only 4 queries should be made for the above requests + # 2 queries at wake, 2 extra to check if there are more rows to be processed - # if the new requests are rollbacked/aborted, then no new queries will be made by the bg worker - sess.execute(text( + # if the new requests are rollbacked/aborted, then no new queries will be made by the bg worker + sess.execute(text( + """ + select net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,100); """ - select net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,100); - """ - )) + )) - sess.rollback() + sess.rollback() - # wait for requests - time.sleep(2) + # wait for requests + time.sleep(2) - (rollback_calls,) = sess.execute(text( + (rollback_calls,) = sess.execute(text( + """ + select coalesce(sum(calls), 0) + from pg_stat_statements + where + query ilike '%DELETE FROM net._http_response r %' or + query ilike '%DELETE FROM net.http_request_queue%'; """ - select coalesce(sum(calls), 0) - from pg_stat_statements - where - query ilike '%DELETE FROM net._http_response r %' or - query ilike '%DELETE FROM net.http_request_queue%'; - """ - )).fetchone() - - assert rollback_calls == commit_calls - - sess.execute(text( + )).fetchone() + + assert rollback_calls == commit_calls + finally: + # Always drop the extension, even on assertion failure - otherwise it + # leaks into whatever test runs next. + sess.execute(text( + """ + select pg_stat_statements_reset(); + drop extension if exists pg_stat_statements; """ - select pg_stat_statements_reset(); - drop extension pg_stat_statements; - """ - )) + )) - sess.commit() + sess.commit() diff --git a/test/test_user_db.py b/test/test_user_db.py index faf06c3..bb121c7 100644 --- a/test/test_user_db.py +++ b/test/test_user_db.py @@ -11,20 +11,24 @@ def test_net_with_different_username_dbname(sess, autocommit_sess): autocommit_sess.execute(text("select net.worker_restart()")) autocommit_sess.execute(text("select net.wait_until_running()")) - (username,datname) = sess.execute( - text( - """ - select usename, datname from pg_stat_activity where backend_type ilike '%pg_net%'; - """ - ) - ).fetchone() - assert username == 'pre_existing' - assert datname == 'pre_existing' - - autocommit_sess.execute(text("alter system reset pg_net.username")) - autocommit_sess.execute(text("alter system reset pg_net.database_name")) - autocommit_sess.execute(text("select net.worker_restart()")) - autocommit_sess.execute(text("select net.wait_until_running()")) + try: + (username,datname) = sess.execute( + text( + """ + select usename, datname from pg_stat_activity where backend_type ilike '%pg_net%'; + """ + ) + ).fetchone() + assert username == 'pre_existing' + assert datname == 'pre_existing' + finally: + # Always restore username/database_name, even on assertion failure - + # they're set via ALTER SYSTEM, and a leaked value here would point the + # worker at the wrong database for every other test in the suite. + autocommit_sess.execute(text("alter system reset pg_net.username")) + autocommit_sess.execute(text("alter system reset pg_net.database_name")) + autocommit_sess.execute(text("select net.worker_restart()")) + autocommit_sess.execute(text("select net.wait_until_running()")) def test_net_appname(sess): """Check that pg_stat_activity has appname set""" diff --git a/test/test_worker_behavior.py b/test/test_worker_behavior.py index 2682e31..42aa2f2 100644 --- a/test/test_worker_behavior.py +++ b/test/test_worker_behavior.py @@ -51,7 +51,7 @@ def test_success_when_worker_is_up(sess): assert result == '' -def test_worker_will_process_queue_when_up(sess): +def test_worker_will_process_queue_when_up(sess, wait_until): """when pg background worker is down and requests arrive, it will process them once it wakes up""" # check worker up @@ -68,12 +68,17 @@ def test_worker_will_process_queue_when_up(sess): assert restarted is not None assert restarted == True - time.sleep(0.1) - - # check worker down - (up,) = sess.execute(text(""" - select is_worker_up(); - """)).fetchone() + # check worker down. pg_stat_activity needs a fresh transaction on `sess` to + # see the worker's row disappear. + def _fetch_worker_up_1(): + sess.rollback() + return sess.execute(text("select is_worker_up()")).fetchone() + (up,) = wait_until( + fetch=_fetch_worker_up_1, + predicate=lambda r: r[0] == False, + timeout=5, + description="worker to be down after kill_worker()", + ) assert up is not None assert up == False @@ -103,24 +108,27 @@ def test_worker_will_process_queue_when_up(sess): sess.commit() - # wait until up - time.sleep(2.1) - - # check worker up - (up,) = sess.execute(text(""" - select is_worker_up(); - """)).fetchone() + # wait until up (the postmaster's restart backoff can take a few seconds). + # pg_stat_activity needs a fresh transaction on `sess` to see the new worker's row. + def _fetch_worker_up(): + sess.rollback() + return sess.execute(text("select is_worker_up()")).fetchone() + (up,) = wait_until( + fetch=_fetch_worker_up, + predicate=lambda r: r[0] == True, + timeout=15, + description="worker to come back up", + ) assert up is not None assert up == True # wait until new requests are done - time.sleep(1.1) - - (count,) = sess.execute(text( - """ - select count(*) from net.http_request_queue; - """ - )).fetchone() + (count,) = wait_until( + fetch=lambda: sess.execute(text("select count(*) from net.http_request_queue")).fetchone(), + predicate=lambda r: r[0] == 0, + timeout=5, + description="queued requests to be processed", + ) assert count == 0 @@ -141,29 +149,35 @@ def test_can_delete_rows_while_processing_queue(sess, autocommit_sess): autocommit_sess.execute(text("select net.worker_restart();")) autocommit_sess.execute(text("select net.wait_until_running();")) - sess.execute(text( + try: + sess.execute(text( + """ + select net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,10); """ - select net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,10); - """ - )) + )) - sess.commit() + sess.commit() - # leave time for some processing - time.sleep(0.1) + # This intentionally races a short, fixed sleep against the (slow, + # batch_size=1) worker to catch it mid-queue, then asserts most rows are + # still unprocessed. Polling for "some processing happened" would fight + # that intent, so this stays a fixed sleep. + time.sleep(0.1) - (count,) = sess.execute(text( + (count,) = sess.execute(text( + """ + WITH deleted AS (DELETE FROM net.http_request_queue RETURNING *) SELECT count(*) FROM deleted; """ - WITH deleted AS (DELETE FROM net.http_request_queue RETURNING *) SELECT count(*) FROM deleted; - """ - )).fetchone() - assert count > 1 - - sess.commit() + )).fetchone() + assert count > 1 - autocommit_sess.execute(text("alter system reset pg_net.batch_size")) - autocommit_sess.execute(text("select net.worker_restart()")) - autocommit_sess.execute(text("select net.wait_until_running()")) + sess.commit() + finally: + # Always restore batch_size, even on assertion failure - it's set via + # ALTER SYSTEM so it would otherwise leak into and slow down later tests. + autocommit_sess.execute(text("alter system reset pg_net.batch_size")) + autocommit_sess.execute(text("select net.worker_restart()")) + autocommit_sess.execute(text("select net.wait_until_running()")) def test_truncate_wait_while_processing_queue(sess, autocommit_sess): @@ -175,34 +189,37 @@ def test_truncate_wait_while_processing_queue(sess, autocommit_sess): autocommit_sess.execute(text("select net.worker_restart();")) autocommit_sess.execute(text("select net.wait_until_running();")) - sess.execute(text( + try: + sess.execute(text( + """ + select net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,10); """ - select net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,10); - """ - )) - sess.commit() + )) + sess.commit() - # truncate succeeds fast, despite the worker still processing the queue 1 by 1 - sess.execute(text( + # truncate succeeds fast, despite the worker still processing the queue 1 by 1 + sess.execute(text( + """ + truncate net.http_request_queue; """ - truncate net.http_request_queue; - """ - )) + )) - # now the queue will be empty - (count,) = sess.execute(text( + # now the queue will be empty + (count,) = sess.execute(text( + """ + select count(*) from net.http_request_queue; """ - select count(*) from net.http_request_queue; - """ - )).fetchone() - assert count == 0 - - autocommit_sess.execute(text("alter system reset pg_net.batch_size")) - autocommit_sess.execute(text("select net.worker_restart()")) - autocommit_sess.execute(text("select net.wait_until_running()")) + )).fetchone() + assert count == 0 + finally: + # Always restore batch_size, even on assertion failure - it's set via + # ALTER SYSTEM so it would otherwise leak into and slow down later tests. + autocommit_sess.execute(text("alter system reset pg_net.batch_size")) + autocommit_sess.execute(text("select net.worker_restart()")) + autocommit_sess.execute(text("select net.wait_until_running()")) -def test_no_failure_on_drop_extension(sess): +def test_no_failure_on_drop_extension(sess, wait_until): """while waiting for a slow request, a drop extension should wait and not crash the worker""" (request_id,) = sess.execute(text(""" @@ -212,7 +229,9 @@ def test_no_failure_on_drop_extension(sess): sess.commit() - # wait until processing + # wait until processing. There's no externally observable signal for + # "the worker has picked up this request" (it's a black box until the + # slow reply completes), so this stays a fixed sleep. time.sleep(1) sess.execute(text(""" @@ -221,80 +240,95 @@ def test_no_failure_on_drop_extension(sess): sess.commit() - # wait until request is finished - time.sleep(3) - - (up,) = sess.execute(text(""" - select is_worker_up(); - """)).fetchone() + # wait until the in-flight request finishes and the worker recovers. + # pg_stat_activity needs a fresh transaction on `sess` to see the new worker's row. + def _fetch_worker_up(): + sess.rollback() + return sess.execute(text("select is_worker_up()")).fetchone() + (up,) = wait_until( + fetch=_fetch_worker_up, + predicate=lambda r: r[0] == True, + timeout=10, + description="worker to recover after drop extension", + ) assert up is not None assert up == True -def test_worker_will_keep_processing_queue_when_restarted(sess, autocommit_sess): +def test_worker_will_keep_processing_queue_when_restarted(sess, autocommit_sess, wait_until): """when the background worker is restarted while working, it will pick up the remaining requests""" autocommit_sess.execute(text("alter system set pg_net.batch_size to '1';")) autocommit_sess.execute(text("select net.worker_restart();")) autocommit_sess.execute(text("select net.wait_until_running();")) - sess.execute(text( + try: + sess.execute(text( + """ + select net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,5); """ - select net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,5); - """ - )) + )) - sess.commit() + sess.commit() - # one restart will likely keep the worker awake since the wake signal could still be on, so do two restarts - # to ensure the wake signal is cleared - sess.execute(text( + # one restart will likely keep the worker awake since the wake signal could still be on, so do two restarts + # to ensure the wake signal is cleared. There's no observable signal for "the wake flag is + # cleared", so these stay fixed sleeps. + sess.execute(text( + """ + select net.worker_restart(); + select net.wait_until_running(); """ - select net.worker_restart(); - select net.wait_until_running(); - """ - )) + )) - time.sleep(0.1) + time.sleep(0.1) - sess.execute(text( + sess.execute(text( + """ + select net.worker_restart(); + select net.wait_until_running(); """ - select net.worker_restart(); - select net.wait_until_running(); - """ - )) + )) - time.sleep(0.1) - - (status_code,count) = sess.execute(text( - """ - select status_code, count(*) from net._http_response group by status_code; - """ - )).fetchone() - - # at most 2 requests should have finished by now because of the low batch_size - assert count <= 2 - assert count > 0 # at least 1 request should be finished - assert status_code == 200 - - # if we sleep for 4 seconds the whole 5 requests should be finished - time.sleep(4) - - (status_code,count) = sess.execute(text( - """ - select status_code, count(*) from net._http_response group by status_code; - """ - )).fetchone() - - assert status_code == 200 - assert count == 5 - - autocommit_sess.execute(text("alter system reset pg_net.batch_size")) - autocommit_sess.execute(text("select net.worker_restart()")) - autocommit_sess.execute(text("select net.wait_until_running()")) + time.sleep(0.1) + (status_code,count) = sess.execute(text( + """ + select status_code, count(*) from net._http_response group by status_code; + """ + )).fetchone() -def test_new_requests_get_attended_asap(sess): + # at most 2 requests should have finished by now because of the low batch_size. This is an + # intentional snapshot of a partial, in-flight state, so it isn't converted to polling. + assert count <= 2 + assert count > 0 # at least 1 request should be finished + assert status_code == 200 + + # wait until the whole 5 requests are finished. Generous timeout, matching + # the slack used elsewhere in this file for restart/throttle-heavy cases: + # with batch_size=1 each request costs ~1s of deliberate throttling (see + # worker.c's post-batch wait), plus the two restarts above each pay a + # postmaster restart backoff, so this needs more slack than a normal poll. + (status_code,count) = wait_until( + fetch=lambda: sess.execute(text( + "select status_code, count(*) from net._http_response group by status_code" + )).fetchone(), + predicate=lambda r: r is not None and r[1] == 5, + timeout=30, + description="all 5 requests to finish processing", + ) + + assert status_code == 200 + assert count == 5 + finally: + # Always restore batch_size, even on assertion failure - it's set via + # ALTER SYSTEM so it would otherwise leak into and slow down later tests. + autocommit_sess.execute(text("alter system reset pg_net.batch_size")) + autocommit_sess.execute(text("select net.worker_restart()")) + autocommit_sess.execute(text("select net.wait_until_running()")) + + +def test_new_requests_get_attended_asap(sess, wait_until): """new requests get attended as soon as possible""" sess.execute(text( @@ -305,20 +339,20 @@ def test_new_requests_get_attended_asap(sess): sess.commit() - # less than a second - time.sleep(0.1) - - (status_code,count) = sess.execute(text( - """ - select status_code, count(*) from net._http_response group by status_code; - """ - )).fetchone() + (status_code,count) = wait_until( + fetch=lambda: sess.execute(text( + "select status_code, count(*) from net._http_response group by status_code" + )).fetchone(), + predicate=lambda r: r is not None and r[1] == 10, + timeout=5, + description="all 10 new requests to be attended to", + ) assert status_code == 200 assert count == 10 -def test_direct_inserts_no_requests(sess): +def test_direct_inserts_no_requests(sess, wait_until): """direct insertions to the net.http_request_queue doesn't trigger new requests""" sess.execute(text( @@ -335,7 +369,8 @@ def test_direct_inserts_no_requests(sess): sess.commit() - # wait for req + # This proves a negative (no processing happens without a wake), so there's no positive + # condition to poll for - it stays a fixed observation window. time.sleep(0.1) # no response @@ -367,19 +402,20 @@ def test_direct_inserts_no_requests(sess): sess.commit() # wait for req - time.sleep(0.1) - - (status_code, count) = sess.execute(text( - """ - select status_code, count(*) from net._http_response group by status_code; - """ - )).fetchone() + (status_code, count) = wait_until( + fetch=lambda: sess.execute(text( + "select status_code, count(*) from net._http_response group by status_code" + )).fetchone(), + predicate=lambda r: r is not None and r[1] == 1, + timeout=5, + description="the woken request to be processed", + ) assert status_code == 200 assert count == 1 -def test_processing_survives_postmaster_crash(): +def test_processing_survives_postmaster_crash(wait_until): """the queue will continue processing even when a postmaster crash or restart happens""" engine = create_engine("postgresql:///postgres") @@ -392,58 +428,88 @@ def test_processing_survives_postmaster_crash(): tmp_sess.execute(text("select net.worker_restart();")) tmp_sess.execute(text("select net.wait_until_running();")) - tmp_sess.execute(text( + try: + tmp_sess.execute(text( + """ + select net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,10); """ - select net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,10); - """ - )).fetchone() - - (count,) = tmp_sess.execute(text( - """ - select count(*) from net.http_request_queue; - """ - )).fetchone() - assert count == 10 - - engine.dispose() - - pgdata_env = os.getenv('PGDATA') - subprocess.run(["pg_ctl", "restart", "-D", pgdata_env]) - - # give it some time to finish restart - time.sleep(1) - - engine = create_engine("postgresql:///postgres") - ac_engine = engine.execution_options(isolation_level="AUTOCOMMIT") - tmp_sess = Session(ac_engine) - - # give it enough time to finish processing the queue - time.sleep(1) - - (count,) = tmp_sess.execute(text( - """ - select count(*) from net.http_request_queue; - """ - )).fetchone() - assert count == 0 - - (status_code,count) = tmp_sess.execute(text( - """ - select status_code, count(*) from net._http_response group by status_code; - """ - )).fetchone() - - assert status_code == 200 - assert count == 10 - - tmp_sess.execute(text("alter system reset pg_net.batch_size")) - tmp_sess.execute(text("select net.worker_restart()")) - tmp_sess.execute(text("select net.wait_until_running()")) - - engine.dispose() + )).fetchone() + # The worker is already running (see wait_until_running() above), so by the time this + # count is read it may have already drained a batch into net._http_response. Check both + # tables together instead of assuming the queue alone still holds all 10. + (queue_count,) = tmp_sess.execute(text( + """ + select count(*) from net.http_request_queue; + """ + )).fetchone() + (response_count,) = tmp_sess.execute(text( + """ + select count(*) from net._http_response; + """ + )).fetchone() + assert queue_count + response_count == 10 + + engine.dispose() + + pgdata_env = os.getenv('PGDATA') + subprocess.run(["pg_ctl", "restart", "-D", pgdata_env]) + + # wait for postgres to accept connections again after the restart + engine = create_engine("postgresql:///postgres") + ac_engine = engine.execution_options(isolation_level="AUTOCOMMIT") + + def _try_connect(): + try: + s = Session(ac_engine) + s.execute(text("select 1")) + return s + except Exception: + return None + + tmp_sess = wait_until( + fetch=_try_connect, + predicate=lambda s: s is not None, + timeout=15, + description="postgres to accept connections after restart", + ) + + # give it enough time to finish processing the queue + (count,) = wait_until( + fetch=lambda: tmp_sess.execute(text("select count(*) from net.http_request_queue")).fetchone(), + predicate=lambda r: r[0] == 0, + timeout=10, + description="queued requests to be processed after restart", + ) + assert count == 0 + + (status_code,count) = tmp_sess.execute(text( + """ + select status_code, count(*) from net._http_response group by status_code; + """ + )).fetchone() -def test_worker_writes_increment_pgstat_counters(sess, autocommit_sess): + assert status_code == 200 + assert count == 10 + finally: + # Always restore batch_size, even on assertion failure - it's set via + # ALTER SYSTEM so it would otherwise leak into and slow down later tests. + # `tmp_sess`/`engine` may point at the pre-restart or post-restart + # connection depending on where a failure happened above; best-effort + # cleanup with whichever is currently live, without masking the + # original failure if the connection itself is unusable (e.g. postgres + # never came back up after the restart). + try: + tmp_sess.execute(text("alter system reset pg_net.batch_size")) + tmp_sess.execute(text("select net.worker_restart()")) + tmp_sess.execute(text("select net.wait_until_running()")) + except Exception: + pass + + engine.dispose() + + +def test_worker_writes_increment_pgstat_counters(sess, autocommit_sess, wait_until): """the worker's INSERTs into net._http_response must be reflected in pg_stat_user_tables. Without this, autovacuum/autoanalyze can never be scheduled and the table silently bloats. @@ -468,13 +534,13 @@ def test_worker_writes_increment_pgstat_counters(sess, autocommit_sess): # Wait until the worker has actually drained the queue and written all # responses to net._http_response. Don't assume "30 rows" - the worker # may pick up the queue in chunks depending on wake() coalescing. - for _ in range(20): - time.sleep(0.5) - (queue_count,) = sess.execute(text( - "select count(*) from net.http_request_queue;" - )).fetchone() - if queue_count == 0: - break + wait_until( + fetch=lambda: sess.execute(text("select count(*) from net.http_request_queue")).fetchone(), + predicate=lambda r: r[0] == 0, + timeout=10, + interval=0.5, + description="worker to drain net.http_request_queue", + ) # Confirm the worker actually wrote rows, otherwise the pgstat assertion # below would be meaningless. @@ -488,17 +554,16 @@ def test_worker_writes_increment_pgstat_counters(sess, autocommit_sess): # worker's first flush attempt is normally a hit (last_flush is far in # the past after a long idle), but we allow generous slack here so an # off-by-a-tick scheduling doesn't flake the suite. - deadline = time.time() + 30.0 - resp_ins = 0 - resp_mod = 0 - while time.time() < deadline: - (resp_ins, resp_mod) = autocommit_sess.execute(text(""" + (resp_ins, resp_mod) = wait_until( + fetch=lambda: autocommit_sess.execute(text(""" select n_tup_ins, n_mod_since_analyze from pg_stat_user_tables where relname='_http_response'; - """)).fetchone() - if resp_ins > 0: - break - time.sleep(0.5) + """)).fetchone(), + predicate=lambda r: r[0] > 0, + timeout=30, + interval=0.5, + description="net._http_response pgstat counters to reflect worker INSERTs", + ) assert resp_ins > 0, ( f"net._http_response.n_tup_ins is still 0 after 30s. " @@ -512,7 +577,7 @@ def test_worker_writes_increment_pgstat_counters(sess, autocommit_sess): ) -def test_worker_writes_trigger_autoanalyze_on_http_response(sess, autocommit_sess): +def test_worker_writes_trigger_autoanalyze_on_http_response(sess, autocommit_sess, wait_until): """autoanalyze on net._http_response must fire after the worker writes enough rows. Without working pgstat counters, autovacuum/autoanalyze never get scheduled and the table bloats - this is the primary symptom @@ -525,71 +590,79 @@ def test_worker_writes_trigger_autoanalyze_on_http_response(sess, autocommit_ses # Make autovacuum eager *before* generating traffic so the launcher is # already running on a 1s naptime by the time stats threshold is crossed. # autovacuum_naptime is PGC_SIGHUP (reloadable). Give the reload a moment - # to propagate to the launcher. + # to propagate to the launcher - there's no SQL-observable signal for + # "the launcher picked up the new naptime", so this stays a fixed sleep. autocommit_sess.execute(text("alter system set autovacuum_naptime = '1s';")) autocommit_sess.execute(text("select pg_reload_conf();")) - time.sleep(1) - # Per-table: trip the autoanalyze threshold after a handful of rows. - # Reloptions take effect immediately; no reload required. - autocommit_sess.execute(text(""" - alter table net._http_response set ( - autovacuum_analyze_threshold = 10, - autovacuum_analyze_scale_factor = 0, - autovacuum_vacuum_threshold = 10, - autovacuum_vacuum_scale_factor = 0 - ); - """)) - - autocommit_sess.execute(text( - "select pg_stat_reset_single_table_counters('net._http_response'::regclass);" - )) - - # Drive 30 inserts through the worker. 30 is well above the threshold (10). - sess.execute(text(""" - select net.http_get('http://localhost:8080/pathological?status=200') - from generate_series(1,30); - """)) - sess.commit() - - # 30s budget covers worst-case worker pgstat flush (PGSTAT_MIN_INTERVAL - # = 1s slack) + worst-case launcher cycle (autovacuum_max_workers=3, - # 3 databases at 1s naptime each ~= 3s/cycle, with 2-3 cycle slack) + - # autoanalyze worker spawn + ANALYZE on a tiny table (sub-second). - # Real wall time on a clean rig is typically ~2-5s; the slack is to - # absorb test-rig load and not flake. - deadline = time.time() + 30.0 - autoanalyze_count = 0 - while time.time() < deadline: - (autoanalyze_count,) = autocommit_sess.execute(text(""" - select autoanalyze_count - from pg_stat_user_tables where relname='_http_response'; - """)).fetchone() - if autoanalyze_count > 0: - break - time.sleep(0.5) - - assert autoanalyze_count > 0, ( - "autoanalyze never fired on net._http_response within 30s. " - "Worker writes are not making pgstat threshold visible to the " - "autovacuum launcher - the customer-facing symptom (silent bloat) " - "would manifest in production." - ) - - # Cleanup: restore defaults so we don't bleed into other tests. - autocommit_sess.execute(text(""" - alter table net._http_response reset ( - autovacuum_analyze_threshold, - autovacuum_analyze_scale_factor, - autovacuum_vacuum_threshold, - autovacuum_vacuum_scale_factor - ); - """)) - autocommit_sess.execute(text("alter system reset autovacuum_naptime;")) - autocommit_sess.execute(text("select pg_reload_conf();")) - - -def test_worker_reports_activity_in_pg_stat_activity(sess, autocommit_sess): + try: + time.sleep(1) + + # Per-table: trip the autoanalyze threshold after a handful of rows. + # Reloptions take effect immediately; no reload required. + autocommit_sess.execute(text(""" + alter table net._http_response set ( + autovacuum_analyze_threshold = 10, + autovacuum_analyze_scale_factor = 0, + autovacuum_vacuum_threshold = 10, + autovacuum_vacuum_scale_factor = 0 + ); + """)) + + autocommit_sess.execute(text( + "select pg_stat_reset_single_table_counters('net._http_response'::regclass);" + )) + + # Drive 30 inserts through the worker. 30 is well above the threshold (10). + sess.execute(text(""" + select net.http_get('http://localhost:8080/pathological?status=200') + from generate_series(1,30); + """)) + sess.commit() + + # 30s budget covers worst-case worker pgstat flush (PGSTAT_MIN_INTERVAL + # = 1s slack) + worst-case launcher cycle (autovacuum_max_workers=3, + # 3 databases at 1s naptime each ~= 3s/cycle, with 2-3 cycle slack) + + # autoanalyze worker spawn + ANALYZE on a tiny table (sub-second). + # Real wall time on a clean rig is typically ~2-5s; the slack is to + # absorb test-rig load and not flake. + (autoanalyze_count,) = wait_until( + fetch=lambda: autocommit_sess.execute(text(""" + select autoanalyze_count + from pg_stat_user_tables where relname='_http_response'; + """)).fetchone(), + predicate=lambda r: r[0] > 0, + timeout=30, + interval=0.5, + description="autoanalyze to fire on net._http_response", + ) + + assert autoanalyze_count > 0, ( + "autoanalyze never fired on net._http_response within 30s. " + "Worker writes are not making pgstat threshold visible to the " + "autovacuum launcher - the customer-facing symptom (silent bloat) " + "would manifest in production." + ) + finally: + # Always restore defaults, even on assertion failure - `autovacuum_naptime` + # is a system-wide GUC that would otherwise leak into every other test/ + # database for the rest of the suite. The table reloptions are lower risk + # (the `sess` fixture drops+recreates the table with the extension), but + # are restored too for a clean baseline regardless of where net._http_response + # ends up. + autocommit_sess.execute(text(""" + alter table net._http_response reset ( + autovacuum_analyze_threshold, + autovacuum_analyze_scale_factor, + autovacuum_vacuum_threshold, + autovacuum_vacuum_scale_factor + ); + """)) + autocommit_sess.execute(text("alter system reset autovacuum_naptime;")) + autocommit_sess.execute(text("select pg_reload_conf();")) + + +def test_worker_reports_activity_in_pg_stat_activity(sess, autocommit_sess, wait_until): """the pg_net worker must call pgstat_report_activity() so its row in pg_stat_activity has a valid state column. """ @@ -599,15 +672,14 @@ def test_worker_reports_activity_in_pg_stat_activity(sess, autocommit_sess): # Wait for the worker to drain any leftover work from previous tests # and settle into idle. Polling makes this robust regardless of what # ran before. - deadline = time.time() + 5.0 - state = None - while time.time() < deadline: - (state,) = autocommit_sess.execute(text( - "select state from pg_stat_activity where backend_type ilike '%pg_net%';" - )).fetchone() - if state == 'idle': - break - time.sleep(0.1) + (state,) = wait_until( + fetch=lambda: autocommit_sess.execute(text( + "select state from pg_stat_activity where backend_type ilike '%pg_net%'" + )).fetchone(), + predicate=lambda r: r[0] == 'idle', + timeout=5, + description="pg_net worker to settle into idle", + ) assert state == 'idle', ( f"pg_net worker state expected 'idle' at rest, got {state!r}. " "Without pgstat_report_activity(STATE_IDLE, ...) the state column " @@ -622,16 +694,19 @@ def test_worker_reports_activity_in_pg_stat_activity(sess, autocommit_sess): # Poll for 'active' for up to 5s. The slow request keeps the worker # busy for ~2s, so we have a wide observation window. - deadline = time.time() + 5.0 saw_active = False - while time.time() < deadline: - (state,) = autocommit_sess.execute(text( - "select state from pg_stat_activity where backend_type ilike '%pg_net%';" - )).fetchone() - if state == 'active': - saw_active = True - break - time.sleep(0.1) + try: + wait_until( + fetch=lambda: autocommit_sess.execute(text( + "select state from pg_stat_activity where backend_type ilike '%pg_net%'" + )).fetchone(), + predicate=lambda r: r[0] == 'active', + timeout=5, + description="pg_net worker to be observed as active", + ) + saw_active = True + except AssertionError: + pass assert saw_active, ( "pg_net worker state was never observed as 'active' during a slow "