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
13 changes: 13 additions & 0 deletions crates/taskito-core/src/storage/diesel_common/jobs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1976,6 +1976,19 @@ macro_rules! impl_diesel_job_ops {
Ok(count)
}

/// Count pending jobs on a queue (for the `max_pending` admission cap).
pub fn count_pending_by_queue(&self, queue_name: &str) -> Result<i64> {
let mut conn = self.conn()?;

let count: i64 = jobs::table
.filter(jobs::queue.eq(queue_name))
.filter(jobs::status.eq(JobStatus::Pending as i32))
.count()
.get_result(&mut conn)?;

Ok(count)
}

/// Purge job errors older than the given timestamp.
///
/// Deletes in bounded batches, each its own txn — see
Expand Down
9 changes: 9 additions & 0 deletions crates/taskito-core/src/storage/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -835,6 +835,12 @@ macro_rules! impl_storage {
) -> $crate::error::Result<i64> {
self.count_running_by_task(task_name)
}
fn count_pending_by_queue(
&self,
queue_name: &str,
) -> $crate::error::Result<i64> {
self.count_pending_by_queue(queue_name)
}
fn stats_by_queue(
&self,
queue_name: &str,
Expand Down Expand Up @@ -1444,6 +1450,9 @@ impl Storage for StorageBackend {
fn count_running_by_task(&self, task_name: &str) -> Result<i64> {
delegate!(self, count_running_by_task, task_name)
}
fn count_pending_by_queue(&self, queue_name: &str) -> Result<i64> {
delegate!(self, count_pending_by_queue, queue_name)
}
fn stats_by_queue(&self, queue_name: &str) -> Result<QueueStats> {
delegate!(self, stats_by_queue, queue_name)
}
Expand Down
7 changes: 7 additions & 0 deletions crates/taskito-core/src/storage/redis_backend/jobs/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -254,6 +254,13 @@ impl RedisStorage {
self.count_in_status(&mut conn, &by_task_key, JobStatus::Running)
}

/// Count pending jobs on a queue (for the `max_pending` admission cap).
pub fn count_pending_by_queue(&self, queue_name: &str) -> Result<i64> {
let mut conn = self.conn()?;
let by_queue_key = self.key(&["jobs", "by_queue", queue_name]);
self.count_in_status(&mut conn, &by_queue_key, JobStatus::Pending)
}

pub fn stats_by_queue(&self, queue_name: &str) -> Result<QueueStats> {
let mut conn = self.conn()?;
self.queue_stats(&mut conn, queue_name)
Expand Down
21 changes: 21 additions & 0 deletions crates/taskito-core/src/storage/sqlite/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -806,6 +806,27 @@ fn test_count_running_by_task() {
assert_eq!(storage.count_running_by_task("no_such_task").unwrap(), 0);
}

#[test]
fn test_count_pending_by_queue() {
let storage = test_storage();
assert_eq!(storage.count_pending_by_queue("default").unwrap(), 0);

storage.enqueue(make_job("task_a")).unwrap();
storage.enqueue(make_job("task_a")).unwrap();
let mut other = make_job("task_b");
other.queue = "other".to_string();
storage.enqueue(other).unwrap();

assert_eq!(storage.count_pending_by_queue("default").unwrap(), 2);
assert_eq!(storage.count_pending_by_queue("other").unwrap(), 1);
assert_eq!(storage.count_pending_by_queue("empty").unwrap(), 0);

// Dequeue drops the job out of Pending → count decreases.
let now = now_millis() + 1000;
storage.dequeue("default", now, None).unwrap().unwrap();
assert_eq!(storage.count_pending_by_queue("default").unwrap(), 1);
}

#[test]
fn test_enqueue_rejects_missing_dependency() {
let storage = test_storage();
Expand Down
4 changes: 4 additions & 0 deletions crates/taskito-core/src/storage/traits.rs
Original file line number Diff line number Diff line change
Expand Up @@ -329,6 +329,10 @@ pub trait Storage: Send + Sync + Clone {

// ── Per-queue stats ──────────────────────────────────────────

/// Cheap count of pending jobs on a queue — the admission-cap primitive.
/// Single-status, unlike the full-breakdown `stats_by_queue`.
fn count_pending_by_queue(&self, queue_name: &str) -> Result<i64>;

fn stats_by_queue(&self, queue_name: &str) -> Result<QueueStats>;
fn stats_all_queues(&self) -> Result<std::collections::HashMap<String, QueueStats>>;

Expand Down
3 changes: 3 additions & 0 deletions crates/taskito-core/tests/rust/storage_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,8 @@ fn test_stats_by_queue_and_task(s: &impl Storage) {
assert_eq!(st.pending, 3);
assert_eq!(st.running, 0);
assert_eq!(s.count_running_by_task(task).unwrap(), 0);
// Lean pending-count primitive agrees with the full breakdown.
assert_eq!(s.count_pending_by_queue(q).unwrap(), 3);

// Run two of them.
let d1 = s.dequeue(q, now_millis() + 1000, None).unwrap().unwrap();
Expand All @@ -181,6 +183,7 @@ fn test_stats_by_queue_and_task(s: &impl Storage) {
let st = s.stats_by_queue(q).unwrap();
assert_eq!(st.running, 2);
assert_eq!(st.pending, 1);
assert_eq!(s.count_pending_by_queue(q).unwrap(), 1);

// Complete one — running drops, completed rises.
s.complete(&d1.id, None).unwrap();
Expand Down
18 changes: 18 additions & 0 deletions crates/taskito-java/src/queue/inspect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,24 @@ pub extern "system" fn Java_org_byteveda_taskito_internal_NativeQueue_statsByQue
})
}

/// `long countPendingByQueue(long handle, String queue)` — the lean primitive
/// behind the `maxPending` admission cap.
#[no_mangle]
pub extern "system" fn Java_org_byteveda_taskito_internal_NativeQueue_countPendingByQueue<
'local,
>(
mut env: JNIEnv<'local>,
_class: JClass<'local>,
handle: jlong,
queue_name: JString<'local>,
) -> jlong {
guard(&mut env, 0, |env| {
let queue = unsafe { borrow_queue(handle) };
let name = read_string(env, &queue_name)?;
Ok(queue.storage.count_pending_by_queue(&name)?)
})
}

/// `String statsAllQueues(long handle)` — a JSON map of queue name to counts.
#[no_mangle]
pub extern "system" fn Java_org_byteveda_taskito_internal_NativeQueue_statsAllQueues<'local>(
Expand Down
10 changes: 10 additions & 0 deletions crates/taskito-node/src/queue/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,16 @@ impl JsQueue {
Ok(created.into_iter().map(|job| job.id).collect())
}

/// Count pending jobs on a queue — the lean primitive behind the
/// `maxPending` admission cap. Sync so the producer can gate a sync
/// `enqueue`/`enqueueMany` without a round trip to the event loop.
#[napi]
pub fn count_pending_by_queue(&self, queue: String) -> Result<i64> {
self.storage
.count_pending_by_queue(&queue)
.map_err(to_napi_err)
}

/// Fetch a job by id, or `null` if no such job exists.
#[napi]
pub fn get_job(&self, id: String) -> Result<Option<JsJob>> {
Expand Down
8 changes: 8 additions & 0 deletions crates/taskito-python/src/py_queue/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -589,6 +589,14 @@ impl PyQueue {
})
}

/// Count pending jobs on a queue — the lean primitive behind the
/// `max_pending` admission cap (avoids the full `stats_by_queue` breakdown).
pub fn count_pending_by_queue(&self, queue_name: &str) -> PyResult<i64> {
self.storage
.count_pending_by_queue(queue_name)
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
}

/// Get queue statistics broken down by queue name.
pub fn stats_all_queues(&self) -> PyResult<Py<PyAny>> {
let all = self
Expand Down
44 changes: 44 additions & 0 deletions sdks/java/src/main/java/org/byteveda/taskito/DefaultTaskito.java
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import org.byteveda.taskito.errors.EnqueueSkippedException;
import org.byteveda.taskito.errors.InterceptionException;
import org.byteveda.taskito.errors.PredicateRejectedException;
import org.byteveda.taskito.errors.QueueFullException;
import org.byteveda.taskito.errors.SerializationException;
import org.byteveda.taskito.errors.WorkflowException;
import org.byteveda.taskito.interception.Interception;
Expand Down Expand Up @@ -91,6 +92,8 @@ final class DefaultTaskito implements Taskito {
private final Map<String, List<EnqueueGate>> gates = new ConcurrentHashMap<>();
private final List<Interceptor> interceptors = new CopyOnWriteArrayList<>();
private final List<SubscriptionConfig> subscriptions = new CopyOnWriteArrayList<>();
// Opt-in per-queue admission caps (queue -> max pending). Absent = uncapped.
private final Map<String, Integer> maxPending = new ConcurrentHashMap<>();

DefaultTaskito(QueueBackend backend, Serializer serializer, Map<String, PayloadCodec> codecs) {
this.backend = backend;
Expand Down Expand Up @@ -145,6 +148,37 @@ public Map<String, ResourceStat> resourceMetrics() {
return resources.metrics();
}

@Override
public Taskito maxPending(String queue, int cap) {
if (cap < 0) {
// A negative cap makes `pending + incoming > cap` always true and
// would permanently reject every enqueue for the queue.
throw new IllegalArgumentException("cap must be non-negative");
}
maxPending.put(queue, cap);
return this;
}

/**
* Enforce the opt-in {@code maxPending} admission cap for a queue. Throws
* {@link QueueFullException} when admitting {@code incoming} jobs would push
* the queue's pending backlog past its cap; {@code incoming} is the batch
* size, so a batch is rejected as a whole rather than overshooting the cap by
* its full size. A no-op (and no query) for uncapped queues. Non-atomic
* count-then-insert, like the rate limiter.
*/
private void rejectIfQueueFull(String queueOrNull, int incoming) {
String queue = queueOrNull == null ? "default" : queueOrNull;
Integer cap = maxPending.get(queue);
if (cap == null) {
return;
}
long pending = backend.countPendingByQueue(queue);
if (pending + incoming > cap) {
throw new QueueFullException(queue, pending, cap);
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

@Override
public Taskito predicate(String taskName, Predicate predicate) {
return gate(
Expand Down Expand Up @@ -267,6 +301,8 @@ private Optional<String> dispatchEnqueue(
.metadata(encode(context.metadata()))
.build();
}
// Admission cap: reject before serializing/inserting if the target queue is full.
rejectIfQueueFull(finalOptions.queue(), 1);
// Serialize before codec-encoding so the idempotency key hashes the deterministic
// pre-codec payload — a non-deterministic codec (e.g. an AES-GCM nonce) must not
// change the dedup key.
Expand Down Expand Up @@ -376,6 +412,9 @@ public <T> List<String> enqueueMany(Task<T> task, List<T> payloads, EnqueueOptio
bytes[i] = encodeCodecs(payloadBytes, task.codecNames());
perJob.add(jobOptions);
}
// The whole batch targets one queue (a single `options`); reject before
// inserting if admitting all of it would push that queue past its cap.
rejectIfQueueFull(options.queue(), payloads.size());
return Arrays.asList(backend.enqueueMany(task.name(), bytes, encode(perJob)));
}

Expand Down Expand Up @@ -460,6 +499,11 @@ public QueueStats statsByQueue(String queue) {
return decode(backend.statsByQueueJson(queue), QueueStats.class);
}

@Override
public long countPendingByQueue(String queue) {
return backend.countPendingByQueue(queue);
}

@Override
public Map<String, QueueStats> statsAllQueues() {
return decodeMap(backend.statsAllQueuesJson(), QueueStats.class);
Expand Down
12 changes: 12 additions & 0 deletions sdks/java/src/main/java/org/byteveda/taskito/Taskito.java
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,15 @@ static Builder builder() {
*/
Taskito intercept(Interceptor interceptor);

/**
* Set an opt-in admission cap on {@code queue}'s pending backlog. Once the
* queue holds {@code cap} pending jobs, {@link #enqueue} throws
* {@link org.byteveda.taskito.errors.QueueFullException}. Enforced
* producer-side (a non-atomic count-then-insert), so it applies even with no
* worker running. Returns {@code this}.
*/
Taskito maxPending(String queue, int cap);

// ── Producer ────────────────────────────────────────────────────

/** Enqueue a typed payload using the task's default options; returns the job id. */
Expand Down Expand Up @@ -176,6 +185,9 @@ static Builder builder() {

QueueStats statsByQueue(String queue);

/** Count pending jobs on {@code queue} — the primitive behind the {@code maxPending} cap. */
long countPendingByQueue(String queue);

Map<String, QueueStats> statsAllQueues();

List<Job> listJobs(JobFilter filter);
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
package org.byteveda.taskito.errors;

import org.byteveda.taskito.TaskitoException;

/**
* An enqueue was rejected because the target queue reached its {@code maxPending}
* admission cap, so no job was created. Enforced producer-side (a non-atomic
* count-then-insert), so it applies even with no worker running.
*/
public class QueueFullException extends TaskitoException {
private final String queue;
private final long pending;
private final long cap;

public QueueFullException(String queue, long pending, long cap) {
super("queue '" + queue + "' is full: " + pending + " pending >= maxPending " + cap);
this.queue = queue;
this.pending = pending;
this.cap = cap;
}

/** The queue that rejected the enqueue. */
public String queue() {
return queue;
}

/** Pending count observed at rejection time. */
public long pending() {
return pending;
}

/** The configured cap. */
public long cap() {
return cap;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,11 @@ public String statsByQueueJson(String queue) {
return withOpenHandle(() -> NativeQueue.statsByQueue(handle, queue));
}

@Override
public long countPendingByQueue(String queue) {
return withOpenHandle(() -> NativeQueue.countPendingByQueue(handle, queue));
}

@Override
public String statsAllQueuesJson() {
return withOpenHandle(() -> NativeQueue.statsAllQueues(handle));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@ private NativeQueue() {}

public static native String statsByQueue(long handle, String queue);

public static native long countPendingByQueue(long handle, String queue);

public static native String statsAllQueues(long handle);

public static native String listJobs(long handle, String filterJson);
Expand Down
10 changes: 10 additions & 0 deletions sdks/java/src/main/java/org/byteveda/taskito/spi/QueueBackend.java
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,16 @@ public interface QueueBackend extends AutoCloseable {

String statsByQueueJson(String queue);

/**
* Count pending jobs on {@code queue} — the primitive behind the
* {@code maxPending} cap. Defaults to unsupported so this optional capability
* doesn't break existing third-party {@code QueueBackend} implementations at
* compile time; it is only invoked when a queue actually has a cap set.
*/
default long countPendingByQueue(String queue) {
throw new UnsupportedOperationException("countPendingByQueue not supported by this backend");
}

String statsAllQueuesJson();

String listJobsJson(String filterJson);
Expand Down
Loading