Worker Concurrency Fix (#336)

Fix an issue where workflows of the wrong version count against worker
concurrency limits.
(#333)

Author: kraftp

dbos/_sys_db.py

@@ -1739,15 +1739,30 @@ def start_queued_workflows(
                 available_tasks = max(0, queue.concurrency - total_running_tasks)
                 max_tasks = min(max_tasks, available_tasks)
 
-            # Lookup unstarted/uncompleted tasks (not running)
+            # Retrieve the first max_tasks workflows in the queue.
+            # Only retrieve workflows of the appropriate version (or without version set)
             query = (
                 sa.select(
                     SystemSchema.workflow_queue.c.workflow_uuid,
                 )
+                .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.started_at_epoch_ms == None)
                 .where(SystemSchema.workflow_queue.c.completed_at_epoch_ms == None)
                 .order_by(SystemSchema.workflow_queue.c.created_at_epoch_ms.asc())
+                .where(
+                    sa.or_(
+                        SystemSchema.workflow_status.c.application_version
+                        == app_version,
+                        SystemSchema.workflow_status.c.application_version.is_(None),
+                    )
+                )
                 .with_for_update(nowait=True)  # Error out early
             )
             # Apply limit only if max_tasks is finite
@@ -1779,15 +1794,6 @@ def start_queued_workflows(
                         SystemSchema.workflow_status.c.status
                         == WorkflowStatusString.ENQUEUED.value
                     )
-                    .where(
-                        sa.or_(
-                            SystemSchema.workflow_status.c.application_version
-                            == app_version,
-                            SystemSchema.workflow_status.c.application_version.is_(
-                                None
-                            ),
-                        )
-                    )
                     .values(
                         status=WorkflowStatusString.PENDING.value,
                         application_version=app_version,

tests/test_queue.py

@@ -14,6 +14,7 @@
 from dbos import (
     DBOS,
     ConfigFile,
+    DBOSClient,
     DBOSConfiguredInstance,
     Queue,
     SetEnqueueOptions,
@@ -1210,3 +1211,29 @@ async def test_workflow(var1: str) -> str:
         with SetWorkflowID(wfid2):
             handle2 = await queue.enqueue_async(test_workflow, "def")
     assert (await handle2.get_result()) == "def-c-p"
+
+def test_worker_concurrency_across_versions(dbos: DBOS, client: DBOSClient) -> None:
+    queue = Queue("test_worker_concurrency_across_versions", worker_concurrency=1)
+
+    @DBOS.workflow()
+    def test_workflow() -> str:
+        return DBOS.workflow_id
+
+    # First enqueue a workflow on the other version, then on the current version
+    other_version = "other_version"
+    other_version_handle: WorkflowHandle[None] = client.enqueue(
+        {
+            "queue_name": "test_worker_concurrency_across_versions",
+            "workflow_name": test_workflow.__qualname__,
+            "app_version": other_version,
+        }
+    )
+    handle = queue.enqueue(test_workflow)
+
+    # Verify the workflow on the current version completes, but the other version is still ENQUEUED
+    assert handle.get_result()
+    assert other_version_handle.get_status().status == "ENQUEUED"
+
+    # Change the version, verify the other version complets
+    GlobalParams.app_version = other_version
+    assert other_version_handle.get_result()