Use the same column to manage queue executor_id (#235)

Author: maxdml

dbos/_schemas/system_database.py

@@ -154,7 +154,7 @@ class SystemSchema:
             nullable=False,
             primary_key=True,
         ),
-        Column("executor_id", Text),
+        # Column("executor_id", Text), # This column is deprecated. Do *not* use it.
         Column("queue_name", Text, nullable=False),
         Column(
             "created_at_epoch_ms",

dbos/_sys_db.py

@@ -1324,24 +1324,32 @@ def start_queued_workflows(self, queue: "Queue", executor_id: str) -> List[str]:
             # If there is a global or local concurrency limit N, select only the N oldest enqueued
             # functions, else select all of them.
 
-            # First lets figure out how many tasks the worker can dequeue
+            # First lets figure out how many tasks are eligible for dequeue.
+            # This means figuring out how many unstarted tasks are within the local and global concurrency limits
             running_tasks_query = (
                 sa.select(
-                    SystemSchema.workflow_queue.c.executor_id,
+                    SystemSchema.workflow_status.c.executor_id,
                     sa.func.count().label("task_count"),
                 )
+                .select_from(
+                    SystemSchema.workflow_queue.join(
+                        SystemSchema.workflow_status,
+                        SystemSchema.workflow_queue.c.workflow_uuid
+                        == SystemSchema.workflow_status.c.workflow_uuid,
+                    )
+                )
                 .where(SystemSchema.workflow_queue.c.queue_name == queue.name)
                 .where(
-                    SystemSchema.workflow_queue.c.executor_id.isnot(
+                    SystemSchema.workflow_queue.c.started_at_epoch_ms.isnot(
                         None
-                    )  # Task is dequeued
+                    )  # Task is started
                 )
                 .where(
                     SystemSchema.workflow_queue.c.completed_at_epoch_ms.is_(
                         None
-                    )  # Task is not completed
+                    )  # Task is not completed.
                 )
-                .group_by(SystemSchema.workflow_queue.c.executor_id)
+                .group_by(SystemSchema.workflow_status.c.executor_id)
             )
             running_tasks_result = c.execute(running_tasks_query).fetchall()
             running_tasks_result_dict = {row[0]: row[1] for row in running_tasks_result}
@@ -1351,12 +1359,6 @@ def start_queued_workflows(self, queue: "Queue", executor_id: str) -> List[str]:
 
             max_tasks = float("inf")
             if queue.worker_concurrency is not None:
-                # Worker local concurrency limit should always be >= running_tasks_for_this_worker
-                # This should never happen but a check + warning doesn't hurt
-                if running_tasks_for_this_worker > queue.worker_concurrency:
-                    dbos_logger.warning(
-                        f"Number of tasks on this worker ({running_tasks_for_this_worker}) exceeds the worker concurrency limit ({queue.worker_concurrency})"
-                    )
                 max_tasks = max(
                     0, queue.worker_concurrency - running_tasks_for_this_worker
                 )
@@ -1371,16 +1373,14 @@ def start_queued_workflows(self, queue: "Queue", executor_id: str) -> List[str]:
                 available_tasks = max(0, queue.concurrency - total_running_tasks)
                 max_tasks = min(max_tasks, available_tasks)
 
-            # Lookup tasks
+            # Lookup unstarted/uncompleted tasks (not running)
             query = (
                 sa.select(
                     SystemSchema.workflow_queue.c.workflow_uuid,
-                    SystemSchema.workflow_queue.c.started_at_epoch_ms,
-                    SystemSchema.workflow_queue.c.executor_id,
                 )
                 .where(SystemSchema.workflow_queue.c.queue_name == queue.name)
+                .where(SystemSchema.workflow_queue.c.started_at_epoch_ms == None)
                 .where(SystemSchema.workflow_queue.c.completed_at_epoch_ms == None)
-                .where(SystemSchema.workflow_queue.c.executor_id == None)
                 .order_by(SystemSchema.workflow_queue.c.created_at_epoch_ms.asc())
                 .with_for_update(nowait=True)  # Error out early
             )
@@ -1423,7 +1423,7 @@ def start_queued_workflows(self, queue: "Queue", executor_id: str) -> List[str]:
                 c.execute(
                     SystemSchema.workflow_queue.update()
                     .where(SystemSchema.workflow_queue.c.workflow_uuid == id)
-                    .values(started_at_epoch_ms=start_time_ms, executor_id=executor_id)
+                    .values(started_at_epoch_ms=start_time_ms)
                 )
                 ret_ids.append(id)
 
@@ -1468,23 +1468,26 @@ def clear_queue_assignment(self, workflow_id: str) -> bool:
 
         with self.engine.connect() as conn:
             with conn.begin() as transaction:
+                # Reset the start time in the queue to mark it as not started
                 res = conn.execute(
                     sa.update(SystemSchema.workflow_queue)
                     .where(SystemSchema.workflow_queue.c.workflow_uuid == workflow_id)
-                    .values(executor_id=None, started_at_epoch_ms=None)
+                    .where(
+                        SystemSchema.workflow_queue.c.completed_at_epoch_ms.is_(None)
+                    )
+                    .values(started_at_epoch_ms=None)
                 )
 
-                # If no rows were affected, the workflow is not anymore in the queue
+                # If no rows were affected, the workflow is not anymore in the queue or was already completed
                 if res.rowcount == 0:
                     transaction.rollback()
                     return False
 
+                # Reset the status of the task to "ENQUEUED"
                 res = conn.execute(
                     sa.update(SystemSchema.workflow_status)
                     .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
-                    .values(
-                        executor_id=None, status=WorkflowStatusString.ENQUEUED.value
-                    )
+                    .values(status=WorkflowStatusString.ENQUEUED.value)
                 )
                 if res.rowcount == 0:
                     # This should never happen