sashatrask
75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
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()
|