ms
102 lines
5.5 KiB
Python
102 lines
5.5 KiB
Python
"""Outbound-only host metrics and optional Celery events collector."""
|
|
import json,logging,os,signal,socket,threading,time,uuid
|
|
from datetime import datetime,timezone
|
|
from pathlib import Path
|
|
import psutil,requests
|
|
from celery import Celery
|
|
from dotenv import load_dotenv
|
|
from .outbox import Outbox
|
|
log=logging.getLogger('dilib.agent')
|
|
STOP=threading.Event()
|
|
ALLOWED_TASK={'uuid','name','hostname','runtime','retries'}
|
|
ALLOWED_WORKER={'hostname','active','processed'}
|
|
KINDS={'worker-online','worker-offline','worker-heartbeat','task-sent','task-received','task-started','task-succeeded','task-failed','task-retried','task-revoked','task-rejected'}
|
|
def envelope(kind,payload,at=None):
|
|
return {'id':str(uuid.uuid4()),'kind':kind,'occurred_at':(at or datetime.now(timezone.utc)).isoformat(),'payload':payload}
|
|
def safe_event(e):
|
|
kind=e.get('type')
|
|
if kind not in KINDS:return None
|
|
keep=ALLOWED_TASK if kind.startswith('task-') else ALLOWED_WORKER
|
|
p={k:e[k] for k in keep if k in e and e[k] is not None}
|
|
if kind.startswith('task-') and not p.get('uuid'):return None
|
|
if kind.startswith('worker-') and not p.get('hostname'):return None
|
|
return envelope(kind,p,datetime.fromtimestamp(float(e.get('timestamp',time.time())),timezone.utc))
|
|
def collect(outbox,state):
|
|
app=Celery('dilib-agent',broker=os.environ['AGENT_BROKER_URL'])
|
|
app.conf.update(accept_content=['json'],event_queue_ttl=300,event_queue_expires=86400)
|
|
while not STOP.is_set():
|
|
try:
|
|
with app.connection() as conn:
|
|
conn.ensure_connection(max_retries=3)
|
|
def handler(event):
|
|
state['last_event']=time.monotonic()
|
|
cleaned=safe_event(event)
|
|
if cleaned:outbox.add(cleaned)
|
|
receiver=app.events.Receiver(conn,handlers={'*':handler},node_id=os.getenv('AGENT_COLLECTOR_ID','dilib-monitor'),accept=['json'])
|
|
with receiver.consumer_context():
|
|
state['connected']=True
|
|
while not STOP.is_set():
|
|
try:conn.drain_events(timeout=2)
|
|
except socket.timeout:pass
|
|
except Exception as exc:
|
|
state['connected']=False
|
|
log.warning('Collector reconnect required (%s)',type(exc).__name__)
|
|
STOP.wait(5)
|
|
finally:state['connected']=False
|
|
|
|
def send_batch(session,outbox,url,token,verify=True):
|
|
batch=outbox.batch()
|
|
if not batch:return 0
|
|
r=session.post(url,json={'events':[e for _,e in batch]},headers={'Authorization':'Agent '+token},timeout=(5,30),verify=verify)
|
|
r.raise_for_status()
|
|
if r.json().get('received')!=len(batch): raise RuntimeError('Server acknowledgement mismatch')
|
|
outbox.ack([seq for seq,_ in batch]);return len(batch)
|
|
def sender(outbox):
|
|
session=requests.Session();backoff=1
|
|
while not STOP.is_set():
|
|
try:
|
|
n=send_batch(session,outbox,os.environ['DASHBOARD_URL'],os.environ['AGENT_TOKEN'],os.getenv('AGENT_CA_BUNDLE') or True)
|
|
backoff=1
|
|
if not n:STOP.wait(1)
|
|
except Exception as exc:
|
|
# Tokens, broker URLs, payloads and raw errors are never logged.
|
|
code=getattr(getattr(exc,'response',None),'status_code',None)
|
|
log.warning('Upload failed (%s, HTTP %s); data remains in outbox',type(exc).__name__,code)
|
|
STOP.wait(backoff);backoff=min(backoff*2,60)
|
|
def check_url(url,auth=None):
|
|
try:
|
|
r=requests.get(url,auth=auth,timeout=(2,3));return 200<=r.status_code<300
|
|
except requests.RequestException:return False
|
|
def main():
|
|
load_dotenv(os.getenv('AGENT_ENV_FILE','.env'))
|
|
logging.basicConfig(level=logging.INFO,format='%(asctime)s %(levelname)s %(message)s')
|
|
for key in ('AGENT_TOKEN','DASHBOARD_URL'):
|
|
if not os.getenv(key):raise RuntimeError(key+' missing')
|
|
if not os.environ['DASHBOARD_URL'].startswith('https://') and os.getenv('AGENT_ALLOW_HTTP')!='1':
|
|
raise RuntimeError('HTTPS required (AGENT_ALLOW_HTTP=1 only for local testing)')
|
|
os.umask(0o077)
|
|
outbox=Outbox(os.getenv('AGENT_SPOOL','var/outbox.sqlite3'),int(os.getenv('AGENT_MAX_PENDING','1000000')))
|
|
for sig in (signal.SIGTERM,signal.SIGINT):signal.signal(sig,lambda *_:STOP.set())
|
|
threads=[threading.Thread(target=sender,args=(outbox,),daemon=True)]
|
|
state={'connected':False};collecting=os.getenv('AGENT_COLLECT_EVENTS','0')=='1'
|
|
if collecting:threads.append(threading.Thread(target=collect,args=(outbox,state),daemon=True))
|
|
for t in threads:t.start()
|
|
urls=json.loads(os.getenv('AGENT_HEALTH_URLS','{}'))
|
|
boot=str(int(psutil.boot_time()))
|
|
psutil.cpu_percent()
|
|
try:
|
|
while not STOP.is_set():
|
|
metrics={'cpu':psutil.cpu_percent(),'ram':psutil.virtual_memory().percent,'disk':psutil.disk_usage('/').percent,'boot_id':boot,'services':{name:check_url(url) for name,url in urls.items()},'outbox':outbox.count()}
|
|
outbox.add(envelope('metric',metrics))
|
|
if collecting:outbox.add(envelope('collector-status',{'online':state['connected']}))
|
|
if os.getenv('AGENT_FLOWER_URL'):
|
|
auth=(os.getenv('AGENT_FLOWER_USER',''),os.getenv('AGENT_FLOWER_PASSWORD',''))
|
|
ok=check_url(os.environ['AGENT_FLOWER_URL'].rstrip('/')+'/api/workers',auth)
|
|
outbox.add(envelope('flower-status',{'online':ok}))
|
|
STOP.wait(max(1,float(os.getenv('AGENT_INTERVAL','5'))))
|
|
finally:
|
|
STOP.set()
|
|
for t in threads:t.join(timeout=40)
|
|
if not any(t.is_alive() for t in threads):outbox.close()
|
|
if __name__=='__main__':main()
|