sashatrask
252 lines
9.4 KiB
Python
252 lines
9.4 KiB
Python
import concurrent.futures
|
|
import os
|
|
from pathlib import Path
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
APP_DIR = ROOT / 'app'
|
|
sys.path.insert(0, str(APP_DIR))
|
|
|
|
import scanner
|
|
|
|
|
|
class ScanSlotReleaseReliabilityTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
scanner.initialize_scanner_runtime(preflight_complete=True, register_cleanup=False)
|
|
|
|
def setUp(self):
|
|
self._reset_fatal_state()
|
|
self.addCleanup(self._reset_fatal_state)
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp_dir.cleanup)
|
|
self.db_path = os.path.join(self.temp_dir.name, 'scan-limiter.sqlite')
|
|
patchers = (
|
|
mock.patch.object(scanner.scan_config, 'max_active_scans', 1),
|
|
mock.patch.object(scanner.scan_config, 'scan_limiter_db', self.db_path),
|
|
mock.patch.object(scanner.scan_config, 'scan_slot_wait_sec', 0.01),
|
|
mock.patch.object(scanner.scan_config, 'scan_slot_heartbeat_sec', 1.0, create=True),
|
|
mock.patch.object(scanner.ScanSlotLease, 'MIN_HEARTBEAT_INTERVAL_SEC', 0.01),
|
|
mock.patch.object(scanner.ScanSlotLease, 'RELEASE_PENDING_INTERVAL_SEC', 0.02),
|
|
mock.patch.object(scanner.ScanSlotLease, 'DB_RETRY_DELAY_SEC', 0),
|
|
mock.patch.object(scanner.ScanSlotLease, 'HEARTBEAT_JOIN_TIMEOUT_SEC', 1.0),
|
|
)
|
|
for patcher in patchers:
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
|
|
@staticmethod
|
|
def _reset_fatal_state():
|
|
with scanner._scan_slot_fatal_lock:
|
|
scanner._scan_slot_fatal_detail = None
|
|
scanner._scan_slot_fatal_event.clear()
|
|
|
|
def _acquire(self):
|
|
lease = scanner.acquire_scan_slot(['trufflehog', 'fixture'], timeout_sec=1)
|
|
self.addCleanup(self._cleanup_lease, lease)
|
|
return lease
|
|
|
|
def _cleanup_lease(self, lease):
|
|
if not lease.released:
|
|
lease.release()
|
|
lease._heartbeat_wake.set()
|
|
self._wait_until(lambda: lease.released, timeout=1)
|
|
if not lease.released:
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute(
|
|
'''DELETE FROM scan_slots
|
|
WHERE slot_id = ? AND owner_pid = ? AND owner_thread = ?''',
|
|
(lease.slot_id, lease.owner_pid, lease.owner_thread),
|
|
)
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
lease._complete_release()
|
|
lease._join_heartbeat()
|
|
|
|
def _row_count(self):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
return conn.execute('SELECT COUNT(*) FROM scan_slots').fetchone()[0]
|
|
finally:
|
|
conn.close()
|
|
|
|
@staticmethod
|
|
def _wait_until(predicate, timeout=2):
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if predicate():
|
|
return True
|
|
time.sleep(0.01)
|
|
return predicate()
|
|
|
|
def test_non_lock_operational_error_is_retried_by_release_caller(self):
|
|
lease = self._acquire()
|
|
delete_once = lease._delete_slot_once
|
|
calls = 0
|
|
|
|
def fail_once():
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls == 1:
|
|
raise sqlite3.OperationalError('disk I/O error')
|
|
return delete_once()
|
|
|
|
with mock.patch.object(lease, '_delete_slot_once', side_effect=fail_once):
|
|
lease.release()
|
|
|
|
self.assertEqual(calls, 2)
|
|
self.assertTrue(lease.released)
|
|
self.assertFalse(lease.release_pending)
|
|
self.assertEqual(self._row_count(), 0)
|
|
self.assertFalse(lease.heartbeat_thread.is_alive())
|
|
|
|
def test_persistent_db_failure_recovers_in_release_pending_mode(self):
|
|
lease = self._acquire()
|
|
connect = lease._connect
|
|
db_restored = threading.Event()
|
|
|
|
def unavailable_then_restored():
|
|
if not db_restored.is_set():
|
|
raise OSError('scan-slot volume unavailable')
|
|
return connect()
|
|
|
|
with mock.patch.object(lease, '_connect', side_effect=unavailable_then_restored):
|
|
with self.assertLogs(scanner.logger, level='ERROR') as captured:
|
|
lease.release()
|
|
|
|
self.assertTrue(any('background retries will continue' in line for line in captured.output))
|
|
self.assertTrue(lease.release_pending)
|
|
self.assertFalse(lease.released)
|
|
self.assertTrue(lease.heartbeat_thread.is_alive())
|
|
self.assertEqual(self._row_count(), 1)
|
|
|
|
time.sleep(0.08)
|
|
self.assertFalse(lease.released)
|
|
self.assertEqual(self._row_count(), 1)
|
|
|
|
db_restored.set()
|
|
lease._heartbeat_wake.set()
|
|
self.assertTrue(self._wait_until(lambda: lease.released))
|
|
|
|
self.assertEqual(self._row_count(), 0)
|
|
next_lease = self._acquire()
|
|
self.assertNotEqual(next_lease.slot_id, lease.slot_id)
|
|
next_lease.release()
|
|
self.assertTrue(next_lease.released)
|
|
self.assertEqual(self._row_count(), 0)
|
|
|
|
def test_concurrent_heartbeat_and_release_are_serialized_and_exact(self):
|
|
lease = self._acquire()
|
|
identity = scanner.current_process_identity()
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute(
|
|
'''INSERT INTO scan_slots(
|
|
slot_id, owner_pid, owner_thread, owner_source, owner_creation_time,
|
|
owner_executable, acquired_at, updated_at
|
|
) VALUES (?, ?, ?, 'test', ?, ?, ?, ?)''',
|
|
(
|
|
'other-slot', os.getpid(), threading.get_ident(),
|
|
identity.creation_time, identity.executable, time.time(), time.time(),
|
|
),
|
|
)
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
connect = lease._connect
|
|
heartbeat_inside_db_operation = threading.Event()
|
|
allow_heartbeat = threading.Event()
|
|
blocked = threading.Event()
|
|
|
|
def block_one_heartbeat_connection():
|
|
if threading.current_thread() is lease.heartbeat_thread and not blocked.is_set():
|
|
blocked.set()
|
|
heartbeat_inside_db_operation.set()
|
|
if not allow_heartbeat.wait(2):
|
|
raise OSError('test heartbeat was not released')
|
|
return connect()
|
|
|
|
with mock.patch.object(lease, '_connect', side_effect=block_one_heartbeat_connection):
|
|
lease._heartbeat_wake.set()
|
|
self.assertTrue(heartbeat_inside_db_operation.wait(1))
|
|
release_thread = threading.Thread(target=lease.release)
|
|
release_thread.start()
|
|
time.sleep(0.05)
|
|
self.assertTrue(release_thread.is_alive())
|
|
allow_heartbeat.set()
|
|
release_thread.join(timeout=2)
|
|
|
|
self.assertFalse(release_thread.is_alive())
|
|
self.assertTrue(lease.released)
|
|
self.assertFalse(lease.heartbeat_thread.is_alive())
|
|
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
remaining = {row[0] for row in conn.execute('SELECT slot_id FROM scan_slots')}
|
|
self.assertEqual(remaining, {'other-slot'})
|
|
conn.execute('DELETE FROM scan_slots WHERE slot_id = ?', ('other-slot',))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
self.assertEqual(self._row_count(), 0)
|
|
|
|
def test_heartbeat_thread_start_failure_rolls_back_row_and_next_acquire_succeeds(self):
|
|
original_start = scanner.threading.Thread.start
|
|
failed = False
|
|
|
|
def fail_first_heartbeat(thread):
|
|
nonlocal failed
|
|
if not failed and thread.name.startswith('scan-slot-heartbeat-'):
|
|
failed = True
|
|
raise RuntimeError('injected thread start failure')
|
|
return original_start(thread)
|
|
|
|
with mock.patch.object(scanner.threading.Thread, 'start', new=fail_first_heartbeat):
|
|
with self.assertRaisesRegex(RuntimeError, 'injected thread start failure'):
|
|
scanner.acquire_scan_slot(['trufflehog', 'first'], timeout_sec=1)
|
|
|
|
self.assertTrue(failed)
|
|
self.assertEqual(self._row_count(), 0)
|
|
|
|
lease = self._acquire()
|
|
lease.release()
|
|
self.assertTrue(lease.released)
|
|
self.assertEqual(self._row_count(), 0)
|
|
|
|
def test_unconfirmed_start_failure_rollback_raises_fatal_fail_closed_error(self):
|
|
with mock.patch.object(
|
|
scanner.ScanSlotLease,
|
|
'start_heartbeat',
|
|
side_effect=RuntimeError('heartbeat unavailable'),
|
|
), mock.patch.object(
|
|
scanner.ScanSlotLease,
|
|
'_delete_slot_once',
|
|
side_effect=OSError('limiter unavailable'),
|
|
), self.assertLogs(scanner.logger, level='CRITICAL') as captured:
|
|
with self.assertRaisesRegex(scanner.ScanSlotFatalError, 'FATAL'):
|
|
scanner.acquire_scan_slot(['trufflehog', 'fatal-fixture'], timeout_sec=1)
|
|
|
|
self.assertEqual(self._row_count(), 1)
|
|
self.assertTrue(scanner._scan_slot_fatal_event.is_set())
|
|
self.assertTrue(any('capacity remains fail-closed' in line for line in captured.output))
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute('DELETE FROM scan_slots')
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|