Fault-tolerant work queue

import heapq
from collections import deque


class WorkQueue:
    """Tasks are QUEUED -> LEASED -> DONE. A lease that runs out, or a failure,
    sends the task back to QUEUED, or to DEAD once its attempts are used up."""

    def __init__(self, lease_time, max_attempts):
        self.lease_time = lease_time
        self.max_attempts = max_attempts
        self.state = {}           # task -> "queued" | "leased" | "done" | "dead"
        self.attempts = {}        # task -> how many times it was claimed
        self.ready = deque()      # queued tasks, oldest first
        self.leases = {}          # live lease_id -> task
        self.deadlines = []       # heap of (deadline, lease_id); may hold stale entries
        self.next_id = 1

    def _expire(self, now):
        # Runs first in every call: leases that ran out end before anything else.
        while self.deadlines and self.deadlines[0][0] <= now:
            deadline, lease_id = heapq.heappop(self.deadlines)
            if lease_id in self.leases:            # skip stale entries (lazy deletion)
                self._release(lease_id)

    def _release(self, lease_id):
        task = self.leases.pop(lease_id)
        if self.attempts[task] >= self.max_attempts:
            self.state[task] = "dead"
        else:
            self.state[task] = "queued"
            self.ready.append(task)                # back of the line

    def _grant(self, now, task):
        lease_id, self.next_id = self.next_id, self.next_id + 1
        self.leases[lease_id] = task
        heapq.heappush(self.deadlines, (now + self.lease_time, lease_id))
        return lease_id

    def add(self, now, task):
        self._expire(now)
        if task in self.state:
            return False
        self.state[task], self.attempts[task] = "queued", 0
        self.ready.append(task)
        return True

    def claim(self, now):
        self._expire(now)
        if not self.ready:
            return None
        task = self.ready.popleft()
        self.state[task] = "leased"
        self.attempts[task] += 1
        return task, self._grant(now, task)

    def heartbeat(self, now, lease_id):
        self._expire(now)
        task = self.leases.pop(lease_id, None)
        if task is None:
            return None                            # unknown, expired or finished
        return self._grant(now, task)              # a new id: the old heap entry goes stale

    def complete(self, now, lease_id):
        self._expire(now)
        task = self.leases.pop(lease_id, None)
        if task is None:
            return False
        self.state[task] = "done"
        return True

    def fail(self, now, lease_id):
        self._expire(now)
        if lease_id not in self.leases:
            return False
        self._release(lease_id)
        return True

    def status(self, now, task):
        self._expire(now)
        return self.state.get(task, "unknown")

#include <deque>
#include <functional>
#include <optional>
#include <queue>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
using namespace std;

// Tasks are queued -> leased -> done. A lease that runs out, or a failure,
// sends the task back to queued, or to dead once its attempts are used up.
class WorkQueue {
    long long lease_time;
    int max_attempts;
    unordered_map<string, string> state;          // task -> "queued" | "leased" | "done" | "dead"
    unordered_map<string, int> attempts;          // task -> how many times it was claimed
    deque<string> ready;                          // queued tasks, oldest first
    unordered_map<long long, string> leases;      // live lease_id -> task
    using Entry = pair<long long, long long>;     // (deadline, lease_id); may be stale
    priority_queue<Entry, vector<Entry>, greater<>> deadlines;
    long long next_id = 1;

    // Runs first in every call: leases that ran out end before anything else.
    void expire(long long now) {
        while (!deadlines.empty() && deadlines.top().first <= now) {
            long long lease_id = deadlines.top().second;
            deadlines.pop();
            if (leases.count(lease_id)) release(lease_id);   // skip stale entries (lazy deletion)
        }
    }

    void release(long long lease_id) {
        string task = leases[lease_id];
        leases.erase(lease_id);
        if (attempts[task] >= max_attempts) {
            state[task] = "dead";
        } else {
            state[task] = "queued";
            ready.push_back(task);                // back of the line
        }
    }

    long long grant(long long now, const string& task) {
        long long lease_id = next_id++;
        leases[lease_id] = task;
        deadlines.push({now + lease_time, lease_id});
        return lease_id;
    }

public:
    WorkQueue(long long lease_time, int max_attempts) : lease_time(lease_time), max_attempts(max_attempts) {}

    bool add(long long now, const string& task) {
        expire(now);
        if (state.count(task)) return false;
        state[task] = "queued";
        attempts[task] = 0;
        ready.push_back(task);
        return true;
    }

    optional<pair<string, long long>> claim(long long now) {
        expire(now);
        if (ready.empty()) return nullopt;
        string task = ready.front();
        ready.pop_front();
        state[task] = "leased";
        attempts[task]++;
        return make_pair(task, grant(now, task));
    }

    optional<long long> heartbeat(long long now, long long lease_id) {
        expire(now);
        auto it = leases.find(lease_id);
        if (it == leases.end()) return nullopt;   // unknown, expired or finished
        string task = it->second;
        leases.erase(it);
        return grant(now, task);                  // a new id: the old heap entry goes stale
    }

    bool complete(long long now, long long lease_id) {
        expire(now);
        auto it = leases.find(lease_id);
        if (it == leases.end()) return false;
        state[it->second] = "done";
        leases.erase(it);
        return true;
    }

    bool fail(long long now, long long lease_id) {
        expire(now);
        if (!leases.count(lease_id)) return false;
        release(lease_id);
        return true;
    }

    string status(long long now, const string& task) {
        expire(now);
        auto it = state.find(task);
        return it == state.end() ? "unknown" : it->second;
    }
};

import java.util.*;

// Tasks are queued -> leased -> done. A lease that runs out, or a failure,
// sends the task back to queued, or to dead once its attempts are used up.
class WorkQueue {
    private final long leaseTime;
    private final int maxAttempts;
    private final Map<String, String> state = new HashMap<>();      // task -> "queued" | "leased" | "done" | "dead"
    private final Map<String, Integer> attempts = new HashMap<>();  // task -> how many times it was claimed
    private final ArrayDeque<String> ready = new ArrayDeque<>();    // queued tasks, oldest first
    private final Map<Long, String> leases = new HashMap<>();       // live lease_id -> task
    // (deadline, lease_id); may hold stale entries
    private final PriorityQueue<long[]> deadlines = new PriorityQueue<>((a, b) ->
            a[0] != b[0] ? Long.compare(a[0], b[0]) : Long.compare(a[1], b[1]));
    private long nextId = 1;

    WorkQueue(long leaseTime, int maxAttempts) {
        this.leaseTime = leaseTime;
        this.maxAttempts = maxAttempts;
    }

    // Runs first in every call: leases that ran out end before anything else.
    private void expire(long now) {
        while (!deadlines.isEmpty() && deadlines.peek()[0] <= now) {
            long leaseId = deadlines.poll()[1];
            if (leases.containsKey(leaseId)) release(leaseId);   // skip stale entries (lazy deletion)
        }
    }

    private void release(long leaseId) {
        String task = leases.remove(leaseId);
        if (attempts.get(task) >= maxAttempts) {
            state.put(task, "dead");
        } else {
            state.put(task, "queued");
            ready.addLast(task);                  // back of the line
        }
    }

    private long grant(long now, String task) {
        long leaseId = nextId++;
        leases.put(leaseId, task);
        deadlines.add(new long[]{now + leaseTime, leaseId});
        return leaseId;
    }

    boolean add(long now, String task) {
        expire(now);
        if (state.containsKey(task)) return false;
        state.put(task, "queued");
        attempts.put(task, 0);
        ready.addLast(task);
        return true;
    }

    // Returns {task, lease id} or null.
    String[] claim(long now) {
        expire(now);
        if (ready.isEmpty()) return null;
        String task = ready.pollFirst();
        state.put(task, "leased");
        attempts.merge(task, 1, Integer::sum);
        return new String[]{task, String.valueOf(grant(now, task))};
    }

    Long heartbeat(long now, long leaseId) {
        expire(now);
        String task = leases.remove(leaseId);
        if (task == null) return null;            // unknown, expired or finished
        return grant(now, task);                  // a new id: the old heap entry goes stale
    }

    boolean complete(long now, long leaseId) {
        expire(now);
        String task = leases.remove(leaseId);
        if (task == null) return false;
        state.put(task, "done");
        return true;
    }

    boolean fail(long now, long leaseId) {
        expire(now);
        if (!leases.containsKey(leaseId)) return false;
        release(leaseId);
        return true;
    }

    String status(long now, String task) {
        expire(now);
        return state.getOrDefault(task, "unknown");
    }
}