The only thing this doesnt handle is when a client goes into offline due to a transient link failure but it comes back online. We want to avoid updating the DB every time a pong is RX'd.. Perhaps we do a seocndary corss checks with offline hosts?
80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
# app/tasks/reconcile_online_workers.py
|
|
"""
|
|
Celery task: reconcile worker status with Redis heartbeats.
|
|
|
|
This task scans all workers marked as "online" in the database.
|
|
If a worker has not sent a heartbeat within the timeout window,
|
|
it will be marked as "offline" via API or direct update.
|
|
|
|
This protects against undetected disconnects, e.g., when a WebSocket
|
|
server crashes or workers silently die without a disconnect event.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
import time
|
|
|
|
import redis
|
|
from app import celery_app as celery, db, logger
|
|
from app.models.models import WorkloadHost # Replace with your actual Worker model
|
|
from app import app
|
|
from datetime import datetime, timedelta
|
|
|
|
API_URL = "http://localhost:5000/api/internal/workers" # Adjust to your environment
|
|
redisclient = redis.from_url(app.config["REDIS_URL"])
|
|
|
|
|
|
@celery.task(name="tasks.reconcile_online_workers", bind=True)
|
|
def reconcile_online_workers(self) -> None:
|
|
"""
|
|
Check Redis for last heartbeat timestamps and reconcile online status of workers.
|
|
|
|
This task ensures that any worker marked 'online' in the database
|
|
but missing a recent heartbeat is transitioned to 'offline'.
|
|
"""
|
|
logger.info("Starting worker heartbeat reconciliation task")
|
|
|
|
now = int(time.time())
|
|
|
|
# Step 1: Get all workers marked as online
|
|
# online_workers = WorkloadHost.query.filter_by(_status="online",deleted=0).all()
|
|
|
|
cutoff = datetime.utcnow() - timedelta(seconds=10)
|
|
|
|
online_workers = WorkloadHost.query.filter(
|
|
WorkloadHost._status == "online",
|
|
WorkloadHost.deleted == 0,
|
|
WorkloadHost.updated_at < cutoff
|
|
).all()
|
|
logger.info(f"Found {len(online_workers)} workers marked as online")
|
|
|
|
stale_workers = []
|
|
|
|
worker: WorkloadHost # Strongly type so that VSCode IDE can autocomplete
|
|
|
|
for worker in online_workers:
|
|
logger.debug(f"Checking worker {worker}")
|
|
redis_key = f"ws_liveness:{worker.id}"
|
|
last_seen = redisclient.get(redis_key)
|
|
|
|
if not last_seen:
|
|
logger.warning(f"No heartbeat found for WorkloadHost {worker.id}")
|
|
is_stale = True
|
|
else:
|
|
try:
|
|
last_seen = int(last_seen)
|
|
is_stale = (now - last_seen) > app.config["PING_HEARTBEAT_TIMEOUT_SECONDS"]
|
|
except Exception:
|
|
logger.exception(f"Invalid heartbeat data for WorkloadHost {worker.id}")
|
|
is_stale = True
|
|
|
|
if is_stale:
|
|
logger.warning(f"WorkloadHost {worker.id} is stale — marking offline")
|
|
stale_workers.append(worker.id)
|
|
|
|
|
|
# Set worker offline
|
|
worker.set_status("offline")
|
|
db.session.commit()
|
|
|
|
logger.info(f"Finished reconciliation. Marked {len(stale_workers)} WorkloadHosts offline")
|