From 3de71d936f28a305e18bdc98bf136b1d2c643eea Mon Sep 17 00:00:00 2001
From: Jan Tuomi <jan@jantuomi.fi>
Date: Thu, 15 May 2025 13:01:52 +0300
Subject: Only pop() tasks added with past timestamp

---
 femtoqueue.py | 25 ++++++++++++++++++-------
 test.py       | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/femtoqueue.py b/femtoqueue.py
index faa6a5b..500115b 100644
--- a/femtoqueue.py
+++ b/femtoqueue.py
@@ -65,12 +65,14 @@ class FemtoQueue:
         makedirs(self.dir_done, exist_ok=True)
         makedirs(self.dir_failed, exist_ok=True)
 
-    def _gen_increasing_uuid(self) -> str:
-        now_us = 1_000_000 * time.time()
+    def _gen_increasing_uuid(self, time_us: int | None) -> str:
+        if not time_us:
+            time_us = int(1_000_000 * time.time())
+
         rand_bytes = urandom(8)
-        return f"{str(int(now_us))}_{rand_bytes.hex}"
+        return f"{str(time_us)}_{rand_bytes.hex}"
 
-    def push(self, data: bytes) -> str:
+    def push(self, data: bytes, time_us: int | None = None) -> str:
         """
         Push a new task into the queue.
 
@@ -78,13 +80,16 @@ class FemtoQueue:
         ----------
         data : bytes
             A bytes object representing the task. For JSON, you can use `json.dumps(obj).encode("utf-8")`.
+        time_us : int or None
+            A timestamp in microseconds. If not None, the current time is used. Only past tasks are available with `pop()`.
+            Future timestamps can be used to schedule tasks.
 
         Returns
         -------
         id : str
             The task identifier, i.e. the file name.
         """
-        id = self._gen_increasing_uuid()
+        id = self._gen_increasing_uuid(time_us)
         creating_path = path.join(self.dir_creating, id)
         pending_path = path.join(self.dir_pending, id)
 
@@ -136,6 +141,11 @@ class FemtoQueue:
                     continue  # Task may have been moved by another node
 
     def _pop_task_path(self) -> str | None:
+        now_us = time.time() * 1_000_000
+
+        def _only_past(task_name: str) -> bool:
+            return int(task_name.split("_")[0]) <= now_us
+
         # Check cache
         if self.todo_cache:
             try:
@@ -146,7 +156,7 @@ class FemtoQueue:
         # If cache empty, then check assigned tasks in progress (aborted)
         self.todo_cache = (
             path.join(self.dir_in_progress, x)
-            for x in sorted(listdir(self.dir_in_progress))
+            for x in sorted(filter(_only_past, listdir(self.dir_in_progress)))
         )
         try:
             return next(self.todo_cache)
@@ -155,7 +165,8 @@ class FemtoQueue:
 
         # Then check pending tasks
         self.todo_cache = (
-            path.join(self.dir_pending, x) for x in sorted(listdir(self.dir_pending))
+            path.join(self.dir_pending, x)
+            for x in sorted(filter(_only_past, listdir(self.dir_pending)))
         )
         try:
             return next(self.todo_cache)
diff --git a/test.py b/test.py
index 2ae9c48..144d9ce 100644
--- a/test.py
+++ b/test.py
@@ -170,6 +170,40 @@ class TestFemtoQueue(unittest.TestCase):
         # Queue should now be empty
         self.assertIsNone(q.pop())
 
+    def test_only_past_tasks_are_popped(self):
+        dir = mkdtemp()
+        q = FemtoQueue(data_dir=dir, node_id="node1")
+
+        set_time_mock(1000)
+        q.push("foobar".encode("utf-8"))
+
+        set_time_mock(0)
+        task = q.pop()
+        self.assertIsNone(task)
+
+        set_time_mock(2000)
+        task = q.pop()
+        self.assertIsNotNone(task)
+
+        reset_time_mock()
+
+    def test_schedule_for_future(self):
+        dir = mkdtemp()
+        q = FemtoQueue(data_dir=dir, node_id="node1")
+
+        future_ts = time.time() + 60
+        future_ts_us = int(1_000_000 * future_ts)
+        q.push("foobar".encode("utf-8"), time_us=future_ts_us)
+
+        task = q.pop()
+        self.assertIsNone(task)
+
+        set_time_mock(future_ts + 60)
+        task = q.pop()
+        self.assertIsNotNone(task)
+
+        reset_time_mock()
+
 
 if __name__ == "__main__":
     unittest.main()
-- 
cgit v1.3