Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,251 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user