mirror of
https://github.com/yogeshojha/rengine.git
synced 2026-10-20 23:31:31 +02:00
Yogesh Ojha
a700bbed03
Launches past the limit stay pending and start in launch order when a slot frees, so each running scan gets full speed and memory stays bounded. - run_scan admits a pending scan under an advisory lock: running plus waiting scans ahead of it must be under the limit - finalize, cancel, pause, delete and a limit change start the next scans; the reaper re-sends only scans whose turn has come - automatic limit: live worker-scans stage slots over the widest step, so --scale raises it; Settings sets a fixed limit from 1 to 50 - watch probes skip the queue, resume is never held - queued scans show how many are ahead
65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
"""Advisory locks for instance-wide work that must not run twice at once."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import zlib
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from typing import TYPE_CHECKING
|
|
|
|
from sqlalchemy import text
|
|
|
|
if TYPE_CHECKING:
|
|
from sqlalchemy.orm import Session
|
|
|
|
# bounty_platform() spans BOUNTY_SYNC .. BOUNTY_SYNC + 0xFFFF
|
|
BOUNTY_SYNC = 0x624F0001
|
|
BOUNTY_FEED = 0x624F0002
|
|
# bounty_reports() spans BOUNTY_REPORTS .. BOUNTY_REPORTS + 0xFFFF
|
|
BOUNTY_REPORTS = 0x62520001
|
|
IP_RANGES = 0x624E0001
|
|
# secret_mining() spans SECRET_MINING .. SECRET_MINING + 0xFFFF
|
|
SECRET_MINING = 0x53450001
|
|
SOFTWARE_BACKFILL = 0x53570001
|
|
# software_match() spans SOFTWARE_MATCH .. SOFTWARE_MATCH + 0xFFFF
|
|
SOFTWARE_MATCH = 0x53580001
|
|
NEW_CHECKS_SWEEP = 0x4E430001
|
|
SCAN_ADMISSION = 0x53410001
|
|
|
|
|
|
def bounty_platform(platform: str) -> int:
|
|
"""One lock per platform."""
|
|
return BOUNTY_SYNC + (zlib.crc32(platform.encode()) & 0xFFFF)
|
|
|
|
|
|
def bounty_reports(platform: str) -> int:
|
|
"""One lock per platform."""
|
|
return BOUNTY_REPORTS + (zlib.crc32(platform.encode()) & 0xFFFF)
|
|
|
|
|
|
def secret_mining(scan_id: object) -> int:
|
|
"""One lock per scan: the stage, the backfill and a re-mine never overlap."""
|
|
return SECRET_MINING + (zlib.crc32(str(scan_id).encode()) & 0xFFFF)
|
|
|
|
|
|
def software_match(scan_id: object) -> int:
|
|
"""One lock per scan."""
|
|
return SOFTWARE_MATCH + (zlib.crc32(str(scan_id).encode()) & 0xFFFF)
|
|
|
|
|
|
@contextmanager
|
|
def sync_lock(session: Session, key: int) -> Iterator[bool]:
|
|
"""Session-level advisory lock. Yields False when another session holds it."""
|
|
held = bool(
|
|
session.execute(
|
|
text("SELECT pg_try_advisory_lock(:key)"), {"key": key}
|
|
).scalar_one()
|
|
)
|
|
try:
|
|
yield held
|
|
finally:
|
|
if held:
|
|
session.rollback()
|
|
session.execute(text("SELECT pg_advisory_unlock(:key)"), {"key": key})
|
|
session.commit()
|