import multiprocessing import os from pathlib import Path from types import SimpleNamespace import sys import tempfile import threading import time import unittest from unittest import mock import uuid ROOT = Path(__file__).resolve().parents[1] APP_DIR = ROOT / 'app' sys.path.insert(0, str(APP_DIR)) import console_runner from process_identity import current_process_identity from result_spool import ResultSpool, SpoolCapacityError from scanner_db import ScannerDB class SharedPublisherDB: def __init__(self, lease, entered, release_gate, ingest_count, hold=False, owner_state=True): self.conn = SimpleNamespace(is_postgres=True) self.lease = lease self.entered = entered self.release_gate = release_gate self.ingest_count = ingest_count self.hold = hold self.owner_state = owner_state self.owns_lease = False def require_runtime_safety_schema(self): return True def try_acquire_result_spool_publisher(self): self.owns_lease = self.lease.acquire(False) return self.owns_lease def result_spool_publisher_state(self): if not self.owner_state: return { 'status': 'held', 'authenticated': False, 'application_name': 'unknown', 'holder_identity': 'unknown', } return { 'status': 'held', 'authenticated': True, 'application_name': 'truf-source:live_owner', 'pid': 123, 'state': 'active', 'holder_identity': '123:fixture-start', } def release_result_spool_publisher(self): if not self.owns_lease: return False self.owns_lease = False self.lease.release() return True def ingest_scan_event(self, envelope): with self.ingest_count.get_lock(): self.ingest_count.value += 1 if self.hold: self.entered.set() if not self.release_gate.wait(15): raise RuntimeError('fixture publisher release timed out') return { 'ingested': True, 'scan_event_id': envelope['scan_event_id'], 'scan_event_hash': envelope['scan_event_hash'], 'duplicate': False, } def publisher_worker( directory, lease, entered, release_gate, ingest_count, outcomes, hold=False, stop_event=None, ): try: spool = ResultSpool( directory, max_event_bytes=4096, max_events=20, max_total_bytes=65536, min_free_bytes=0, lock_timeout_sec=2, ) db = SharedPublisherDB(lease, entered, release_gate, ingest_count, hold=hold) drained = console_runner.wait_for_result_spool_ready( spool, db, stop_event=stop_event, wait_seconds=0.05, diagnostic_interval=0.2, ) outcomes.put(('ok', len(drained))) except KeyboardInterrupt: outcomes.put(('stopped', 0)) except BaseException as exc: outcomes.put((type(exc).__name__, str(exc))) class CapacityProgressDB: conn = SimpleNamespace(is_postgres=True) def require_runtime_safety_schema(self): return True def result_spool_reservation_progress(self, reservations): return { 'safe_progress': bool(reservations), 'counts': { 'exact_live': len(reservations), 'exact_expired': 0, 'stale_or_reassigned': 0, }, 'earliest_progress_in_sec': 60, } def capacity_worker(directory, index, start, stop_event, outcomes): try: spool = ResultSpool( directory, max_event_bytes=1024, max_events=100, max_total_bytes=2500, min_free_bytes=0, lock_timeout_sec=2, ) start.wait(10) reservation = console_runner.reserve_result_spool_claims( spool, CapacityProgressDB(), f'source-{index}:{os.getpid()}:cycle', 1, 60, stop_event=stop_event, wait_seconds=0.05, diagnostic_interval=1, ) time.sleep(0.05) if not spool.release_reservation(reservation): raise RuntimeError('worker reservation release failed') outcomes.put(('acquired', index)) except KeyboardInterrupt: outcomes.put(('stopped', index)) except BaseException as exc: outcomes.put((type(exc).__name__, index)) class ResultSpoolBackpressureTests(unittest.TestCase): @staticmethod def context(): return multiprocessing.get_context('spawn') def test_completion_handoff_cannot_be_misread_as_reservation_renewal_loss(self): class RacingSpool: def __init__(self): self.write_entered = threading.Event() self.allow_write = threading.Event() self.renew_calls = 0 def write_event(self, *args, **kwargs): self.write_entered.set() if not self.allow_write.wait(5): raise RuntimeError('fixture write release timed out') return SimpleNamespace(event_id='durable-event') def renew_reservation(self, reservation_id, lease_seconds): self.renew_calls += 1 return False spool = RacingSpool() lease_token = 'lease-token' active = {lease_token} reserved = {lease_token} active_lock = threading.Lock() outcomes = {} renewal_attempted = threading.Event() writer = threading.Thread(target=lambda: outcomes.__setitem__('write', console_runner.write_reserved_result_spool_event( spool, {}, 'reservation', 7, lease_token, reserved, active, active_lock, ).event_id )) def renew(): renewal_attempted.set() outcomes['renew'] = console_runner.renew_active_result_spool_reservation( spool, 'reservation', 1800, reserved, active_lock, ) writer.start() self.assertTrue(spool.write_entered.wait(5)) heartbeat = threading.Thread(target=renew) heartbeat.start() self.assertTrue(renewal_attempted.wait(5)) spool.allow_write.set() writer.join(5) heartbeat.join(5) self.assertFalse(writer.is_alive()) self.assertFalse(heartbeat.is_alive()) self.assertEqual(outcomes, {'write': 'durable-event', 'renew': True}) self.assertEqual(spool.renew_calls, 0) self.assertEqual(active, set()) self.assertEqual(reserved, set()) def test_absent_reservation_with_active_token_still_fails_closed(self): spool = SimpleNamespace(renew_reservation=mock.Mock(return_value=False)) reserved = {'lease-token'} self.assertFalse(console_runner.renew_active_result_spool_reservation( spool, 'reservation', 1800, reserved, threading.Lock(), )) spool.renew_reservation.assert_called_once_with('reservation', 1800) def test_two_publishers_wait_handoff_and_ingest_once(self): with tempfile.TemporaryDirectory() as temp_dir: directory = os.path.join(temp_dir, 'spool') spool = ResultSpool( directory, max_event_bytes=4096, max_events=20, max_total_bytes=65536, min_free_bytes=0, ) reservation = spool.reserve_claims('unrelated-live-claim', 1, 60, bytes_per_event=1024) event_id = str(uuid.uuid4()) spool.write_event({'version': 1, 'scan_event_id': event_id, 'result': {}}) context = self.context() lease = context.Lock() entered = context.Event() release_gate = context.Event() ingest_count = context.Value('i', 0) outcomes = context.Queue() owner = context.Process( target=publisher_worker, args=(directory, lease, entered, release_gate, ingest_count, outcomes, True), ) owner.start() self.assertTrue(entered.wait(10)) waiter = context.Process( target=publisher_worker, args=(directory, lease, entered, release_gate, ingest_count, outcomes), ) waiter.start() time.sleep(0.3) release_gate.set() for process in (owner, waiter): process.join(20) self.assertFalse(process.is_alive()) self.assertEqual(process.exitcode, 0) results = sorted(outcomes.get(timeout=5) for _ in range(2)) self.assertEqual(results, [('ok', 0), ('ok', 1)]) self.assertEqual(ingest_count.value, 1) self.assertFalse(os.path.exists(spool._event_path(event_id))) self.assertTrue(os.path.exists(spool._reservation_path(reservation))) self.assertTrue(spool.release_reservation(reservation)) def test_waiter_honors_shutdown_without_touching_pending_event(self): with tempfile.TemporaryDirectory() as temp_dir: directory = os.path.join(temp_dir, 'spool') spool = ResultSpool( directory, max_event_bytes=4096, max_events=20, max_total_bytes=65536, min_free_bytes=0, ) event_id = str(uuid.uuid4()) spool.write_event({'version': 1, 'scan_event_id': event_id, 'result': {}}) context = self.context() lease = context.Lock() entered = context.Event() release_gate = context.Event() stop_event = context.Event() ingest_count = context.Value('i', 0) outcomes = context.Queue() owner = context.Process( target=publisher_worker, args=(directory, lease, entered, release_gate, ingest_count, outcomes, True), ) owner.start() self.assertTrue(entered.wait(10)) waiter = context.Process( target=publisher_worker, args=(directory, lease, entered, release_gate, ingest_count, outcomes, False, stop_event), ) waiter.start() stop_event.set() waiter.join(15) self.assertFalse(waiter.is_alive()) self.assertEqual(waiter.exitcode, 0) waiter_result = outcomes.get(timeout=5) self.assertEqual(waiter_result, ('stopped', 0)) self.assertTrue(os.path.exists(spool._event_path(event_id))) release_gate.set() owner.join(20) self.assertFalse(owner.is_alive()) self.assertEqual(owner.exitcode, 0) self.assertEqual(outcomes.get(timeout=5), ('ok', 1)) self.assertEqual(ingest_count.value, 1) def test_unknown_competing_publisher_identity_is_fatal(self): with tempfile.TemporaryDirectory() as temp_dir: directory = os.path.join(temp_dir, 'spool') spool = ResultSpool( directory, max_event_bytes=4096, max_events=20, max_total_bytes=65536, min_free_bytes=0, ) spool.write_event({'version': 1, 'scan_event_id': str(uuid.uuid4()), 'result': {}}) context = self.context() lease = context.Lock() self.assertTrue(lease.acquire(False)) try: db = SharedPublisherDB( lease, context.Event(), context.Event(), context.Value('i', 0), owner_state=False, ) with self.assertRaisesRegex(RuntimeError, 'unknown or unauthenticated'): console_runner.wait_for_result_spool_ready( spool, db, wait_seconds=0.01, diagnostic_interval=0.01, ) finally: lease.release() @staticmethod def empty_spool(): return SimpleNamespace( next_pending_event=lambda: None, assert_claims_allowed=lambda: True, ) @staticmethod def sequenced_db(acquisitions, states=None, release=True, state_error=None): class DB: conn = SimpleNamespace(is_postgres=True) def __init__(self): self.acquisitions = list(acquisitions) self.states = list(states or []) self.release_result = release self.state_error = state_error self.acquire_calls = 0 self.release_calls = 0 def require_runtime_safety_schema(self): return True def try_acquire_result_spool_publisher(self): self.acquire_calls += 1 return self.acquisitions.pop(0) def result_spool_publisher_state(self): if self.state_error: raise self.state_error return self.states.pop(0) def release_result_spool_publisher(self): self.release_calls += 1 return self.release_result return DB() def test_release_between_try_and_inspect_retries_acquire(self): db = self.sequenced_db( [False, True], states=[{'status': 'free', 'authenticated': True}], ) with mock.patch.object(console_runner.time, 'sleep') as sleep: self.assertEqual( console_runner.wait_for_result_spool_ready( self.empty_spool(), db, wait_seconds=0.1, diagnostic_interval=30, ), {}, ) self.assertEqual(db.acquire_calls, 2) self.assertEqual(db.release_calls, 1) sleep.assert_called_once_with(0.1) def test_holder_handoff_between_observations_retries_safely(self): owner_a = { 'status': 'held', 'authenticated': True, 'application_name': 'truf-source:github', 'holder_identity': '10:start-a', } owner_b = { 'status': 'held', 'authenticated': True, 'application_name': 'truf-source:gitlab', 'holder_identity': '11:start-b', } db = self.sequenced_db([False, False, True], states=[owner_a, owner_b]) with mock.patch.object(console_runner.time, 'sleep'): self.assertEqual( console_runner.wait_for_result_spool_ready( self.empty_spool(), db, wait_seconds=0.1, diagnostic_interval=30, ), {}, ) self.assertEqual(db.acquire_calls, 3) self.assertEqual(db.release_calls, 1) def test_owner_lookup_error_remains_fatal(self): db = self.sequenced_db([False], state_error=OSError('catalog lookup failed')) with self.assertRaisesRegex(OSError, 'catalog lookup failed'): console_runner.wait_for_result_spool_ready( self.empty_spool(), db, wait_seconds=0.1, diagnostic_interval=30, ) self.assertEqual(db.acquire_calls, 1) self.assertEqual(db.release_calls, 0) def test_shutdown_during_release_inspection_race_stops_wait(self): class StopDuringWait: def is_set(self): return False def wait(self, timeout): return True db = self.sequenced_db( [False], states=[{'status': 'free', 'authenticated': True}], ) with self.assertRaises(KeyboardInterrupt): console_runner.wait_for_result_spool_ready( self.empty_spool(), db, stop_event=StopDuringWait(), wait_seconds=0.1, diagnostic_interval=30, ) self.assertEqual(db.acquire_calls, 1) self.assertEqual(db.release_calls, 0) def test_exact_release_is_required_on_acquired_session(self): successful = self.sequenced_db([True], release=True) self.assertEqual( console_runner.drain_result_spool(self.empty_spool(), successful), {}, ) self.assertEqual((successful.acquire_calls, successful.release_calls), (1, 1)) uncertain = self.sequenced_db([True], release=False) with self.assertRaisesRegex(RuntimeError, 'lease release was not confirmed'): console_runner.drain_result_spool(self.empty_spool(), uncertain) self.assertEqual((uncertain.acquire_calls, uncertain.release_calls), (1, 1)) def test_scanner_db_tracks_exact_same_session_acquire_release(self): class Cursor: def __init__(self, row): self.row = row def fetchone(self): return self.row class Connection: is_postgres = True def __init__(self): self.closed = False self.calls = [] def execute(self, sql, params=None): self.calls.append(sql) if 'pg_try_advisory_lock' in sql: return Cursor({'acquired': True}) if 'pg_advisory_unlock' in sql: return Cursor({'released': True}) raise AssertionError(sql) def commit(self): return None def close(self): self.closed = True connection = Connection() db = ScannerDB.__new__(ScannerDB) db.conn = connection db._result_spool_publisher_held = False db.last_error = '' self.assertTrue(db.try_acquire_result_spool_publisher()) with self.assertRaisesRegex(RuntimeError, 'already held'): db.try_acquire_result_spool_publisher() self.assertTrue(db.release_result_spool_publisher()) self.assertFalse(db.release_result_spool_publisher()) self.assertEqual(sum('pg_try_advisory_lock' in sql for sql in connection.calls), 1) self.assertEqual(sum('pg_advisory_unlock' in sql for sql in connection.calls), 1) def test_uncertain_release_closes_session_to_prevent_lock_leak(self): class Cursor: def fetchone(self): return {'released': False} class Connection: is_postgres = True def __init__(self): self.closed = False def execute(self, sql, params=None): return Cursor() def commit(self): return None def close(self): self.closed = True connection = Connection() db = ScannerDB.__new__(ScannerDB) db.conn = connection db._result_spool_publisher_held = True db.last_error = '' self.assertFalse(db.release_result_spool_publisher()) self.assertTrue(connection.closed) self.assertIsNone(db.conn) self.assertFalse(db._result_spool_publisher_held) def test_twelve_sources_wait_for_capacity_then_handoff_without_loss(self): with tempfile.TemporaryDirectory() as temp_dir: directory = os.path.join(temp_dir, 'spool') spool = ResultSpool( directory, max_event_bytes=1024, max_events=100, max_total_bytes=2500, min_free_bytes=0, ) blocker = spool.reserve_claims('unrelated-owner', 1, 60, bytes_per_event=1024) context = self.context() start = context.Event() never_stop = context.Event() outcomes = context.Queue() processes = [ context.Process( target=capacity_worker, args=(directory, index, start, never_stop, outcomes), ) for index in range(12) ] for process in processes: process.start() start.set() time.sleep(0.5) self.assertTrue(os.path.exists(spool._reservation_path(blocker))) self.assertTrue(spool.release_reservation(blocker)) for process in processes: process.join(30) self.assertFalse(process.is_alive()) self.assertEqual(process.exitcode, 0) results = sorted(outcomes.get(timeout=5) for _ in processes) self.assertEqual(results, [('acquired', index) for index in range(12)]) self.assertEqual(list(os.scandir(spool.reservation_directory)), []) self.assertEqual(spool._pending_paths()[0], []) def test_capacity_wait_honors_shutdown_and_retains_owner_reservation(self): with tempfile.TemporaryDirectory() as temp_dir: directory = os.path.join(temp_dir, 'spool') spool = ResultSpool( directory, max_event_bytes=1024, max_events=100, max_total_bytes=2500, min_free_bytes=0, ) blocker = spool.reserve_claims('unrelated-owner', 1, 60, bytes_per_event=1024) context = self.context() start = context.Event() stop_event = context.Event() outcomes = context.Queue() process = context.Process( target=capacity_worker, args=(directory, 1, start, stop_event, outcomes), ) process.start() start.set() time.sleep(0.3) stop_event.set() process.join(15) self.assertFalse(process.is_alive()) self.assertEqual(process.exitcode, 0) self.assertEqual(outcomes.get(timeout=5), ('stopped', 1)) self.assertTrue(os.path.exists(spool._reservation_path(blocker))) self.assertTrue(spool.release_reservation(blocker)) def test_intrinsically_impossible_capacity_request_remains_fatal(self): with tempfile.TemporaryDirectory() as temp_dir: spool = ResultSpool( os.path.join(temp_dir, 'spool'), max_event_bytes=1024, max_events=100, max_total_bytes=2500, min_free_bytes=0, ) with self.assertRaises(SpoolCapacityError): console_runner.reserve_result_spool_claims( spool, CapacityProgressDB(), 'owner', 3, 60, wait_seconds=0.01, diagnostic_interval=0.01, ) def test_capacity_without_safe_progress_evidence_remains_fatal(self): class UnsafeProgressDB(CapacityProgressDB): def result_spool_reservation_progress(self, reservations): return {'safe_progress': False} with tempfile.TemporaryDirectory() as temp_dir: spool = ResultSpool( os.path.join(temp_dir, 'spool'), max_event_bytes=1024, max_events=100, max_total_bytes=2500, min_free_bytes=0, ) blocker = spool.reserve_claims('unrelated-owner', 1, 60, bytes_per_event=1024) with self.assertRaisesRegex(RuntimeError, 'unsafe or indeterminate'): console_runner.reserve_result_spool_claims( spool, UnsafeProgressDB(), 'owner', 1, 60, wait_seconds=0.01, diagnostic_interval=0.01, ) self.assertTrue(os.path.exists(spool._reservation_path(blocker))) self.assertTrue(spool.release_reservation(blocker)) def test_exact_stopped_owner_refund_releases_only_its_unused_reservation(self): with tempfile.TemporaryDirectory() as temp_dir, mock.patch.dict(os.environ, { 'SCANNER_DB_URL': '', 'DATABASE_URL': '', }): spool = ResultSpool( os.path.join(temp_dir, 'spool'), max_event_bytes=1024, max_events=20, max_total_bytes=65536, min_free_bytes=0, ) db = ScannerDB(db_path=os.path.join(temp_dir, 'scanner.db')) identity = current_process_identity() owner = f'source:{identity.pid}:cycle' owned = spool.reserve_claims(owner, 1, 60, bytes_per_event=1024) unrelated = spool.reserve_claims('other-owner', 1, 60, bytes_per_event=1024) target = 'https://github.com/example/stopped-owner' db.enqueue_targets('source', 'github', 'q', [target]) claims = db.claim_targets( 'source', 'github', 1, owner, 60, return_rows=True, claim_batch=owned, ) spool.bind_claims(owned, claims) candidates = spool.stopped_owner_unused_reservations('source', identity) self.assertEqual([row['reservation_id'] for row in candidates], [owned]) self.assertTrue(db.refund_stopped_result_spool_claims(candidates[0])) self.assertTrue(spool.release_stopped_owner_reservation(owned, 'source', identity)) row = db.target_queue_item('source', 'github', target) self.assertEqual((row['status'], row['attempts']), ('pending', 0)) self.assertFalse(os.path.exists(spool._reservation_path(owned))) self.assertTrue(os.path.exists(spool._reservation_path(unrelated))) self.assertTrue(spool.release_reservation(unrelated)) db.close() if __name__ == '__main__': unittest.main()