Use monotonic clock for task ordering

2 files changed, 102 insertions, 44 deletions

diff --git a/femtoqueue.py b/femtoqueue.py
index eb65684..718e824 100644
--- a/femtoqueue.py
+++ b/femtoqueue.py
@@ -18,6 +18,7 @@ class FemtoQueue:
         "pending",
         "done",
         "failed",
+        "scheduled",
     ]
 
     def __init__(
@@ -48,7 +49,7 @@ class FemtoQueue:
 
         assert timeout_stale_ms > 0
         self.timeout_stale_ms = timeout_stale_ms
-        self.latest_stale_check_ts: float | None = None
+        self.latest_stale_check_ts_us: int | None = None
 
         self.sync_after_write = sync_after_write
 
@@ -60,6 +61,7 @@ class FemtoQueue:
         self.dir_in_progress = path.join(data_dir, node_id)
         self.dir_done = path.join(data_dir, "done")
         self.dir_failed = path.join(data_dir, "failed")
+        self.dir_scheduled = path.join(data_dir, "scheduled")
 
         makedirs(self.data_dir, exist_ok=True)
         makedirs(self.dir_creating, exist_ok=True)
@@ -67,11 +69,35 @@ class FemtoQueue:
         makedirs(self.dir_in_progress, exist_ok=True)
         makedirs(self.dir_done, exist_ok=True)
         makedirs(self.dir_failed, exist_ok=True)
+        makedirs(self.dir_scheduled, exist_ok=True)
 
-    def _gen_increasing_uuid(self, time_us: int | None) -> str:
-        if not time_us:
-            time_us = int(1_000_000 * time.time())
+        self.monotonic_time_started_us: int = int(time.monotonic() * 1_000_000)
+        self.reference_time_us: int = self._resolve_reference_time_us()
 
+        self.interval_check_scheduled_us = 1_000_000
+        self.latest_scheduled_check_ts_us: int | None = None
+
+    def _resolve_reference_time_us(self) -> int:
+        tasks_pending = listdir(self.dir_pending)
+        newest_ts_us = None
+        for task_name in tasks_pending:
+            ts_us = int(task_name.split("_")[0])
+            if newest_ts_us is None or ts_us > newest_ts_us:
+                newest_ts_us = ts_us
+
+        if newest_ts_us is None:
+            return int(time.time() * 1_000_000)
+        else:
+            return newest_ts_us
+
+    def _monotonic_time_now_us(self) -> int:
+        return (
+            self.reference_time_us
+            + int(time.monotonic() * 1_000_000)
+            - self.monotonic_time_started_us
+        )
+
+    def _gen_increasing_uuid(self, time_us: int) -> str:
         rand_bytes = urandom(8)
         return f"{str(time_us)}_{rand_bytes.hex}"
 
@@ -107,17 +133,18 @@ class FemtoQueue:
 
         return data
 
-    def push(self, data: bytes, time_us: int | None = None) -> str:
+    def schedule(self, data: bytes, time_us: int) -> str:
         """
-        Push a new task into the queue.
+        Schedule a new task to be pushed to the queue at the given timestamp.
+        Upon calling `pop`, scheduled tasks whose timestamp is in the past (based on time-of-day clock, i.e. `time.time()`)
+        will be moved to the end of the queue. Note that these events will get a different ID when moved.
 
         Parameters
         ----------
         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.
+        time_us : int
+            A timestamp in microseconds.
 
         Returns
         -------
@@ -126,6 +153,34 @@ class FemtoQueue:
         """
         id = self._gen_increasing_uuid(time_us)
         creating_path = path.join(self.dir_creating, id)
+        scheduled_path = path.join(self.dir_scheduled, id)
+
+        with open(creating_path, "wb") as f:
+            self._write_v1(f, data)
+            if self.sync_after_write:
+                fsync(f)
+
+        rename(creating_path, scheduled_path)
+
+        return id
+
+    def push(self, data: bytes) -> str:
+        """
+        Push a new task into the queue.
+
+        Parameters
+        ----------
+        data : bytes
+            A bytes object representing the task. For JSON, you can use `json.dumps(obj).encode("utf-8")`.
+
+        Returns
+        -------
+        id : str
+            The task identifier, i.e. the file name.
+        """
+        m_now_us = self._monotonic_time_now_us()
+        id = self._gen_increasing_uuid(m_now_us)
+        creating_path = path.join(self.dir_creating, id)
         pending_path = path.join(self.dir_pending, id)
 
         with open(creating_path, "wb") as f:
@@ -137,19 +192,17 @@ class FemtoQueue:
 
         return id
 
-    def _release_stale_tasks(self):
-        now = time.time()
-
+    def _release_stale_tasks(self, m_now_us: int):
         # Only run this every `timeout_stale_ms` milliseconds because iterating
         # through all tasks is slow
-        timeout_sec = self.timeout_stale_ms / 1000.0
+        timeout_us = self.timeout_stale_ms * 1_000
         if (
-            self.latest_stale_check_ts is not None
-            and now - self.latest_stale_check_ts < timeout_sec
+            self.latest_stale_check_ts_us is not None
+            and m_now_us - self.latest_stale_check_ts_us < timeout_us
         ):
             return
 
-        self.latest_stale_check_ts = now
+        self.latest_stale_check_ts_us = m_now_us
 
         for dir_name in listdir(self.data_dir):
             full_dir_path = path.join(self.data_dir, dir_name)
@@ -164,9 +217,8 @@ class FemtoQueue:
             for task_file in listdir(full_dir_path):
                 task_path = path.join(full_dir_path, task_file)
                 modified_time_us = int(task_file.split("_")[0])
-                modified_time = modified_time_us / 1_000_000.0
 
-                if now - modified_time < timeout_sec:
+                if m_now_us - modified_time_us < timeout_us:
                     continue
 
                 try:
@@ -175,12 +227,31 @@ class FemtoQueue:
                 except FileNotFoundError:
                     continue  # Task may have been moved by another node
 
-    def _pop_task_path(self) -> str | None:
-        now_us = time.time() * 1_000_000
+    def _trigger_scheduled_tasks(self, m_now_us: int):
+        if (
+            self.latest_scheduled_check_ts_us is not None
+            and m_now_us - self.latest_scheduled_check_ts_us
+            < self.interval_check_scheduled_us
+        ):
+            return
+
+        self.latest_scheduled_check_ts_us = m_now_us
+
+        w_now_us = int(time.time() * 1_000_000)  # wall clock time
 
-        def _only_past(task_name: str) -> bool:
-            return int(task_name.split("_")[0]) <= now_us
+        for task_file in listdir(self.dir_scheduled):
+            scheduled_time_us = int(task_file.split("_")[0])
 
+            if scheduled_time_us < w_now_us:
+                scheduled_path = path.join(self.dir_scheduled, task_file)
+                id = self._gen_increasing_uuid(m_now_us)
+                try:
+                    pending_path = path.join(self.dir_pending, id)
+                    rename(scheduled_path, pending_path)
+                except FileNotFoundError:
+                    continue  # Task may have been moved by another node
+
+    def _pop_task_path(self) -> str | None:
         # Check cache
         if self.todo_cache:
             try:
@@ -191,7 +262,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(filter(_only_past, listdir(self.dir_in_progress)))
+            for x in sorted(listdir(self.dir_in_progress))
         )
         try:
             return next(self.todo_cache)
@@ -200,8 +271,7 @@ class FemtoQueue:
 
         # Then check pending tasks
         self.todo_cache = (
-            path.join(self.dir_pending, x)
-            for x in sorted(filter(_only_past, listdir(self.dir_pending)))
+            path.join(self.dir_pending, x) for x in sorted(listdir(self.dir_pending))
         )
         try:
             return next(self.todo_cache)
@@ -219,7 +289,9 @@ class FemtoQueue:
         -------
         task : FemtoTask or None
         """
-        self._release_stale_tasks()
+        m_now_us = self._monotonic_time_now_us()
+        self._release_stale_tasks(m_now_us)
+        self._trigger_scheduled_tasks(m_now_us)
 
         while True:
             task = self._pop_task_path()
diff --git a/test.py b/test.py
index 144d9ce..13db3c1 100644
--- a/test.py
+++ b/test.py
@@ -7,6 +7,7 @@ from femtoqueue import FemtoQueue, FemtoTask
 from tempfile import mkdtemp
 
 original_time_fn = time.time
+original_monotonic_fn = time.monotonic
 
 
 def set_time_mock(ts: float):
@@ -14,10 +15,12 @@ def set_time_mock(ts: float):
         return ts
 
     time.time = mock_time_fn
+    time.monotonic = mock_time_fn
 
 
 def reset_time_mock():
     time.time = original_time_fn
+    time.monotonic = original_monotonic_fn
 
 
 class TestFemtoQueue(unittest.TestCase):
@@ -170,30 +173,13 @@ 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)
+        q.schedule("foobar".encode("utf-8"), time_us=future_ts_us)
 
         task = q.pop()
         self.assertIsNone(task)