Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 16 additions & 9 deletions crates/taskito-core/src/scheduler/poller.rs
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,7 @@ impl Scheduler {
// If we exceed the cap, roll back: clear the claim row and reset
// status to `Pending` so the job can be dispatched again later.
if !self.check_post_claim_concurrency(&job)? {
self.rollback_claim_and_retry(&job.id, now + CONCURRENCY_RETRY_DELAY_MS)?;
self.rollback_claim_and_reschedule(&job.id, now + CONCURRENCY_RETRY_DELAY_MS)?;
return Ok(false);
}

Expand All @@ -140,14 +140,20 @@ impl Scheduler {
Ok(()) => Ok(true),
Err(TrySendError::Full(_)) => {
warn!("worker channel full; rescheduling job {job_id} (worker pool is behind)",);
self.rollback_claim_and_retry(&job_id, now + CHANNEL_BACKPRESSURE_RETRY_DELAY_MS)?;
self.rollback_claim_and_reschedule(
&job_id,
now + CHANNEL_BACKPRESSURE_RETRY_DELAY_MS,
)?;
Ok(false)
}
Err(TrySendError::Closed(_)) => {
warn!(
"worker channel closed; rescheduling job {job_id} (worker pool shutting down)",
);
self.rollback_claim_and_retry(&job_id, now + CHANNEL_BACKPRESSURE_RETRY_DELAY_MS)?;
self.rollback_claim_and_reschedule(
&job_id,
now + CHANNEL_BACKPRESSURE_RETRY_DELAY_MS,
)?;
Ok(false)
}
}
Expand Down Expand Up @@ -193,7 +199,7 @@ impl Scheduler {
let key = format!("queue:{}", job.queue);
if !self.rate_limiter.try_acquire(&key, rl_config)? {
self.storage
.retry(&job.id, now + RATE_LIMIT_RETRY_DELAY_MS)?;
.reschedule(&job.id, now + RATE_LIMIT_RETRY_DELAY_MS)?;
return Ok(false);
}
}
Expand All @@ -202,14 +208,14 @@ impl Scheduler {
if let Some(config) = self.task_configs.get(&job.task_name) {
if config.circuit_breaker.is_some() && !self.circuit_breaker.allow(&job.task_name)? {
self.storage
.retry(&job.id, now + CIRCUIT_BREAKER_RETRY_DELAY_MS)?;
.reschedule(&job.id, now + CIRCUIT_BREAKER_RETRY_DELAY_MS)?;
return Ok(false);
}

if let Some(ref rl_config) = config.rate_limit {
if !self.rate_limiter.try_acquire(&job.task_name, rl_config)? {
self.storage
.retry(&job.id, now + RATE_LIMIT_RETRY_DELAY_MS)?;
.reschedule(&job.id, now + RATE_LIMIT_RETRY_DELAY_MS)?;
return Ok(false);
}
}
Expand Down Expand Up @@ -265,11 +271,12 @@ impl Scheduler {

/// Reverse a successful `claim_execution` and reschedule the job. Used
/// when a post-claim gate rejects the job after the claim row has
/// already been written.
fn rollback_claim_and_retry(&self, job_id: &str, next_at: i64) -> Result<()> {
/// already been written. Uses `reschedule` (not `retry`) so a job that
/// never executed does not lose retry budget.
fn rollback_claim_and_reschedule(&self, job_id: &str, next_at: i64) -> Result<()> {
if let Err(e) = self.storage.complete_execution(job_id) {
warn!("failed to clear execution claim during rollback for job {job_id}: {e}");
}
self.storage.retry(job_id, next_at)
self.storage.reschedule(job_id, next_at)
}
}
25 changes: 25 additions & 0 deletions crates/taskito-core/src/storage/diesel_common/jobs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -705,6 +705,31 @@ macro_rules! impl_diesel_job_ops {
Ok(())
}

/// Re-schedule a job back to `Pending` without touching
/// `retry_count`. Mirrors [`retry`](Self::retry) for soft-gate
/// reschedules (rate limit, circuit breaker, concurrency cap,
/// backpressure) where the job never executed, so its retry
/// budget must be preserved.
pub fn reschedule(&self, id: &str, next_scheduled_at: i64) -> Result<()> {
let mut conn = self.conn()?;

let affected = diesel::update(jobs::table)
.filter(jobs::id.eq(id))
.set((
jobs::status.eq(JobStatus::Pending as i32),
jobs::scheduled_at.eq(next_scheduled_at),
jobs::started_at.eq(None::<i64>),
jobs::completed_at.eq(None::<i64>),
jobs::error.eq(None::<String>),
))
.execute(&mut conn)?;

if affected == 0 {
return Err(QueueError::JobNotFound(id.to_string()));
}
Ok(())
}

/// Cancel a pending job and cascade-cancel its dependents. The
/// cancelled job moves from `jobs` into `archived_jobs`.
pub fn cancel_job(&self, id: &str) -> Result<bool> {
Expand Down
6 changes: 6 additions & 0 deletions crates/taskito-core/src/storage/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,9 @@ macro_rules! impl_storage {
fn retry(&self, id: &str, next_scheduled_at: i64) -> $crate::error::Result<()> {
self.retry(id, next_scheduled_at)
}
fn reschedule(&self, id: &str, next_scheduled_at: i64) -> $crate::error::Result<()> {
self.reschedule(id, next_scheduled_at)
}
fn cancel_job(&self, id: &str) -> $crate::error::Result<bool> {
self.cancel_job(id)
}
Expand Down Expand Up @@ -679,6 +682,9 @@ impl Storage for StorageBackend {
fn retry(&self, id: &str, next_scheduled_at: i64) -> Result<()> {
delegate!(self, retry, id, next_scheduled_at)
}
fn reschedule(&self, id: &str, next_scheduled_at: i64) -> Result<()> {
delegate!(self, reschedule, id, next_scheduled_at)
}
fn cancel_job(&self, id: &str) -> Result<bool> {
delegate!(self, cancel_job, id)
}
Expand Down
25 changes: 25 additions & 0 deletions crates/taskito-core/src/storage/redis_backend/jobs/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,31 @@ impl RedisStorage {
Ok(())
}

/// Re-schedule a job back to `Pending` without consuming retry budget.
/// Mirrors [`retry`](Self::retry) for soft-gate reschedules where the job
/// never executed, so `retry_count` must be preserved.
pub fn reschedule(&self, id: &str, next_scheduled_at: i64) -> Result<()> {
let mut conn = self.conn()?;
let mut job = self.get_job_required(id)?;
let old_status = job.status;

job.status = JobStatus::Pending;
job.scheduled_at = next_scheduled_at;
job.started_at = None;
job.completed_at = None;
job.error = None;

self.save_job_and_move_status(&mut conn, &job, old_status)?;

// Re-add to pending queue
let queue_key = self.key(&["queue", &job.queue, "pending"]);
let score = dequeue_score(job.priority, job.scheduled_at);
conn.zadd::<_, _, _, ()>(&queue_key, &job.id, score)
.map_err(map_err)?;

Ok(())
}

pub fn cancel_job(&self, id: &str) -> Result<bool> {
let mut conn = self.conn()?;
let job = match self.get_job(id)? {
Expand Down
33 changes: 33 additions & 0 deletions crates/taskito-core/src/storage/sqlite/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -294,6 +294,39 @@ fn test_retry_reschedule() {
assert_eq!(fetched.scheduled_at, future);
}

#[test]
fn test_reschedule_preserves_retry_count() {
// Soft-gate reschedules (rate limit, circuit breaker, concurrency,
// backpressure) must NOT consume the job's retry budget, unlike retry().
let storage = test_storage();
let job = storage.enqueue(make_job("reschedule_task")).unwrap();
storage
.dequeue("default", now_millis() + 1000, None)
.unwrap();

let future = now_millis() + 5000;
storage.reschedule(&job.id, future).unwrap();

let fetched = storage.get_job(&job.id).unwrap().unwrap();
assert_eq!(fetched.status, JobStatus::Pending);
assert_eq!(fetched.scheduled_at, future);
assert_eq!(
fetched.retry_count, 0,
"reschedule must not burn retry budget"
);

// Repeated reschedules still never touch retry_count.
storage
.dequeue("default", now_millis() + 1000, None)
.unwrap();
storage.reschedule(&job.id, future + 1000).unwrap();
let again = storage.get_job(&job.id).unwrap().unwrap();
assert_eq!(again.retry_count, 0);

// Unknown id is reported, matching retry().
assert!(storage.reschedule("missing-id", future).is_err());
}

#[test]
fn test_dead_letter_queue() {
let storage = test_storage();
Expand Down
5 changes: 5 additions & 0 deletions crates/taskito-core/src/storage/traits.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,11 @@ pub trait Storage: Send + Sync + Clone {
fn complete(&self, id: &str, result_bytes: Option<Vec<u8>>) -> Result<()>;
fn fail(&self, id: &str, error: &str) -> Result<()>;
fn retry(&self, id: &str, next_scheduled_at: i64) -> Result<()>;
/// Re-schedule a job back to `Pending` **without** consuming its retry
/// budget. Used for soft-gate reschedules (rate limit, circuit breaker,
/// concurrency cap, channel backpressure) where the job never executed,
/// unlike [`retry`](Self::retry) which increments `retry_count`.
fn reschedule(&self, id: &str, next_scheduled_at: i64) -> Result<()>;
fn cancel_job(&self, id: &str) -> Result<bool>;
fn request_cancel(&self, id: &str) -> Result<bool>;
fn is_cancel_requested(&self, id: &str) -> Result<bool>;
Expand Down
20 changes: 20 additions & 0 deletions crates/taskito-core/tests/rust/storage_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,25 @@ fn test_retry(s: &impl Storage) {
assert_eq!(fetched.scheduled_at, future);
}

fn test_reschedule(s: &impl Storage) {
// reschedule() must restore the job to Pending without incrementing
// retry_count — the soft-gate parity contract across all backends.
let q = "q-reschedule";
let job = s.enqueue(make_job(q, "reschedule_task")).unwrap();
s.dequeue(q, now_millis() + 1000, None).unwrap();

let future = now_millis() + 5000;
s.reschedule(&job.id, future).unwrap();

let fetched = s.get_job(&job.id).unwrap().unwrap();
assert_eq!(fetched.status, JobStatus::Pending);
assert_eq!(fetched.scheduled_at, future);
assert_eq!(
fetched.retry_count, 0,
"reschedule must not burn retry budget"
);
}

fn test_cancel_job(s: &impl Storage) {
let job = s.enqueue(make_job("q-cancel", "cancel_me")).unwrap();
assert!(s.cancel_job(&job.id).unwrap());
Expand Down Expand Up @@ -579,6 +598,7 @@ fn run_storage_tests(s: &impl Storage) {
test_complete(s);
test_fail(s);
test_retry(s);
test_reschedule(s);
test_cancel_job(s);
test_stats(s);
test_stats_by_queue_and_task(s);
Expand Down
8 changes: 6 additions & 2 deletions py_src/taskito/dashboard/handlers/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,12 @@
# Internal namespaces that must never be exposed or mutated through the public
# settings API. ``auth:`` holds password hashes, live session records (with
# their CSRF tokens), and the CSRF secret — reading or overwriting any of these
# is a full auth bypass, so the settings endpoints treat them as if absent.
_PROTECTED_PREFIXES = ("auth:",)
# is a full auth bypass. ``webhooks:`` holds subscription rows with plaintext
# HMAC signing secrets plus delivery logs (full request payloads / response
# bodies); the dedicated webhooks API redacts the secret to ``has_secret``, so
# these must stay out of the generic settings endpoints too. The settings
# handlers treat every protected key as if absent.
_PROTECTED_PREFIXES = ("auth:", "webhooks:")


def _is_protected(key: str) -> bool:
Expand Down
15 changes: 14 additions & 1 deletion py_src/taskito/notes.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
from __future__ import annotations

import json
import math
from typing import Any

from taskito.exceptions import NotesValidationError
Expand Down Expand Up @@ -73,7 +74,12 @@ def validate_and_encode_notes(notes: dict[str, Any] | None) -> str | None:
_validate_key(key)
_validate_value(key, value, depth=1)

encoded = json.dumps(notes, sort_keys=True, ensure_ascii=False, separators=(",", ":"))
# allow_nan=False is a belt-and-suspenders guard: _validate_value already
# rejects non-finite floats, but this ensures the encoder can never emit
# invalid JSON tokens even if a new code path slips one through.
encoded = json.dumps(
notes, sort_keys=True, ensure_ascii=False, separators=(",", ":"), allow_nan=False
)
# For the common all-ASCII payload, byte length == char length, so skip the
# throwaway UTF-8 encode just to measure size.
encoded_bytes = len(encoded) if encoded.isascii() else len(encoded.encode("utf-8"))
Expand Down Expand Up @@ -107,6 +113,13 @@ def _validate_value(path: str, value: Any, *, depth: int) -> None:
f"limit of {MAX_NOTE_VALUE_LENGTH}"
)
return
if isinstance(value, float) and not math.isfinite(value):
# NaN/Infinity serialize to bare ``NaN``/``Infinity`` tokens that are
# invalid JSON — any strict reader (the dashboard's JS JSON.parse,
# Rust serde_json) rejects them, breaking the whole API response.
raise NotesValidationError(
f"note value at {path!r} must be a finite number, got {value!r}"
)
if isinstance(value, (bool, int, float)) or value is None:
return
if isinstance(value, list):
Expand Down
6 changes: 5 additions & 1 deletion py_src/taskito/result.py
Original file line number Diff line number Diff line change
Expand Up @@ -285,7 +285,11 @@ def to_dict(self) -> dict[str, Any]:
"timeout_ms": self._py_job.timeout_ms,
"unique_key": self._py_job.unique_key,
"metadata": self._py_job.metadata,
"notes": self.notes,
# Emit the raw canonical JSON string (the dashboard API contract is
# ``notes: string | null`` and the client JSON.parse-s it). The
# ``notes`` property returns a parsed dict for Python callers, but
# serializing that dict here would break the typed client contract.
"notes": self._py_job.notes,
"namespace": self._py_job.namespace,
}

Expand Down
31 changes: 31 additions & 0 deletions tests/core/test_notes.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

from __future__ import annotations

import json
from typing import Any

import pytest
Expand Down Expand Up @@ -114,6 +115,19 @@ def test_primitives_accepted() -> None:
assert encoded == '{"b":true,"f":1.5,"i":1,"n":null,"s":"x"}'


@pytest.mark.parametrize("bad", [float("inf"), float("-inf"), float("nan")])
def test_non_finite_floats_rejected(bad: float) -> None:
# NaN/Infinity serialize to invalid JSON tokens that break strict readers
# (the dashboard JSON.parse, Rust serde_json) — reject at the boundary.
with pytest.raises(NotesValidationError, match="must be a finite number"):
validate_and_encode_notes({"x": bad})


def test_non_finite_float_rejected_when_nested() -> None:
with pytest.raises(NotesValidationError, match="must be a finite number"):
validate_and_encode_notes({"outer": {"inner": [1, float("nan")]}})


# ── End-to-end through Queue.enqueue ----------------------------------------


Expand All @@ -126,6 +140,23 @@ def noop() -> str:
assert job.notes == {"customer_id": "cus_abc", "tier": "gold"}


def test_to_dict_emits_notes_as_json_string(queue: Queue) -> None:
# The dashboard API contract is ``notes: string | null`` (the client
# JSON.parse-s it). to_dict() must emit the canonical string, not the
# parsed dict the ``notes`` property returns for Python callers.
@queue.task()
def noop() -> str:
return "ok"

job = noop.apply_async(notes={"customer_id": "cus_abc", "tier": "gold"})
raw = job.to_dict()["notes"]
assert isinstance(raw, str)
assert json.loads(raw) == {"customer_id": "cus_abc", "tier": "gold"}

plain = noop.apply_async()
assert plain.to_dict()["notes"] is None


def test_enqueue_without_notes_returns_none(queue: Queue) -> None:
@queue.task()
def noop() -> str:
Expand Down
20 changes: 20 additions & 0 deletions tests/dashboard/test_dashboard_settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,26 @@ def test_get_unknown_setting_returns_404(dashboard_server: tuple[AuthedClient, Q
assert exc_info.value.code == 404


def test_get_settings_hides_webhook_keys(dashboard_server: tuple[AuthedClient, Queue]) -> None:
client, queue = dashboard_server
# Webhook subscriptions carry plaintext HMAC signing secrets; delivery logs
# carry full request payloads. Both live under ``webhooks:`` in the same KV
# store as /api/settings and must never leak through the generic endpoints.
queue.set_setting(
"webhooks:subscriptions",
json.dumps([{"id": "wh1", "url": "https://x.example", "secret": "topsecret"}]),
)
queue.set_setting("webhooks:deliveries:wh1", json.dumps([{"payload": "sensitive"}]))

snapshot = client.get("/api/settings")
assert not any(k.startswith("webhooks:") for k in snapshot)

# Direct fetch of a protected key is reported as absent, not leaked.
with pytest.raises(urllib.error.HTTPError) as exc_info:
client.get("/api/settings/webhooks:subscriptions")
assert exc_info.value.code == 404


def test_put_setting_with_missing_value_field_returns_400(
dashboard_server: tuple[AuthedClient, Queue],
) -> None:
Expand Down
Loading