move to idle instead of offline (#40)

* move to idle instead of offline

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: PythonFZ

laufband/graphband.py

@@ -218,8 +218,10 @@ def _register_worker(self):
                 session.commit()
 
     def __del__(self):
-        if hasattr(self, "_thread_event"):
+        if hasattr(self, "_thread_event") and hasattr(self, "_heartbeat_thread"):
             self._thread_event.set()
+            if self._heartbeat_thread.is_alive():
+                self._heartbeat_thread.join()
 
     def close(self):
         """Exit out of the graphband generator.
@@ -536,5 +538,3 @@ def __iter__(self) -> Iterator[Task[TaskTypeVar]]:
                 break
         if completed_naturally:
             self._iterator_completed = True
-        self._thread_event.set()
-        self._heartbeat_thread.join()

tests/test_graphband.py

@@ -91,7 +91,7 @@ def test_graphband_sequential_success(tmp_path):
     with Session(pbar._engine) as session:
         workers = session.query(WorkerEntry).all()
         assert len(workers) == 1
-        assert workers[0].status == WorkerStatus.OFFLINE
+        assert workers[0].status == WorkerStatus.IDLE
         tasks = session.query(TaskEntry).all()
         assert len(tasks) == 10
         for id, task in enumerate(tasks):
@@ -102,6 +102,15 @@ def test_graphband_sequential_success(tmp_path):
     # if we no iterate again, we yield nothing
     assert list(pbar) == []
 
+    # Test that worker goes offline when garbage collected
+    engine = pbar._engine
+    del pbar
+
+    with Session(engine) as session:
+        workers = session.query(WorkerEntry).all()
+        assert len(workers) == 1
+        assert workers[0].status == WorkerStatus.OFFLINE
+
     pbar = Graphband(
         sequential_task(),
         db=f"sqlite:///{tmp_path}/graphband.sqlite",
@@ -732,3 +741,78 @@ def test_has_more_jobs_with_killed_workers(tmp_path):
     assert retries_worker.has_more_jobs is True
     assert len(list(retries_worker)) == 1
     assert retries_worker.has_more_jobs is False
+
+
+def test_resume_worker(tmp_path):
+    lock_path = f"{tmp_path}/graphband.lock"
+    db_path = f"sqlite:///{tmp_path}/graphband.sqlite"
+
+    worker = Graphband(
+        sequential_task(),
+        db=db_path,
+        lock=Lock(lock_path),
+        heartbeat_timeout=2,
+        heartbeat_interval=1,
+        identifier="worker",
+    )
+    engine = create_engine(worker.db)
+    length = 0
+
+    with Session(engine) as session:
+        worker_entry = session.get(WorkerEntry, "worker")
+        assert worker_entry is not None
+        assert worker_entry.status == WorkerStatus.IDLE
+
+    for item in worker:
+        with Session(engine) as session:
+            worker_entry = session.get(WorkerEntry, "worker")
+            assert worker_entry is not None
+            assert worker_entry.status == WorkerStatus.BUSY
+        length += 1
+        if item.id == "task_5":
+            break
+    assert length == 6
+
+    with Session(engine) as session:
+        worker_entry = session.get(WorkerEntry, "worker")
+        assert worker_entry is not None
+        assert worker_entry.status == WorkerStatus.IDLE
+
+    for item in worker:
+        with Session(engine) as session:
+            worker_entry = session.get(WorkerEntry, "worker")
+            assert worker_entry is not None
+            assert worker_entry.status == WorkerStatus.BUSY
+        length += 1
+
+    assert length == 10
+
+    with Session(engine) as session:
+        worker_entry = session.get(WorkerEntry, "worker")
+        assert worker_entry is not None
+        assert worker_entry.status == WorkerStatus.IDLE
+
+    del worker
+
+    with Session(engine) as session:
+        worker_entry = session.get(WorkerEntry, "worker")
+        assert worker_entry is not None
+        assert worker_entry.status == WorkerStatus.OFFLINE
+
+    proc = multiprocessing.Process(
+        target=task_worker,
+        args=(sequential_task, lock_path, db_path, tmp_path / "test.txt", 0.1),
+        kwargs={
+            "heartbeat_timeout": 1,
+            "heartbeat_interval": 0.5,
+            "max_killed_retries": 0,  # No retries allowed for killed tasks
+            "identifier": "killed-worker",
+        },
+    )
+    proc.start()
+    proc.join(timeout=5)
+
+    with Session(engine) as session:
+        worker_entry = session.get(WorkerEntry, "killed-worker")
+        assert worker_entry is not None
+        assert worker_entry.status == WorkerStatus.OFFLINE