retry faileds, remove jobs and check completion

Author: leodisarli

.gitmodules

@@ -2,15 +2,15 @@
 import redis
 
 from dataclasses import dataclass
-from typing import Optional, Any
+from typing import Optional, Any, List
 from threading import Lock
 
 from .clock import now_ms
 from .ids import new_ulid
-from .types import ReservePaused, ReserveJob, ReserveResult, AckFailResult
+from .types import ReservePaused, ReserveJob, ReserveResult, AckFailResult, BatchRemoveResult, BatchRetryFailedResult
 from .transport import RedisLike
 from .scripts import OmniqScripts
-from .helper import queue_base, queue_anchor
+from .helper import queue_base, queue_anchor, check_completion_anchor
 
 @dataclass
 class OmniqOps:
@@ -300,6 +300,195 @@ def job_timeout_ms(self, *, queue: str, job_id: str, default_ms: int = 60_000) -
         except Exception:
             n = 0
         return n if n > 0 else int(default_ms)
+    
+    def retry_failed(self, *, queue: str, job_id: str, now_ms_override: int = 0) -> None:
+        anchor = queue_anchor(queue)
+        nms = now_ms_override or now_ms()
+
+        res = self._evalsha_with_noscript_fallback(
+            self.scripts.retry_failed.sha,
+            self.scripts.retry_failed.src,
+            1,
+            anchor,
+            job_id,
+            str(int(nms)),
+        )
+
+        if not isinstance(res, list) or len(res) < 1:
+            raise RuntimeError(f"Unexpected RETRY_FAILED response: {res}")
+
+        if res[0] == "OK":
+            return
+
+        if res[0] == "ERR":
+            reason = str(res[1]) if len(res) > 1 else "UNKNOWN"
+            raise RuntimeError(f"RETRY_FAILED failed: {reason}")
+
+        raise RuntimeError(f"Unexpected RETRY_FAILED response: {res}")
+
+    def retry_failed_batch(
+        self,
+        *,
+        queue: str,
+        job_ids: List[str],
+        now_ms_override: int = 0,
+    ) -> BatchRetryFailedResult:
+        if len(job_ids) > 100:
+            raise ValueError("retry_failed_batch max is 100 job_ids per call")
+
+        anchor = queue_anchor(queue)
+        nms = now_ms_override or now_ms()
+
+        argv: list[str] = [str(int(nms)), str(len(job_ids))]
+        argv.extend([str(j) for j in job_ids])
+
+        res = self._evalsha_with_noscript_fallback(
+            self.scripts.retry_failed_batch.sha,
+            self.scripts.retry_failed_batch.src,
+            1,
+            anchor,
+            *argv,
+        )
+
+        if not isinstance(res, list):
+            raise RuntimeError(f"Unexpected RETRY_FAILED_BATCH response: {res}")
+
+        if len(res) >= 2 and str(res[0]) == "ERR":
+            reason = str(res[1])
+            extra = str(res[2]) if len(res) > 2 else ""
+            raise RuntimeError(f"RETRY_FAILED_BATCH failed: {reason} {extra}".strip())
+
+        out: BatchRetryFailedResult = []
+        i = 0
+        while i < len(res):
+            job_id = str(res[i] or "")
+            status = str(res[i + 1] or "")
+            reason: Optional[str] = None
+            if status == "ERR":
+                reason = str(res[i + 2] or "UNKNOWN")
+                i += 3
+            else:
+                i += 2
+            out.append((job_id, status, reason))
+        return out
+
+    def remove_job(self, *, queue: str, job_id: str, lane: str) -> str:
+        anchor = queue_anchor(queue)
+
+        res = self._evalsha_with_noscript_fallback(
+            self.scripts.remove_job.sha,
+            self.scripts.remove_job.src,
+            1,
+            anchor,
+            job_id,
+            lane,
+        )
+
+        if not isinstance(res, list) or len(res) < 1:
+            raise RuntimeError(f"Unexpected REMOVE_JOB response: {res}")
+
+        if res[0] == "OK":
+            return str(res[0] or "")
+
+        if res[0] == "ERR":
+            reason = str(res[1]) if len(res) > 1 else "UNKNOWN"
+            raise RuntimeError(f"REMOVE_JOB failed: {reason}")
+
+        raise RuntimeError(f"Unexpected REMOVE_JOB response: {res}")
+
+    def remove_jobs_batch(
+        self,
+        *,
+        queue: str,
+        lane: str,
+        job_ids: List[str],
+    ) -> BatchRemoveResult:
+        if len(job_ids) > 100:
+            raise ValueError("remove_jobs_batch max is 100 job_ids per call")
+
+        anchor = queue_anchor(queue)
+
+        argv: list[str] = [str(lane), str(len(job_ids))]
+        argv.extend([str(j) for j in job_ids])
+
+        res = self._evalsha_with_noscript_fallback(
+            self.scripts.remove_jobs_batch.sha,
+            self.scripts.remove_jobs_batch.src,
+            1,
+            anchor,
+            *argv,
+        )
+
+        if not isinstance(res, list):
+            raise RuntimeError(f"Unexpected REMOVE_JOBS_BATCH response: {res}")
+
+        if len(res) >= 2 and str(res[0]) == "ERR":
+            reason = str(res[1])
+            extra = str(res[2]) if len(res) > 2 else ""
+            raise RuntimeError(f"REMOVE_JOBS_BATCH failed: {reason} {extra}".strip())
+
+        out: BatchRemoveResult = []
+        i = 0
+        while i < len(res):
+            job_id = str(res[i] or "")
+            status = str(res[i + 1] or "")
+            reason: Optional[str] = None
+            if status == "ERR":
+                reason = str(res[i + 2] or "UNKNOWN")
+                i += 3
+            else:
+                i += 2
+            out.append((job_id, status, reason))
+        return out
+    
+    def check_completion_init_job_counter(self, *, key: str, expected: int) -> None:
+        anchor = check_completion_anchor(key)
+
+        res = self._evalsha_with_noscript_fallback(
+            self.scripts.check_completion_init.sha,
+            self.scripts.check_completion_init.src,
+            1,
+            anchor,
+            str(int(expected)),
+        )
+
+        if not isinstance(res, list) or len(res) < 1:
+            raise RuntimeError(f"Unexpected CHECK_COMPLETION_INIT response: {res}")
+
+        if res[0] == "OK":
+            return
+
+        if res[0] == "ERR":
+            reason = str(res[1]) if len(res) > 1 else "UNKNOWN"
+            raise RuntimeError(f"CHECK_COMPLETION_INIT failed: {reason}")
+
+        raise RuntimeError(f"Unexpected CHECK_COMPLETION_INIT response: {res}")
+
+    def check_completion_job_decrement(self, *, key: str, child_id: str) -> int:
+        anchor = check_completion_anchor(key)
+
+        res = self._evalsha_with_noscript_fallback(
+            self.scripts.check_completion_decrement.sha,
+            self.scripts.check_completion_decrement.src,
+            1,
+            anchor,
+            str(child_id),
+        )
+
+        if not isinstance(res, list) or len(res) < 2:
+            raise RuntimeError(f"Unexpected CHECK_COMPLETION_DECREMENT response: {res}")
+
+        if res[0] == "OK":
+            try:
+                return int(res[1])
+            except Exception:
+                raise RuntimeError(f"Unexpected CHECK_COMPLETION_DECREMENT remaining: {res}")
+
+        if res[0] == "ERR":
+            reason = str(res[1]) if len(res) > 1 else "UNKNOWN"
+            raise RuntimeError(f"CHECK_COMPLETION_DECREMENT failed: {reason}")
+
+        raise RuntimeError(f"Unexpected CHECK_COMPLETION_DECREMENT response: {res}")
 
     @staticmethod
     def paused_backoff_s(poll_interval_s: float) -> float:

src/omniq/_ops.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from ._ops import OmniqOps
+
+
+@dataclass(frozen=True)
+class CheckCompletion:
+    ops: OmniqOps
+    default_child_id: str
+
+    def InitJobCounter(self, key: str, expected: int) -> None:
+        self.ops.check_completion_init_job_counter(key=key, expected=int(expected))
+
+    def JobDecrement(self, key: str, child_id: Optional[str] = None) -> int:
+        cid = (child_id or self.default_child_id or "").strip()
+        if not cid:
+            raise ValueError("child_id is required (or provide default_child_id)")
+        return int(self.ops.check_completion_job_decrement(key=key, child_id=cid))

src/omniq/check_completion.py

@@ -102,6 +102,18 @@ def resume(self, *, queue: str) -> int:
     def is_paused(self, *, queue: str) -> bool:
         return self._ops.is_paused(queue=queue)
 
+    def retry_failed(self, *, queue: str, job_id: str, now_ms_override: int = 0) -> None:
+        return self._ops.retry_failed(queue=queue, job_id=job_id, now_ms_override=now_ms_override)
+
+    def retry_failed_batch(self, *, queue: str, job_ids: list[str], now_ms_override: int = 0) -> None:
+        return self._ops.retry_failed_batch(queue=queue, job_ids=job_ids, now_ms_override=now_ms_override)
+
+    def remove_job(self, *, queue: str, job_id: str, lane: str) -> None:
+        return self._ops.remove_job(queue=queue, job_id=job_id, lane=lane)
+    
+    def remove_jobs_batch(self, *, queue: str, lane: str, job_ids: list[str]):
+        return self._ops.remove_jobs_batch(queue=queue, lane=lane, job_ids=job_ids)
+
     def consume(
         self,
         *,

src/omniq/client.py

@@ -7,6 +7,7 @@
 
 from ._ops import OmniqOps
 from .types import JobCtx, ReserveJob
+from .check_completion import CheckCompletion
 
 @dataclass
 class StopController:
@@ -176,6 +177,7 @@ def on_sigint(signum, frame):
             except Exception:
                 payload_obj = res.payload
 
+            cc = CheckCompletion(ops=ops, default_child_id=res.job_id)
             ctx = JobCtx(
                 queue=queue,
                 job_id=res.job_id,
@@ -185,6 +187,7 @@ def on_sigint(signum, frame):
                 lock_until_ms=res.lock_until_ms,
                 lease_token=res.lease_token,
                 gid=res.gid,
+                check_completion=cc,
             )
 
             if verbose:

src/omniq/consumer.py

@@ -0,0 +1,112 @@
+local anchor      = KEYS[1]
+local job_id      = ARGV[1]
+local now_ms      = tonumber(ARGV[2] or "0")
+local lease_token = ARGV[3]
+local err_msg     = ARGV[4]
+
+local DEFAULT_GROUP_LIMIT = 1
+local MAX_ERR_BYTES = 4096
+
+local function derive_base(a)
+  if a == nil or a == "" then return "" end
+  if string.sub(a, -5) == ":meta" then
+    return string.sub(a, 1, -6)
+  end
+  return a
+end
+
+local base = derive_base(anchor)
+
+local k_job     = base .. ":job:" .. job_id
+local k_active  = base .. ":active"
+local k_delayed = base .. ":delayed"
+local k_failed  = base .. ":failed"
+local k_gready  = base .. ":groups:ready"
+
+local function to_i(v)
+  if v == false or v == nil or v == '' then return 0 end
+  local n = tonumber(v)
+  if n == nil then return 0 end
+  return math.floor(n)
+end
+
+local function dec_floor0(key)
+  local v = to_i(redis.call("DECR", key))
+  if v < 0 then
+    redis.call("SET", key, "0")
+    return 0
+  end
+  return v
+end
+
+local function group_limit_for(gid)
+  local k_glimit = base .. ":g:" .. gid .. ":limit"
+  local lim = to_i(redis.call("GET", k_glimit))
+  if lim <= 0 then return DEFAULT_GROUP_LIMIT end
+  return lim
+end
+
+local function maybe_store_last_error()
+  if err_msg == nil or err_msg == "" then return end
+  if string.len(err_msg) > MAX_ERR_BYTES then
+    err_msg = string.sub(err_msg, 1, MAX_ERR_BYTES)
+  end
+  redis.call("HSET", k_job,
+    "last_error", err_msg,
+    "last_error_ms", tostring(now_ms)
+  )
+end
+
+if lease_token == nil or lease_token == "" then
+  return {"ERR", "TOKEN_REQUIRED"}
+end
+
+local cur_token = redis.call("HGET", k_job, "lease_token") or ""
+if cur_token ~= lease_token then
+  return {"ERR", "TOKEN_MISMATCH"}
+end
+
+if redis.call("ZREM", k_active, job_id) ~= 1 then
+  return {"ERR", "NOT_ACTIVE"}
+end
+
+maybe_store_last_error()
+
+local gid = redis.call("HGET", k_job, "gid")
+if gid and gid ~= "" then
+  local k_ginflight = base .. ":g:" .. gid .. ":inflight"
+  local inflight = dec_floor0(k_ginflight)
+  local limit = group_limit_for(gid)
+  local k_gwait = base .. ":g:" .. gid .. ":wait"
+  if inflight < limit and to_i(redis.call("LLEN", k_gwait)) > 0 then
+    redis.call("ZADD", k_gready, now_ms, gid)
+  end
+end
+
+local attempt      = to_i(redis.call("HGET", k_job, "attempt"))
+local max_attempts = to_i(redis.call("HGET", k_job, "max_attempts"))
+if max_attempts <= 0 then max_attempts = 1 end
+local backoff_ms   = to_i(redis.call("HGET", k_job, "backoff_ms"))
+
+if attempt >= max_attempts then
+  redis.call("HSET", k_job,
+    "state", "failed",
+    "updated_ms", tostring(now_ms),
+    "lease_token", "",
+    "lock_until_ms", ""
+  )
+  redis.call("LPUSH", k_failed, job_id)
+  return {"FAILED"}
+end
+
+local due_ms = now_ms + backoff_ms
+redis.call("HSET", k_job,
+  "state", "delayed",
+  "due_ms", tostring(due_ms),
+  "updated_ms", tostring(now_ms),
+  "lease_token", "",
+  "lock_until_ms", ""
+)
+redis.call("ZADD", k_delayed, due_ms, job_id)
+
+return {"RETRY", tostring(due_ms)}