diff options
| -rw-r--r-- | femtoqueue.py | 14 | ||||
| -rw-r--r-- | test.py | 29 |
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 @@ -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() |