init

agent/main.py

@@ -0,0 +1,101 @@
+"""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()