Initial server source import
This commit is contained in:
@@ -0,0 +1,74 @@
|
||||
import glob
|
||||
import os
|
||||
import sys
|
||||
import sysconfig
|
||||
|
||||
|
||||
for path in glob.glob('/proc/[0-9]*/cmdline'):
|
||||
try:
|
||||
command = open(path, 'rb').read().replace(b'\0', b' ')
|
||||
except (FileNotFoundError, PermissionError, ProcessLookupError):
|
||||
continue
|
||||
if b'/opt/truf/app/result_ingester.py' in command or b'/opt/truf/app/jsonl_projector.py' in command:
|
||||
raise SystemExit('a pipeline worker process is still active')
|
||||
|
||||
sys.path.append(sysconfig.get_paths()['purelib'])
|
||||
sys.path.insert(0, '/opt/truf/app')
|
||||
|
||||
import psycopg
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
from db_backend import DatabaseConnection
|
||||
from scanner_db import (
|
||||
PIPELINE_ADVISORY_CLASS,
|
||||
PIPELINE_ADVISORY_OBJECTS,
|
||||
ScannerDB,
|
||||
)
|
||||
|
||||
|
||||
connection = psycopg.connect(
|
||||
dbname='truf',
|
||||
user='truf',
|
||||
host='/run/truf-postgres',
|
||||
row_factory=dict_row,
|
||||
options='-c search_path=public -c statement_timeout=30000 -c lock_timeout=5000',
|
||||
)
|
||||
database = ScannerDB(enabled=False)
|
||||
database.conn = DatabaseConnection('postgres', connection, application_schema='public')
|
||||
released = 0
|
||||
try:
|
||||
rows = database.conn.execute(
|
||||
"SELECT worker_name, generation, lease_token FROM pipeline_leases "
|
||||
"WHERE state NOT IN ('released','failed') ORDER BY worker_name"
|
||||
).fetchall()
|
||||
expected = {'result_ingester', 'jsonl_projector'}
|
||||
if not {row['worker_name'] for row in rows} <= expected:
|
||||
raise RuntimeError('an unexpected pipeline lease is active')
|
||||
for row in rows:
|
||||
name = row['worker_name']
|
||||
advisory = PIPELINE_ADVISORY_OBJECTS[name]
|
||||
locked = database.conn.execute(
|
||||
'SELECT pg_try_advisory_lock(?, ?) AS acquired',
|
||||
(PIPELINE_ADVISORY_CLASS, advisory),
|
||||
).fetchone()
|
||||
if not locked or not locked['acquired']:
|
||||
raise RuntimeError('a stopped pipeline worker still owns its advisory lock')
|
||||
database._pipeline_advisory_held.add(name)
|
||||
if not database.release_pipeline_lease(
|
||||
name,
|
||||
row['generation'],
|
||||
row['lease_token'],
|
||||
state='released',
|
||||
error='deployment_quiesce',
|
||||
):
|
||||
raise RuntimeError('pipeline lease identity changed during release')
|
||||
released += 1
|
||||
remaining = database.conn.execute(
|
||||
"SELECT count(*) AS count FROM pipeline_leases "
|
||||
"WHERE state NOT IN ('released','failed')"
|
||||
).fetchone()['count']
|
||||
if remaining:
|
||||
raise RuntimeError('pipeline lease reconciliation is incomplete')
|
||||
print(f'released_pipeline_leases={released}')
|
||||
finally:
|
||||
database.close()
|
||||
Reference in New Issue
Block a user