Add initial impl

5 files changed, 350 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..24a346b
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,3 @@
+*.pyc
+venv/
+profile.*
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..0310e64
--- /dev/null
+++ b/README.md
@@ -0,0 +1,51 @@
+# FemtoQueue
+
+Ever wanted a zero-dependency, filesystem-backed, durable, concurrent, retrying task queue implementation? No?
+
+## Example
+
+```python
+from femtoqueue import FemtoQueue, FemtoTask
+
+q = FemtoQueue(data_dir = "fq", node_id = "node1")
+q.push("foobar".encode("utf-8"))
+
+while task := q.pop():
+    # Do something with `task.data`
+    q.done(task) # or q.fail(task)
+
+print("All tasks processed")
+```
+
+## Installation
+
+Just chuck the `femtoqueue.py` file into your Python 3 project. There are no dependencies other than the standard library.
+
+## Features
+
+This mini-library provides the `FemtoQueue` class with the standard queue interface:
+
+| Method                         | Description                         |
+| :----------------------------- | :---------------------------------- |
+| `push(task: FemtoTask) -> str` | Add a task to the queue, returns id |
+| `pop() -> FemtoTask`           | Get a task from the queue           |
+
+Each task corresponds to one file in the `data_dir` directory. State changes are atomic since they use `mv` (or its Python equivalent `os.rename`).
+
+Each concurrent worker node (library user) must have a stable identifier `node_id`. This way workers can automatically retry a task if they unexpectedly crash in the middle of processing.
+
+Stale tasks (i.e. in progress for too long) are moved back to `pending` automatically when a timeout is reached (default: 30s).
+
+## But isn't this slow?
+
+I wouldn't migrate away from your production queue system just yet, but this is faster than you'd expect. Easily fast enough for some small or medium project. Turns out, creating and renaming files is pretty snappy.
+
+## Unit tests
+
+```bash
+python test.py
+```
+
+## Author and license
+
+Jan Tuomi <<jan@jantuomi.fi>>. Licensed under AGPL 3.0. All rights reserved.
diff --git a/benchmark.py b/benchmark.py
new file mode 100644
index 0000000..a28dc0d
--- /dev/null
+++ b/benchmark.py
@@ -0,0 +1,34 @@
+import time
+import tempfile
+import shutil
+from femtoqueue import FemtoQueue, FemtoTask
+
+def benchmark_femtoqueue(num_tasks: int = 1000):
+    tmpdir = tempfile.mkdtemp()
+    queue = FemtoQueue(data_dir=tmpdir, node_name="node1")
+
+    data = b"x" * 100  # 100-byte payload
+
+    print(f"Pushing {num_tasks} tasks...")
+    start = time.time()
+    for _ in range(num_tasks):
+        queue.push(data)
+    push_duration = time.time() - start
+    print(f"Pushed in {push_duration:.4f}s ({num_tasks / push_duration:.2f} tasks/sec)")
+
+    print("Processing tasks (pop + done)...")
+    start = time.time()
+    processed = 0
+    while True:
+        task = queue.pop()
+        if not task:
+            break
+        queue.done(task)
+        processed += 1
+    process_duration = time.time() - start
+    print(f"Processed in {process_duration:.4f}s ({processed / process_duration:.2f} tasks/sec)")
+
+    shutil.rmtree(tmpdir)
+
+if __name__ == "__main__":
+    benchmark_femtoqueue(500)
diff --git a/femtoqueue.py b/femtoqueue.py
new file mode 100644
index 0000000..4003336
--- /dev/null
+++ b/femtoqueue.py
@@ -0,0 +1,134 @@
+from os import makedirs, path, listdir, rename
+from dataclasses import dataclass
+from uuid import uuid4
+from time import time
+
+@dataclass
+class FemtoTask:
+    id: str
+    data: bytes
+
+class FemtoQueue:
+    def __init__(
+        self,
+        data_dir: str,
+        node_name: str,
+        timeout_stale_ms: int = 30_000,
+    ):
+        assert node_name != "pending" \
+           and node_name != "done" \
+           and node_name != "failed"
+        self.node_name = node_name
+
+        assert timeout_stale_ms > 0
+        self.timeout_stale_ms = timeout_stale_ms
+        self.latest_stale_check_ts: float | None = None
+
+        self.data_dir = data_dir
+        self.dir_pending = path.join(data_dir, "pending")
+        self.dir_in_progress = path.join(data_dir, node_name)
+        self.dir_done = path.join(data_dir, "done")
+        self.dir_failed = path.join(data_dir, "failed")
+
+        makedirs(self.data_dir, exist_ok=True)
+        makedirs(self.dir_pending, exist_ok=True)
+        makedirs(self.dir_in_progress, exist_ok=True)
+        makedirs(self.dir_done, exist_ok=True)
+        makedirs(self.dir_failed, exist_ok=True)
+
+    def push(self, data: bytes) -> str:
+        id = uuid4().hex
+        pending_path = path.join(self.dir_pending, id)
+
+        with open(pending_path, "wb") as f:
+            f.write(data)
+
+        return id
+
+    def _release_stale_tasks(self):
+        now = time()
+
+        # Only run this every `timeout_stale_ms` milliseconds because iterating
+        # through all tasks is slow
+        timeout_sec = self.timeout_stale_ms / 1000.0
+        if self.latest_stale_check_ts is not None and now - self.latest_stale_check_ts < timeout_sec:
+            return
+
+        for dir_name in listdir(self.data_dir):
+            full_dir_path = path.join(self.data_dir, dir_name)
+
+            # Skip non-directories and reserved names
+            if not path.isdir(full_dir_path):
+                continue
+            if dir_name in ("pending", "done", "failed", self.node_name):
+                continue
+
+            # Check tasks in this node's in-progress directory
+            for task_file in listdir(full_dir_path):
+                task_path = path.join(full_dir_path, task_file)
+                try:
+                    modified_time = path.getmtime(task_path)
+                except FileNotFoundError:
+                    continue  # Task may have been moved concurrently
+
+                if now - modified_time < timeout_sec:
+                    continue
+
+                try:
+                    pending_path = path.join(self.dir_pending, task_file)
+                    rename(task_path, pending_path)
+                except FileNotFoundError:
+                    continue  # Task may have been moved by another node
+
+    def _pop_task_path(self) -> str | None:
+        # First check assigned tasks in progress
+        tasks = [path.join(self.dir_in_progress, p) for p in listdir(self.dir_in_progress)]
+        #tasks.sort(key=lambda x: path.getmtime(x)) # oldest first
+        if len(tasks) > 0:
+            return tasks[0]
+
+        # Then check pending tasks
+        tasks = [path.join(self.dir_pending, p) for p in listdir(self.dir_pending)]
+        #tasks.sort(key=lambda x: path.getmtime(x)) # oldest first
+        if len(tasks) > 0:
+            return tasks[0]
+
+        return None
+
+    def pop(self) -> FemtoTask | None:
+        #self._release_stale_tasks()
+
+        while True:
+            task = self._pop_task_path()
+            if task is None: return None
+
+            id = path.basename(task)
+            in_progress_path = path.join(self.dir_in_progress, id)
+
+            try:
+                rename(task, in_progress_path)
+            except FileNotFoundError:
+                # If another node grabbed the task, just get another one
+                continue
+
+            with open(in_progress_path, "rb") as f:
+                content = f.read()
+                return FemtoTask(id = id, data = content)
+
+    def done(self, task: FemtoTask):
+        in_progress_path = path.join(self.dir_in_progress, task.id)
+        done_path = path.join(self.dir_done, task.id)
+
+        try:
+            rename(in_progress_path, done_path)
+        except FileNotFoundError as e:
+            raise Exception(f"Tried to complete a task that is not in progress, id={task.id}") from e
+
+    def fail(self, task: FemtoTask):
+        in_progress_path = path.join(self.dir_in_progress, task.id)
+        failed_path = path.join(self.dir_failed, task.id)
+
+        try:
+            rename(in_progress_path, failed_path)
+        except FileNotFoundError as e:
+            raise Exception(f"Tried to fail a task that is not in progress, id={task.id}") from e
diff --git a/test.py b/test.py
new file mode 100644
index 0000000..127e53f
--- /dev/null
+++ b/test.py
@@ -0,0 +1,128 @@
+from typing import cast
+import os
+import time
+import unittest
+import json
+from femtoqueue import FemtoQueue, FemtoTask
+from tempfile import mkdtemp
+
+class TestFemtoQueue(unittest.TestCase):
+    def test_basic(self):
+        dir = mkdtemp()
+        q = FemtoQueue(data_dir = dir, node_name = "node1")
+
+        # Add a JSON task
+        some_data = json.dumps({ "foo": "bar" })
+        q.push(some_data.encode("utf-8"))
+
+        # The task should be in the queue with the correct payload
+        task = q.pop()
+        self.assertIsNotNone(task)
+        task = cast(FemtoTask, task)
+        parsed_data = json.loads(task.data.decode("utf-8"))
+        self.assertEqual(parsed_data["foo"], "bar")
+
+        # Mark the task as done
+        q.done(task)
+
+        # There should be no tasks available
+        task = q.pop()
+        self.assertIsNone(task)
+
+    def test_aborted(self):
+        dir = mkdtemp()
+        q = FemtoQueue(data_dir = dir, node_name = "node1")
+
+        # Add a JSON task
+        some_data = json.dumps({ "foo": "bar" })
+        q.push(some_data.encode("utf-8"))
+
+        # The task should be in the queue with the correct payload
+        task = q.pop()
+        self.assertIsNotNone(task)
+        task = cast(FemtoTask, task)
+        parsed_data = json.loads(task.data.decode("utf-8"))
+        self.assertEqual(parsed_data["foo"], "bar")
+
+        # Simulate a fault and start over
+        q = FemtoQueue(data_dir = dir, node_name = "node1")
+
+        # The task should still be assigned to this node
+        task = q.pop()
+        self.assertIsNotNone(task)
+        task = cast(FemtoTask, task)
+        parsed_data = json.loads(task.data.decode("utf-8"))
+        self.assertEqual(parsed_data["foo"], "bar")
+
+        # Mark the task as done
+        q.done(task)
+
+        # There should be no tasks available
+        task = q.pop()
+        self.assertIsNone(task)
+
+    def test_release_stale_tasks(self):
+        dir = mkdtemp()
+
+        # Node1 creates and claims the task
+        q1 = FemtoQueue(data_dir=dir, node_name="node1", timeout_stale_ms=100)
+        q1.push(b"stuck")
+        task = q1.pop()
+        self.assertIsNotNone(task)
+        task = cast(FemtoTask, task)
+
+        # Simulate the task becoming stale by changing mtime
+        task_path = os.path.join(dir, "node1", task.id)
+        old_time = time.time() - 9999
+        os.utime(task_path, (old_time, old_time))
+
+        # Now node2 should see the task as stale and reclaim it
+        q2 = FemtoQueue(data_dir=dir, node_name="node2", timeout_stale_ms=100)
+        revived_task = q2.pop()
+        self.assertIsNotNone(revived_task)
+        revived_task = cast(FemtoTask, revived_task)
+        self.assertEqual(revived_task.data, b"stuck")
+        self.assertEqual(revived_task.id, task.id)
+
+    def test_mark_task_failed(self):
+        dir = mkdtemp()
+        q = FemtoQueue(data_dir=dir, node_name="node1")
+
+        # Push and pop a task
+        q.push(b"will fail")
+        task = q.pop()
+        self.assertIsNotNone(task)
+        task = cast(FemtoTask, task)
+
+        # Mark the task as failed
+        q.fail(task)
+
+        # Ensure the file now exists in the failed directory
+        failed_path = os.path.join(dir, "failed", task.id)
+        self.assertTrue(os.path.exists(failed_path))
+
+        # The task should not reappear in pop
+        self.assertIsNone(q.pop())
+
+    def test_mark_task_done(self):
+        dir = mkdtemp()
+        q = FemtoQueue(data_dir=dir, node_name="node1")
+
+        # Push and pop a task
+        q.push(b"complete me")
+        task = q.pop()
+        self.assertIsNotNone(task)
+        task = cast(FemtoTask, task)
+
+        # Mark the task as done
+        q.done(task)
+
+        # Ensure the file is now in the 'done' directory
+        done_path = os.path.join(dir, "done", task.id)
+        self.assertTrue(os.path.exists(done_path))
+
+        # The task should not be returned again
+        self.assertIsNone(q.pop())
+
+if __name__ == '__main__':
+    unittest.main()