from datetime import datetime
from ccat_ops_db.models import OperationKind
def _operation_type_str(operation_type) -> str:
"""Render an operation identity to its on-the-wire Redis string.
Callers key off OperationKind members; Redis stores the identity as a plain
string. An OperationKind member's VALUE is byte-identical to the legacy
breadcrumb string, so the keys are unchanged on the wire. Using ``.value``
explicitly (NOT an f-string of the member) is mandatory: under Python 3.12
``str(member)`` / ``f"{member}"`` render the member NAME, which would build a
wrong key. Plain strings pass through unchanged.
"""
if isinstance(operation_type, OperationKind):
return operation_type.value
return operation_type
[docs]
class TaskStateManager:
"""Manager for tracking and recovering task states across all operation types."""
[docs]
def __init__(self, redis_client):
self.redis = redis_client
[docs]
def register_task(
self, task_id, operation_type, operation_id, additional_info=None, max_retries=3
):
"""
Register a task in Redis with its metadata.
Args:
task_id (str): Celery task ID
operation_type (str): Type of operation (transfer, archive, package, delete, verify)
operation_id (int): Database ID of the operation
additional_info (dict, optional): Additional context about the operation
max_retries (int, optional): Maximum retry count for this task
"""
key = f"task:{task_id}"
# Base data for all task types
data = {
"operation_type": operation_type,
"operation_id": str(operation_id),
"status": "RUNNING",
"start_time": datetime.now().isoformat(),
"heartbeat": datetime.now().isoformat(),
"retry_count": "0",
"max_retries": str(max_retries),
}
# Add additional info if provided
if additional_info:
for k, v in additional_info.items():
data[k] = str(v)
# Store in Redis with TTL
self.redis.hmset(key, data)
self.redis.expire(key, 86400 * 2) # 48 hour TTL
# Maintain indices for each operation type and ID
self.redis.sadd(f"running_tasks:{operation_type}", task_id)
self.redis.sadd(f"tasks_for_operation:{operation_type}:{operation_id}", task_id)
[docs]
def update_heartbeat(self, task_id):
"""Update task heartbeat to indicate it's still running."""
key = f"task:{task_id}"
if self.redis.exists(key):
self.redis.hset(key, "heartbeat", datetime.now().isoformat())
[docs]
def complete_task(self, task_id):
"""Mark task as completed and remove from tracking."""
key = f"task:{task_id}"
# Get operation info before deleting
if not self.redis.exists(key):
return
operation_type = self.redis.hget(key, "operation_type")
operation_id = self.redis.hget(key, "operation_id")
if operation_type and operation_id:
# Remove from indices
self.redis.srem(f"running_tasks:{operation_type}", task_id)
self.redis.srem(
f"tasks_for_operation:{operation_type}:{operation_id}", task_id
)
# Mark as completed
self.redis.hset(key, "status", "COMPLETED")
self.redis.hset(key, "end_time", datetime.now().isoformat())
# Keep completed task info for a while before deleting
self.redis.expire(key, 86400) # 24 hour TTL for completed tasks
[docs]
def fail_task(self, task_id, error_message, is_retryable=True):
"""
Mark task as failed.
Returns:
tuple: (can_retry, operation_type, operation_id)
"""
key = f"task:{task_id}"
if not self.redis.exists(key):
return False, None, None
# Get operation info
operation_type = self.redis.hget(key, "operation_type")
operation_id = self.redis.hget(key, "operation_id")
# Update status
self.redis.hset(key, "status", "FAILED")
self.redis.hset(key, "error", error_message)
self.redis.hset(key, "end_time", datetime.now().isoformat())
# Remove from running tasks
if operation_type:
self.redis.srem(f"running_tasks:{operation_type}", task_id)
if is_retryable:
# Increment retry count
retry_count = int(self.redis.hincrby(key, "retry_count", 1))
max_retries = int(self.redis.hget(key, "max_retries") or 3)
if retry_count <= max_retries:
# Ready to be retried
return True, operation_type, operation_id
else:
# Max retries exceeded
self.redis.hset(key, "status", "FAILED_PERMANENT")
return False, operation_type, operation_id
else:
# Not retryable
self.redis.hset(key, "status", "FAILED_PERMANENT")
return False, operation_type, operation_id
[docs]
def get_stalled_tasks(self, heartbeat_timeout=300):
"""
Find tasks that haven't updated their heartbeat recently.
Returns:
list: List of dicts with task information
"""
now = datetime.now()
stalled_tasks = []
# Get all running tasks
all_task_keys = self.redis.keys("task:*")
for task_key in all_task_keys:
task_data = self.redis.hgetall(task_key)
if task_data.get("status") != "RUNNING":
continue
# Check heartbeat
try:
last_heartbeat = datetime.fromisoformat(task_data.get("heartbeat"))
task_id = task_key.replace("task:", "")
if (now - last_heartbeat).total_seconds() > heartbeat_timeout:
stalled_tasks.append(
{
"task_id": task_id,
"operation_type": task_data.get("operation_type"),
"operation_id": task_data.get("operation_id"),
"last_heartbeat": last_heartbeat,
"stalled_for": (now - last_heartbeat).total_seconds(),
}
)
except (ValueError, TypeError):
# Invalid heartbeat format, consider stalled
stalled_tasks.append(
{
"task_id": task_key.replace("task:", ""),
"operation_type": task_data.get("operation_type"),
"operation_id": task_data.get("operation_id"),
"last_heartbeat": None,
"stalled_for": None,
}
)
return stalled_tasks
[docs]
def is_operation_alive(self, operation_type, operation_id, heartbeat_timeout=None):
"""
Report whether an operation still has a live worker behind it.
Reverse-lookup of the per-operation index: read the sibling task_ids in
``tasks_for_operation:{operation_type}:{operation_id}`` and inspect each
``task:{task_id}`` hash. The operation is alive only if at least one
sibling is ``RUNNING`` with a heartbeat fresher than ``heartbeat_timeout``.
This is the liveness probe the archive reconciler gates on (#161): the
primary dead-worker path is ``task_monitor`` heartbeat recovery, and this
lets the reconciler stand down whenever a registered task is still beating.
Args:
operation_type: Operation identity — an ``OperationKind`` member
(preferred) or its breadcrumb string (e.g.
``OperationKind.ARCHIVE`` / ``"long_term_archive"``). The member's
value is the on-the-wire key, so both build the same Redis key.
operation_id (int): Database ID of the operation.
heartbeat_timeout (int, optional): Max heartbeat age in seconds to
count a sibling as alive. Defaults to
``TASK_RECOVERY.heartbeat_timeout``.
Returns:
bool: True if a sibling task is RUNNING with a fresh heartbeat;
False if the index set is empty, the hashes are missing, or every
sibling is stale or non-RUNNING.
"""
operation_type = _operation_type_str(operation_type)
if heartbeat_timeout is None:
from .config.config import ccat_data_transfer_settings
heartbeat_timeout = (
ccat_data_transfer_settings.TASK_RECOVERY.heartbeat_timeout
)
now = datetime.now()
task_ids = self.redis.smembers(
f"tasks_for_operation:{operation_type}:{operation_id}"
)
for task_id in task_ids or []:
task_data = self.redis.hgetall(f"task:{task_id}")
if not task_data or task_data.get("status") != "RUNNING":
continue
# Mirror get_stalled_tasks() staleness logic: an unparseable or
# missing heartbeat is treated as not-alive (i.e. stale).
try:
last_heartbeat = datetime.fromisoformat(task_data.get("heartbeat"))
except (ValueError, TypeError):
continue
if (now - last_heartbeat).total_seconds() < heartbeat_timeout:
return True
return False
