Ensure ordering at microsecond resolution

2 files changed, 30 insertions, 13 deletions

diff --git a/femtoqueue.py b/femtoqueue.py
index e7ffcf8..500089c 100644
--- a/femtoqueue.py
+++ b/femtoqueue.py
@@ -1,7 +1,7 @@
 from os import makedirs, path, listdir, rename
 from dataclasses import dataclass
 from uuid import uuid4
-from time import time
+import time
 
 @dataclass
 class FemtoTask:
@@ -37,7 +37,8 @@ class FemtoQueue:
         makedirs(self.dir_failed, exist_ok=True)
 
     def push(self, data: bytes) -> str:
-        id = f"{uuid4().hex[-12:]}_{str(int(time()))}"
+        now_us = 1_000_000 * time.time()
+        id = f"{str(int(now_us))}_{uuid4().hex[-12:]}"
         pending_path = path.join(self.dir_pending, id)
 
         with open(pending_path, "wb") as f:
@@ -46,7 +47,7 @@ class FemtoQueue:
         return id
 
     def _release_stale_tasks(self):
-        now = time()
+        now = time.time()
 
         # Only run this every `timeout_stale_ms` milliseconds because iterating
         # through all tasks is slow
@@ -67,7 +68,8 @@ class FemtoQueue:
             # 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)
-                modified_time = int(task_file[13:])
+                modified_time_us = int(task_file.split("_")[0])
+                modified_time = modified_time_us / 1_000_000.0
 
                 if now - modified_time < timeout_sec:
                     continue
@@ -82,13 +84,13 @@ class FemtoQueue:
         # First check assigned tasks in progress
         tasks = listdir(self.dir_in_progress)
         if len(tasks) > 0:
-            id = min(tasks, key=lambda x: x[13:])
+            id = min(tasks)
             return path.join(self.dir_in_progress, id)
 
         # Then check pending tasks
         tasks = listdir(self.dir_pending)
         if len(tasks) > 0:
-            id = min(tasks, key=lambda x: x[13:])
+            id = min(tasks)
             return path.join(self.dir_pending, id)
 
         return None
diff --git a/test.py b/test.py
index 127e53f..23ae883 100644
--- a/test.py
+++ b/test.py
@@ -6,6 +6,16 @@ import json
 from femtoqueue import FemtoQueue, FemtoTask
 from tempfile import mkdtemp
 
+original_time_fn = time.time
+def set_time_mock(ts: float):
+    def mock_time_fn():
+        return ts
+
+    time.time = mock_time_fn
+
+def reset_time_mock():
+    time.time = original_time_fn
+
 class TestFemtoQueue(unittest.TestCase):
     def test_basic(self):
         dir = mkdtemp()
@@ -63,26 +73,31 @@ class TestFemtoQueue(unittest.TestCase):
 
     def test_release_stale_tasks(self):
         dir = mkdtemp()
+        timeout_stale_ms = 100
 
+        set_time_mock(0)
         # Node1 creates and claims the task
-        q1 = FemtoQueue(data_dir=dir, node_name="node1", timeout_stale_ms=100)
+        q1 = FemtoQueue(data_dir=dir, node_name="node1", timeout_stale_ms=timeout_stale_ms)
         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))
+        # Assert that node2 can not see the task
+        q2 = FemtoQueue(data_dir=dir, node_name="node2", timeout_stale_ms=timeout_stale_ms)
+        non_existant_task = q2.pop()
+        self.assertIsNone(non_existant_task)
+
+        # Simulate the task becoming stale by skipping forward in time
+        set_time_mock(10)
 
         # 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)
+
+        reset_time_mock()
 
     def test_mark_task_failed(self):
         dir = mkdtemp()