capp_queue/backend/
fjall_round_robin.rs

1use std::{collections::HashSet, marker::PhantomData, path::Path, sync::Mutex};
2
3use async_trait::async_trait;
4use fjall::{Database, Keyspace, KeyspaceCreateOptions, PersistMode};
5use serde::{Serialize, de::DeserializeOwned};
6use uuid::Uuid;
7
8use crate::{HasTagKey, Task, TaskId, TaskQueue, TaskQueueError, TaskSerializer};
9
10/// Round-robin Fjall task queue.
11///
12/// Pops scan forward from a cursor across all known tags, skipping any tag
13/// whose sub-queue is currently empty, and stop at the first tag with a
14/// queued task. A tag is never dropped from rotation just because it's
15/// momentarily empty — it's checked again on every subsequent pop until it
16/// either has work or the queue is added to. This keeps one busy tag from
17/// starving others, even when tags refill at very different rates. FIFO
18/// within a tag is preserved by sorting on the task's UUID v7 suffix.
19///
20/// Layout:
21/// - `tasks`: task_id (uuid bytes) -> serialized `Task<D>`
22/// - `queue`: `<u16 BE tag-len><tag utf8><uuid v7 bytes>` -> empty
23/// - `inflight`: task_id -> `<u16 BE tag-len><tag utf8>` (so recover knows the bucket)
24/// - `dlq`: task_id -> serialized `Task<D>`
25pub struct FjallRoundRobinTaskQueue<D, S>
26where
27    D: HasTagKey,
28    S: TaskSerializer,
29{
30    db: Database,
31    tasks: Keyspace,
32    queue: Keyspace,
33    inflight: Keyspace,
34    dlq: Keyspace,
35    state: Mutex<RotationState>,
36    _marker: PhantomData<(D, S)>,
37}
38
39#[derive(Debug, Default)]
40struct RotationState {
41    active: Vec<String>,
42    cursor: usize,
43    present: HashSet<String>,
44}
45
46impl RotationState {
47    fn ensure(&mut self, tag: &str) {
48        if self.present.insert(tag.to_string()) {
49            self.active.push(tag.to_string());
50        }
51    }
52}
53
54impl<D, S> FjallRoundRobinTaskQueue<D, S>
55where
56    D: HasTagKey,
57    S: TaskSerializer,
58{
59    pub fn open(path: impl AsRef<Path>) -> Result<Self, TaskQueueError> {
60        let db = Database::builder(path).open()?;
61        let tasks = db.keyspace("tasks", KeyspaceCreateOptions::default)?;
62        let queue = db.keyspace("queue", KeyspaceCreateOptions::default)?;
63        let inflight = db.keyspace("inflight", KeyspaceCreateOptions::default)?;
64        let dlq = db.keyspace("dlq", KeyspaceCreateOptions::default)?;
65
66        let state = Self::rebuild_state(&queue)?;
67
68        Ok(Self {
69            db,
70            tasks,
71            queue,
72            inflight,
73            dlq,
74            state: Mutex::new(state),
75            _marker: PhantomData,
76        })
77    }
78
79    fn rebuild_state(queue: &Keyspace) -> Result<RotationState, TaskQueueError> {
80        let mut present: HashSet<String> = HashSet::new();
81        let mut active: Vec<String> = Vec::new();
82        for item in queue.iter() {
83            let (key, _) = item.into_inner()?;
84            let tag = decode_tag_from_key(&key)?.to_string();
85            if present.insert(tag.clone()) {
86                active.push(tag);
87            }
88        }
89        // Deterministic order across restarts.
90        active.sort();
91        Ok(RotationState {
92            active,
93            cursor: 0,
94            present,
95        })
96    }
97
98    fn task_id_to_bytes(task_id: &TaskId) -> [u8; 16] {
99        *task_id.get().as_bytes()
100    }
101}
102
103#[async_trait]
104impl<D, S> TaskQueue<D> for FjallRoundRobinTaskQueue<D, S>
105where
106    D: HasTagKey
107        + std::fmt::Debug
108        + Clone
109        + Serialize
110        + DeserializeOwned
111        + Send
112        + Sync
113        + 'static,
114    S: TaskSerializer + Send + Sync,
115{
116    async fn push(&self, task: &Task<D>) -> Result<(), TaskQueueError> {
117        let tag = task.payload.get_tag_value().to_string();
118        let task_id_bytes = Self::task_id_to_bytes(&task.task_id);
119        let task_bytes = S::serialize_task(task)?;
120        let queue_key = encode_queue_key(&tag, &task_id_bytes);
121
122        let mut state = self.state.lock().unwrap();
123        self.tasks.insert(task_id_bytes, &task_bytes)?;
124        self.queue.insert(queue_key, &[] as &[u8])?;
125        self.inflight.remove(task_id_bytes)?;
126        self.db.persist(PersistMode::SyncAll)?;
127
128        state.ensure(&tag);
129        Ok(())
130    }
131
132    async fn pop(&self) -> Result<Task<D>, TaskQueueError> {
133        let mut state = self.state.lock().unwrap();
134
135        if state.active.is_empty() {
136            return Err(TaskQueueError::QueueEmpty);
137        }
138        let n = state.active.len();
139        for step in 0..n {
140            let idx = (state.cursor + step) % n;
141            let tag = state.active[idx].clone();
142            let prefix = encode_tag_prefix(&tag);
143
144            let Some(guard) = self.queue.prefix(&prefix).next() else {
145                // Transiently empty: this tag stays in rotation and will be
146                // checked again on the next pop, not evicted.
147                continue;
148            };
149            let (key, _) = guard.into_inner()?;
150
151            let task_id = extract_task_id_from_key(&key)?;
152            let task_id_bytes = Self::task_id_to_bytes(&task_id);
153
154            self.queue.remove(key.as_ref())?;
155            self.inflight.insert(task_id_bytes, encode_tag_blob(&tag))?;
156            self.db.persist(PersistMode::SyncAll)?;
157
158            let task_bytes = self
159                .tasks
160                .get(task_id_bytes)?
161                .ok_or(TaskQueueError::TaskNotFound(task_id))?;
162            let task = S::deserialize_task::<D>(&task_bytes)?;
163
164            state.cursor = (idx + 1) % n;
165            return Ok(task);
166        }
167        Err(TaskQueueError::QueueEmpty)
168    }
169
170    async fn ack(&self, task_id: &TaskId) -> Result<(), TaskQueueError> {
171        let _guard = self.state.lock().unwrap();
172        let task_id_bytes = Self::task_id_to_bytes(task_id);
173        self.tasks.remove(task_id_bytes)?;
174        self.inflight.remove(task_id_bytes)?;
175        self.db.persist(PersistMode::SyncAll)?;
176        Ok(())
177    }
178
179    async fn nack(&self, task: &Task<D>) -> Result<(), TaskQueueError> {
180        let _guard = self.state.lock().unwrap();
181        let task_id_bytes = Self::task_id_to_bytes(&task.task_id);
182        let task_bytes = S::serialize_task(task)?;
183
184        self.dlq.insert(task_id_bytes, &task_bytes)?;
185        self.tasks.remove(task_id_bytes)?;
186        self.inflight.remove(task_id_bytes)?;
187        self.db.persist(PersistMode::SyncAll)?;
188        Ok(())
189    }
190
191    async fn set(&self, task: &Task<D>) -> Result<(), TaskQueueError> {
192        let _guard = self.state.lock().unwrap();
193        let task_id_bytes = Self::task_id_to_bytes(&task.task_id);
194        let task_bytes = S::serialize_task(task)?;
195        self.tasks.insert(task_id_bytes, &task_bytes)?;
196        Ok(())
197    }
198
199    async fn recover_inflight(&self) -> Result<u64, TaskQueueError> {
200        let mut state = self.state.lock().unwrap();
201
202        let mut entries: Vec<([u8; 16], String)> = Vec::new();
203        for item in self.inflight.iter() {
204            let (key, value) = item.into_inner()?;
205            if key.len() != 16 {
206                return Err(TaskQueueError::QueueError(
207                    "inflight key is not a 16-byte task id".into(),
208                ));
209            }
210            let mut id = [0u8; 16];
211            id.copy_from_slice(&key);
212            let tag = decode_tag_blob(&value)?.to_string();
213            entries.push((id, tag));
214        }
215
216        let mut recovered = 0u64;
217        for (task_id_bytes, tag) in entries {
218            // Recover to the original FIFO position by reusing the task's
219            // uuid v7 suffix (oldest within its tag's sub-queue).
220            self.queue
221                .insert(encode_queue_key(&tag, &task_id_bytes), &[] as &[u8])?;
222            self.inflight.remove(&task_id_bytes[..])?;
223            state.ensure(&tag);
224            recovered += 1;
225        }
226
227        if recovered > 0 {
228            self.db.persist(PersistMode::SyncAll)?;
229        }
230        Ok(recovered)
231    }
232}
233
234impl<D, S> std::fmt::Debug for FjallRoundRobinTaskQueue<D, S>
235where
236    D: HasTagKey,
237    S: TaskSerializer,
238{
239    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
240        f.debug_struct("FjallRoundRobinTaskQueue").finish()
241    }
242}
243
244fn encode_queue_key(tag: &str, task_id_bytes: &[u8; 16]) -> Vec<u8> {
245    let tag_bytes = tag.as_bytes();
246    let len =
247        u16::try_from(tag_bytes.len()).expect("tag too long (max 65535 bytes)");
248    let mut key = Vec::with_capacity(2 + tag_bytes.len() + 16);
249    key.extend_from_slice(&len.to_be_bytes());
250    key.extend_from_slice(tag_bytes);
251    key.extend_from_slice(task_id_bytes);
252    key
253}
254
255fn encode_tag_prefix(tag: &str) -> Vec<u8> {
256    let tag_bytes = tag.as_bytes();
257    let len =
258        u16::try_from(tag_bytes.len()).expect("tag too long (max 65535 bytes)");
259    let mut prefix = Vec::with_capacity(2 + tag_bytes.len());
260    prefix.extend_from_slice(&len.to_be_bytes());
261    prefix.extend_from_slice(tag_bytes);
262    prefix
263}
264
265fn encode_tag_blob(tag: &str) -> Vec<u8> {
266    encode_tag_prefix(tag)
267}
268
269fn decode_tag_from_key(key: &[u8]) -> Result<&str, TaskQueueError> {
270    if key.len() < 2 {
271        return Err(TaskQueueError::QueueError("queue key too short".into()));
272    }
273    let len = u16::from_be_bytes([key[0], key[1]]) as usize;
274    if key.len() < 2 + len + 16 {
275        return Err(TaskQueueError::QueueError(
276            "queue key shorter than declared tag length".into(),
277        ));
278    }
279    std::str::from_utf8(&key[2..2 + len])
280        .map_err(|e| TaskQueueError::QueueError(format!("invalid tag utf8: {e}")))
281}
282
283fn decode_tag_blob(bytes: &[u8]) -> Result<&str, TaskQueueError> {
284    if bytes.len() < 2 {
285        return Err(TaskQueueError::QueueError("tag blob too short".into()));
286    }
287    let len = u16::from_be_bytes([bytes[0], bytes[1]]) as usize;
288    if bytes.len() != 2 + len {
289        return Err(TaskQueueError::QueueError(
290            "tag blob length mismatch".into(),
291        ));
292    }
293    std::str::from_utf8(&bytes[2..])
294        .map_err(|e| TaskQueueError::QueueError(format!("invalid tag utf8: {e}")))
295}
296
297fn extract_task_id_from_key(key: &[u8]) -> Result<TaskId, TaskQueueError> {
298    if key.len() < 18 {
299        return Err(TaskQueueError::QueueError(
300            "queue key missing task id suffix".into(),
301        ));
302    }
303    let suffix = &key[key.len() - 16..];
304    let uuid = Uuid::from_slice(suffix)
305        .map_err(|e| TaskQueueError::QueueError(format!("invalid uuid: {e}")))?;
306    Ok(TaskId::from_uuid(uuid))
307}
capp_queue/backend/fjall_round_robin.rs

capp_queue/backend/
fjall_round_robin.rs
