invoice_az/invoice_az/background_tasks.py

# ======= GLOBAL BACKGROUND TASK REGISTRY =======
#
# A lightweight, per-user registry of in-flight background jobs, backed by the
# Redis cache. It powers the desk-wide "Background Tasks" widget (bottom-right):
# every long-running worker writes a progress snapshot here, the widget reads it
# back and can request cancellation.
#
# Storage: a per-user Redis HASH, one field per job_id. Each field is written/
# deleted independently (HSET/HDEL are atomic per field), so MULTIPLE concurrent
# jobs for the same user never clobber each other. (An earlier version stored the
# whole registry as one pickled dict and did read-modify-write, which raced under
# concurrent loads — phantom/flickering rows. The hash fixes that.)
#
# Design notes:
#   * Keyed by USER. `get_active_background_tasks()` / `cancel_background_task()`
#     derive the user from `frappe.session.user` — the client never passes a user,
#     so one user can neither see nor cancel another user's jobs.
#   * `hgetall` reads straight from Redis (no frappe.local cache), so the widget
#     always sees a fresh, consistent snapshot. The registry is the source of
#     truth (survives page reloads); realtime pings are just "re-fetch" hints.
#   * Cancellation is best-effort: the worker checks `is_cancel_requested()`
#     between steps and bails out after finishing the current step.

import uuid

import frappe

REALTIME_EVENT = "background_tasks_changed"

_REGISTRY_TTL = 7200  # the per-user hash self-expires after 2h (crash safety net)
_CANCEL_TTL = 3600


# ======= INTERNAL HELPERS =======

def _registry_key(user):
    return f"invoice_az:bg_tasks:{user}"


def _cancel_key(job_id):
    return f"invoice_az:bg_cancel:{job_id}"


def _touch_ttl(user):
    """Keep the per-user hash key alive while jobs are running."""
    try:
        cache = frappe.cache()
        cache.expire(cache.make_key(_registry_key(user)), _REGISTRY_TTL)
    except Exception:
        pass  # TTL is only a crash safety net; finish_task removes tasks normally


def _all_tasks(user):
    """Fresh-from-Redis list of the user's task snapshots (hgetall bypasses
    frappe.local cache)."""
    return list((frappe.cache().hgetall(_registry_key(user)) or {}).values())


def _get_task(user, job_id):
    return frappe.cache().hget(_registry_key(user), job_id)


def _put_task(user, job_id, task):
    frappe.cache().hset(_registry_key(user), job_id, task)
    _touch_ttl(user)


def _del_task(user, job_id):
    frappe.cache().hdel(_registry_key(user), job_id)


def _ping(user, payload=None):
    """Hint the widget to re-fetch. `payload` may carry a finished task snapshot
    so the widget can flash its final state before dropping it."""
    frappe.publish_realtime(REALTIME_EVENT, payload or {}, user=user)


# ======= PUBLIC WORKER-SIDE API =======

def new_job_id():
    return uuid.uuid4().hex


def register_task(job_id, user, feature, label, total=0, cancellable=True):
    """Add a task to the user's registry in the `running` state."""
    if not job_id:
        return
    _put_task(user, job_id, {
        "job_id": job_id,
        "feature": feature,
        "label": label,
        "current": 0,
        "total": total or 0,
        "percent": 0,
        "status": "running",  # running | cancelling | done | error | cancelled
        "cancellable": bool(cancellable),
        "message": "",
    })
    _ping(user)


def update_task(job_id, user, current=None, total=None, percent=None, message=None):
    """Update progress for a running task. Recomputes `percent` from
    current/total unless an explicit `percent` is given."""
    if not job_id:
        return
    task = _get_task(user, job_id)
    if not task:
        return
    if total is not None:
        task["total"] = total or 0
    if current is not None:
        task["current"] = current
    if message is not None:
        task["message"] = message
    if percent is not None:
        task["percent"] = max(0, min(100, int(percent)))
    elif task.get("total"):
        task["percent"] = max(0, min(100, int(round(task["current"] / task["total"] * 100))))
    _put_task(user, job_id, task)
    _ping(user)


def finish_task(job_id, user, status="done", message=None):
    """Remove the task from the registry and flash its final state to the widget.

    Removing it immediately keeps the poll-based fallback consistent: a task that
    is gone from the registry is simply gone from the widget's list."""
    if not job_id:
        return
    task = _get_task(user, job_id)
    _del_task(user, job_id)
    frappe.cache().delete_value(_cancel_key(job_id))
    if task:
        task["status"] = status
        if status == "done":
            task["percent"] = 100
            task["current"] = task.get("total") or task.get("current") or 0
            task["message"] = message if message is not None else "Done"
        elif message is not None:
            task["message"] = message
        _ping(user, {"finished": task})
    else:
        _ping(user)


def request_cancel(job_id, user):
    """Set the cancel flag for a task owned by `user`. Returns False if the task
    is not in this user's registry (so cross-user cancellation is impossible)."""
    task = _get_task(user, job_id)
    if not task:
        return False
    frappe.cache().set_value(_cancel_key(job_id), 1, expires_in_sec=_CANCEL_TTL)
    task["status"] = "cancelling"
    task["message"] = "Cancellation requested"
    _put_task(user, job_id, task)
    _ping(user)
    return True


def is_cancel_requested(job_id):
    """Fresh-from-Redis check the worker calls between steps.

    CRITICAL: plain get_value() reads/writes frappe.local.cache, which pins the
    first result (None) for the whole job — so a cancel flag set later by the web
    request would never be seen. expires=True + use_local_cache=False forces a
    real Redis read every call."""
    return bool(
        frappe.cache().get_value(_cancel_key(job_id), expires=True, use_local_cache=False)
    )


# ======= WHITELISTED CLIENT-SIDE API =======

@frappe.whitelist()
def get_active_background_tasks():
    """Return the current user's in-flight tasks. Source of truth for the widget."""
    return _all_tasks(frappe.session.user)


@frappe.whitelist()
def cancel_background_task(job_id):
    """Best-effort cancel of one of the current user's tasks."""
    user = frappe.session.user
    if not request_cancel(job_id, user):
        return {"success": False, "message": frappe._("Task not found or already finished.")}
    return {
        "success": True,
        "message": frappe._("Cancellation requested. The job will stop after the current step."),
    }