Fix Deactivate (#327) (#328)

Co-authored-by: Peter Kraft <[email protected]>

Author: qianl15

dbos/_admin_server.py

@@ -66,11 +66,11 @@ def do_GET(self) -> None:
         elif self.path == _deactivate_path:
             if not AdminRequestHandler.is_deactivated:
                 dbos_logger.info(
-                    f"Deactivating DBOS executor {GlobalParams.executor_id} with version {GlobalParams.app_version}. This executor will complete existing workflows but will not start new workflows."
+                    f"Deactivating DBOS executor {GlobalParams.executor_id} with version {GlobalParams.app_version}. This executor will complete existing workflows but will not create new workflows."
                 )
                 AdminRequestHandler.is_deactivated = True
-            # Stop all scheduled workflows, queues, and kafka loops
-            for event in self.dbos.stop_events:
+            # Stop all event receivers (scheduler and Kafka threads)
+            for event in self.dbos.poller_stop_events:
                 event.set()
             self.send_response(200)
             self._end_headers()

dbos/_core.py

@@ -296,7 +296,7 @@ def _init_workflow(
 
     if workflow_deadline_epoch_ms is not None:
         evt = threading.Event()
-        dbos.stop_events.append(evt)
+        dbos.background_thread_stop_events.append(evt)
 
         def timeout_func() -> None:
             try:

dbos/_dbos.py

@@ -197,7 +197,7 @@ def register_poller(
         self, evt: threading.Event, func: Callable[..., Any], *args: Any, **kwargs: Any
     ) -> None:
         if self.dbos and self.dbos._launched:
-            self.dbos.stop_events.append(evt)
+            self.dbos.poller_stop_events.append(evt)
             self.dbos._executor.submit(func, *args, **kwargs)
         else:
             self.pollers.append((evt, func, args, kwargs))
@@ -330,7 +330,10 @@ def __init__(
         self._registry: DBOSRegistry = _get_or_create_dbos_registry()
         self._registry.dbos = self
         self._admin_server_field: Optional[AdminServer] = None
-        self.stop_events: List[threading.Event] = []
+        # Stop internal background threads (queue thread, timeout threads, etc.)
+        self.background_thread_stop_events: List[threading.Event] = []
+        # Stop pollers (event receivers) that can create new workflows (scheduler, Kafka)
+        self.poller_stop_events: List[threading.Event] = []
         self.fastapi: Optional["FastAPI"] = fastapi
         self.flask: Optional["Flask"] = flask
         self._executor_field: Optional[ThreadPoolExecutor] = None
@@ -502,7 +505,7 @@ def _launch(self, *, debug_mode: bool = False) -> None:
 
             # Start the queue thread
             evt = threading.Event()
-            self.stop_events.append(evt)
+            self.background_thread_stop_events.append(evt)
             bg_queue_thread = threading.Thread(
                 target=queue_thread, args=(evt, self), daemon=True
             )
@@ -515,7 +518,7 @@ def _launch(self, *, debug_mode: bool = False) -> None:
                     dbos_domain = os.environ.get("DBOS_DOMAIN", "cloud.dbos.dev")
                     self.conductor_url = f"wss://{dbos_domain}/conductor/v1alpha1"
                 evt = threading.Event()
-                self.stop_events.append(evt)
+                self.background_thread_stop_events.append(evt)
                 self.conductor_websocket = ConductorWebsocket(
                     self,
                     conductor_url=self.conductor_url,
@@ -527,7 +530,7 @@ def _launch(self, *, debug_mode: bool = False) -> None:
 
             # Grab any pollers that were deferred and start them
             for evt, func, args, kwargs in self._registry.pollers:
-                self.stop_events.append(evt)
+                self.poller_stop_events.append(evt)
                 poller_thread = threading.Thread(
                     target=func, args=args, kwargs=kwargs, daemon=True
                 )
@@ -583,7 +586,9 @@ def _reset_system_database(self) -> None:
 
     def _destroy(self) -> None:
         self._initialized = False
-        for event in self.stop_events:
+        for event in self.poller_stop_events:
+            event.set()
+        for event in self.background_thread_stop_events:
             event.set()
         self._background_event_loop.stop()
         if self._sys_db_field is not None:

dbos/_recovery.py

@@ -29,7 +29,7 @@ def startup_recovery_thread(
 ) -> None:
     """Attempt to recover local pending workflows on startup using a background thread."""
     stop_event = threading.Event()
-    dbos.stop_events.append(stop_event)
+    dbos.background_thread_stop_events.append(stop_event)
     while not stop_event.is_set() and len(pending_workflows) > 0:
         try:
             for pending_workflow in list(pending_workflows):

tests/test_admin_server.py

@@ -3,6 +3,7 @@
 import threading
 import time
 import uuid
+from datetime import datetime
 
 import pytest
 import requests
@@ -65,8 +66,59 @@ def test_admin_endpoints(dbos: DBOS) -> None:
     response = requests.get("http://localhost:3001/deactivate", timeout=5)
     assert response.status_code == 200
 
-    for event in dbos.stop_events:
-        assert event.is_set(), "Event is not set!"
+    for event in dbos.poller_stop_events:
+        assert event.is_set()
+
+
+def test_deactivate(dbos: DBOS, config: ConfigFile) -> None:
+    wf_counter: int = 0
+
+    queue = Queue("example-queue")
+
+    @DBOS.scheduled("* * * * * *")
+    @DBOS.workflow()
+    def test_workflow(scheduled: datetime, actual: datetime) -> None:
+        nonlocal wf_counter
+        wf_counter += 1
+
+    @DBOS.workflow()
+    def regular_workflow() -> int:
+        return 5
+
+    # Let the scheduled workflow run
+    time.sleep(2)
+    val = wf_counter
+    assert val > 0
+    # Deactivate--scheduled workflow should stop
+    response = requests.get("http://localhost:3001/deactivate", timeout=5)
+    assert response.status_code == 200
+    for event in dbos.poller_stop_events:
+        assert event.is_set()
+    # Verify the scheduled workflow does not run anymore
+    time.sleep(3)
+    assert wf_counter <= val + 1
+    # Enqueue a workflow, verify it still runs
+    assert queue.enqueue(regular_workflow).get_result() == 5
+
+    # Test deferred event receivers
+    DBOS.destroy(destroy_registry=True)
+    dbos = DBOS(config=config)
+
+    @DBOS.scheduled("* * * * * *")
+    @DBOS.workflow()
+    def deferred_workflow(scheduled: datetime, actual: datetime) -> None:
+        nonlocal wf_counter
+        wf_counter += 1
+
+    DBOS.launch()
+    assert len(dbos.poller_stop_events) > 0
+    for event in dbos.poller_stop_events:
+        assert not event.is_set()
+    # Deactivate--scheduled workflow should stop
+    response = requests.get("http://localhost:3001/deactivate", timeout=5)
+    assert response.status_code == 200
+    for event in dbos.poller_stop_events:
+        assert event.is_set()
 
 
 def test_admin_recovery(config: ConfigFile) -> None:

tests/test_scheduler.py

@@ -202,7 +202,8 @@ def test_transaction() -> None:
                 time.sleep(1)
 
     # Stop the scheduled workflow
-    dbos.stop_events[0].set()
+    for evt in dbos.poller_stop_events:
+        evt.set()
 
     dbos._sys_db.update_workflow_status(
         {