Job Queue (#106)

Add a queue abstraction for asynchronous processing.

Example usage:
```python
queue = Queue("example_queue")

@dbos.step()
def process_job(job):
  ...

@DBOS.workflow()
def process_jobs(jobs):
  for job in jobs:
    queue.enqueue(process_job, job)
```

You can enqueue any DBOS-annotated function (workflows, transactions, or
steps).

`enqueue` returns a [workflow
handle](https://docs.dbos.dev/python/reference/workflow_handles) so you
can track the status of enqueued tasks.

You can set a maximum concurrency level for a queue to cap the maximum
number of tasks from that queue that can be processed at once.

Scheduled workflows and Kafka workflows now use queues, making them more
robust in a distributed setting.

Author: kraftp

dbos/__init__.py

@@ -3,6 +3,7 @@
 from .dbos import DBOS, DBOSConfiguredInstance, WorkflowHandle, WorkflowStatus
 from .dbos_config import ConfigFile, get_dbos_database_url, load_config
 from .kafka_message import KafkaMessage
+from .queue import Queue
 from .system_database import GetWorkflowsInput, WorkflowStatusString
 
 __all__ = [
@@ -19,4 +20,5 @@
     "load_config",
     "get_dbos_database_url",
     "error",
+    "Queue",
 ]

dbos/core.py

@@ -63,6 +63,7 @@
     OperationResultInternal,
     WorkflowInputs,
     WorkflowStatusInternal,
+    WorkflowStatusString,
 )
 
 if TYPE_CHECKING:
@@ -126,6 +127,7 @@ def _init_workflow(
     class_name: Optional[str],
     config_name: Optional[str],
     temp_wf_type: Optional[str],
+    queue: Optional[str] = None,
 ) -> WorkflowStatusInternal:
     wfid = (
         ctx.workflow_id
@@ -134,7 +136,11 @@ def _init_workflow(
     )
     status: WorkflowStatusInternal = {
         "workflow_uuid": wfid,
-        "status": "PENDING",
+        "status": (
+            WorkflowStatusString.PENDING.value
+            if queue is None
+            else WorkflowStatusString.ENQUEUED.value
+        ),
         "name": wf_name,
         "class_name": class_name,
         "config_name": config_name,
@@ -150,21 +156,26 @@ def _init_workflow(
             json.dumps(ctx.authenticated_roles) if ctx.authenticated_roles else None
         ),
         "assumed_role": ctx.assumed_role,
+        "queue_name": queue,
     }
 
     # If we have a class name, the first arg is the instance and do not serialize
     if class_name is not None:
         inputs = {"args": inputs["args"][1:], "kwargs": inputs["kwargs"]}
 
-    if temp_wf_type != "transaction":
+    if temp_wf_type != "transaction" or queue is not None:
         # Synchronously record the status and inputs for workflows and single-step workflows
         # We also have to do this for single-step workflows because of the foreign key constraint on the operation outputs table
+        # TODO: Make this transactional (and with the queue step below)
         dbos._sys_db.update_workflow_status(status, False, ctx.in_recovery)
         dbos._sys_db.update_workflow_inputs(wfid, utils.serialize(inputs))
     else:
         # Buffer the inputs for single-transaction workflows, but don't buffer the status
         dbos._sys_db.buffer_workflow_inputs(wfid, utils.serialize(inputs))
 
+    if queue is not None:
+        dbos._sys_db.enqueue(wfid, queue)
+
     return status
 
 
@@ -179,6 +190,8 @@ def _execute_workflow(
         output = func(*args, **kwargs)
         status["status"] = "SUCCESS"
         status["output"] = utils.serialize(output)
+        if status["queue_name"] is not None:
+            dbos._sys_db.remove_from_queue(status["workflow_uuid"])
         dbos._sys_db.buffer_workflow_status(status)
     except DBOSWorkflowConflictIDError:
         # Retrieve the workflow handle and wait for the result.
@@ -191,6 +204,8 @@ def _execute_workflow(
     except Exception as error:
         status["status"] = "ERROR"
         status["error"] = utils.serialize(error)
+        if status["queue_name"] is not None:
+            dbos._sys_db.remove_from_queue(status["workflow_uuid"])
         dbos._sys_db.update_workflow_status(status)
         raise
 
@@ -249,6 +264,8 @@ def _execute_workflow_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[Any]
                 return _start_workflow(
                     dbos,
                     wf_func,
+                    status["queue_name"],
+                    True,
                     dbos._registry.instance_info_map[iname],
                     *inputs["args"],
                     **inputs["kwargs"],
@@ -264,14 +281,21 @@ def _execute_workflow_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[Any]
                 return _start_workflow(
                     dbos,
                     wf_func,
+                    status["queue_name"],
+                    True,
                     dbos._registry.class_info_map[class_name],
                     *inputs["args"],
                     **inputs["kwargs"],
                 )
         else:
             with SetWorkflowID(workflow_id):
                 return _start_workflow(
-                    dbos, wf_func, *inputs["args"], **inputs["kwargs"]
+                    dbos,
+                    wf_func,
+                    status["queue_name"],
+                    True,
+                    *inputs["args"],
+                    **inputs["kwargs"],
                 )
 
 
@@ -331,6 +355,8 @@ def _workflow_decorator(func: F) -> F:
 def _start_workflow(
     dbos: "DBOS",
     func: "Workflow[P, R]",
+    queue_name: Optional[str],
+    execute_workflow: bool,
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> "WorkflowHandle[R]":
@@ -384,8 +410,12 @@ def _start_workflow(
         class_name=get_dbos_class_name(fi, func, gin_args),
         config_name=get_config_name(fi, func, gin_args),
         temp_wf_type=get_temp_workflow_type(func),
+        queue=queue_name,
     )
 
+    if not execute_workflow:
+        return _WorkflowHandlePolling(new_wf_id, dbos)
+
     if fself is not None:
         future = dbos._executor.submit(
             cast(Callable[..., R], _execute_workflow_wthread),
@@ -529,6 +559,7 @@ def temp_wf(*args: Any, **kwargs: Any) -> Any:
         set_dbos_func_name(temp_wf, "<temp>." + func.__qualname__)
         set_temp_workflow_type(temp_wf, "transaction")
         dbosreg.register_wf_function(get_dbos_func_name(temp_wf), wrapped_wf)
+        wrapper.__orig_func = temp_wf  # type: ignore
 
         return cast(F, wrapper)
 
@@ -645,6 +676,7 @@ def temp_wf(*args: Any, **kwargs: Any) -> Any:
         set_dbos_func_name(temp_wf, "<temp>." + func.__qualname__)
         set_temp_workflow_type(temp_wf, "step")
         dbosreg.register_wf_function(get_dbos_func_name(temp_wf), wrapped_wf)
+        wrapper.__orig_func = temp_wf  # type: ignore
 
         return cast(F, wrapper)
 

dbos/dbos.py

@@ -20,7 +20,6 @@
     Tuple,
     Type,
     TypeVar,
-    cast,
 )
 
 from opentelemetry.trace import Span
@@ -40,6 +39,7 @@
     _WorkflowHandlePolling,
 )
 from dbos.decorators import classproperty
+from dbos.queue import Queue, queue_thread
 from dbos.recovery import _recover_pending_workflows, _startup_recovery_thread
 from dbos.registrations import (
     DBOSClassInfo,
@@ -138,6 +138,7 @@ def __init__(self) -> None:
         self.workflow_info_map: dict[str, Workflow[..., Any]] = {}
         self.class_info_map: dict[str, type] = {}
         self.instance_info_map: dict[str, object] = {}
+        self.queue_info_map: dict[str, Queue] = {}
         self.pollers: list[_RegisteredJob] = []
         self.dbos: Optional[DBOS] = None
         self.config: Optional[ConfigFile] = None
@@ -356,6 +357,11 @@ def _launch(self) -> None:
         # Start flush workflow buffers thread
         self._executor.submit(self._sys_db.flush_workflow_buffers)
 
+        # Start the queue thread
+        evt = threading.Event()
+        self.stop_events.append(evt)
+        self._executor.submit(queue_thread, evt, self)
+
         # Grab any pollers that were deferred and start them
         for evt, func, args, kwargs in self._registry.pollers:
             self.stop_events.append(evt)
@@ -488,13 +494,18 @@ def scheduled(cls, cron: str) -> Callable[[ScheduledWorkflow], ScheduledWorkflow
 
     @classmethod
     def kafka_consumer(
-        cls, config: dict[str, Any], topics: list[str]
+        cls,
+        config: dict[str, Any],
+        topics: list[str],
+        in_order: bool = False,
     ) -> Callable[[KafkaConsumerWorkflow], KafkaConsumerWorkflow]:
         """Decorate a function to be used as a Kafka consumer."""
         try:
             from dbos.kafka import kafka_consumer
 
-            return kafka_consumer(_get_or_create_dbos_registry(), config, topics)
+            return kafka_consumer(
+                _get_or_create_dbos_registry(), config, topics, in_order
+            )
         except ModuleNotFoundError as e:
             raise DBOSException(
                 f"{e.name} dependency not found. Please install {e.name} via your package manager."
@@ -508,7 +519,7 @@ def start_workflow(
         **kwargs: P.kwargs,
     ) -> WorkflowHandle[R]:
         """Invoke a workflow function in the background, returning a handle to the ongoing execution."""
-        return _start_workflow(_get_dbos_instance(), func, *args, **kwargs)
+        return _start_workflow(_get_dbos_instance(), func, None, True, *args, **kwargs)
 
     @classmethod
     def get_workflow_status(cls, workflow_id: str) -> Optional[WorkflowStatus]:

dbos/kafka.py

@@ -1,26 +1,30 @@
 import threading
-import traceback
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Generator, NoReturn, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, NoReturn
 
 from confluent_kafka import Consumer, KafkaError, KafkaException
-from confluent_kafka import Message as CTypeMessage
+
+from dbos.queue import Queue
 
 if TYPE_CHECKING:
     from dbos.dbos import _DBOSRegistry
 
 from .context import SetWorkflowID
+from .error import DBOSInitializationError
 from .kafka_message import KafkaMessage
 from .logger import dbos_logger
 
 KafkaConsumerWorkflow = Callable[[KafkaMessage], None]
 
+kafka_queue: Queue
+in_order_kafka_queues: dict[str, Queue] = {}
+
 
 def _kafka_consumer_loop(
     func: KafkaConsumerWorkflow,
     config: dict[str, Any],
     topics: list[str],
     stop_event: threading.Event,
+    in_order: bool,
 ) -> None:
 
     def on_error(err: KafkaError) -> NoReturn:
@@ -70,24 +74,35 @@ def on_error(err: KafkaError) -> NoReturn:
                 with SetWorkflowID(
                     f"kafka-unique-id-{msg.topic}-{msg.partition}-{msg.offset}"
                 ):
-                    try:
-                        func(msg)
-                    except Exception as e:
-                        dbos_logger.error(
-                            f"Exception encountered in Kafka consumer: {traceback.format_exc()}"
-                        )
+                    if in_order:
+                        assert msg.topic is not None
+                        queue = in_order_kafka_queues[msg.topic]
+                        queue.enqueue(func, msg)
+                    else:
+                        kafka_queue.enqueue(func, msg)
 
     finally:
         consumer.close()
 
 
 def kafka_consumer(
-    dbosreg: "_DBOSRegistry", config: dict[str, Any], topics: list[str]
+    dbosreg: "_DBOSRegistry", config: dict[str, Any], topics: list[str], in_order: bool
 ) -> Callable[[KafkaConsumerWorkflow], KafkaConsumerWorkflow]:
     def decorator(func: KafkaConsumerWorkflow) -> KafkaConsumerWorkflow:
+        if in_order:
+            for topic in topics:
+                if topic.startswith("^"):
+                    raise DBOSInitializationError(
+                        f"Error: in-order processing is not supported for regular expression topic selectors ({topic})"
+                    )
+                queue = Queue(f"_dbos_kafka_queue_topic_{topic}", concurrency=1)
+                in_order_kafka_queues[topic] = queue
+        else:
+            global kafka_queue
+            kafka_queue = Queue("_dbos_internal_queue")
         stop_event = threading.Event()
         dbosreg.register_poller(
-            stop_event, _kafka_consumer_loop, func, config, topics, stop_event
+            stop_event, _kafka_consumer_loop, func, config, topics, stop_event, in_order
         )
         return func
 

dbos/migrations/versions/eab0cc1d9a14_job_queue.py

@@ -0,0 +1,55 @@
+"""job_queue
+
+Revision ID: eab0cc1d9a14
+Revises: a3b18ad34abe
+Create Date: 2024-09-13 14:50:00.531294
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "eab0cc1d9a14"
+down_revision: Union[str, None] = "a3b18ad34abe"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "job_queue",
+        sa.Column("workflow_uuid", sa.Text(), nullable=False),
+        sa.Column("queue_name", sa.Text(), nullable=False),
+        sa.Column(
+            "created_at_epoch_ms",
+            sa.BigInteger(),
+            server_default=sa.text(
+                "(EXTRACT(epoch FROM now()) * 1000::numeric)::bigint"
+            ),
+            nullable=False,
+            primary_key=True,
+        ),
+        sa.ForeignKeyConstraint(
+            ["workflow_uuid"],
+            ["dbos.workflow_status.workflow_uuid"],
+            onupdate="CASCADE",
+            ondelete="CASCADE",
+        ),
+        schema="dbos",
+    )
+    op.add_column(
+        "workflow_status",
+        sa.Column(
+            "queue_name",
+            sa.Text(),
+        ),
+        schema="dbos",
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("job_queue", schema="dbos")
+    op.drop_column("workflow_status", "queue_name", schema="dbos")

dbos/queue.py

@@ -0,0 +1,36 @@
+import threading
+import time
+from typing import TYPE_CHECKING, Optional
+
+from dbos.core import P, R, _execute_workflow_id, _start_workflow
+from dbos.error import DBOSInitializationError
+
+if TYPE_CHECKING:
+    from dbos.dbos import DBOS, Workflow, WorkflowHandle
+
+
+class Queue:
+    def __init__(self, name: str, concurrency: Optional[int] = None) -> None:
+        self.name = name
+        self.concurrency = concurrency
+        from dbos.dbos import _get_or_create_dbos_registry
+
+        registry = _get_or_create_dbos_registry()
+        registry.queue_info_map[self.name] = self
+
+    def enqueue(
+        self, func: "Workflow[P, R]", *args: P.args, **kwargs: P.kwargs
+    ) -> "WorkflowHandle[R]":
+        from dbos.dbos import _get_dbos_instance
+
+        dbos = _get_dbos_instance()
+        return _start_workflow(dbos, func, self.name, False, *args, **kwargs)
+
+
+def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
+    while not stop_event.is_set():
+        time.sleep(1)
+        for queue_name, queue in dbos._registry.queue_info_map.items():
+            wf_ids = dbos._sys_db.start_queued_workflows(queue_name, queue.concurrency)
+            for id in wf_ids:
+                _execute_workflow_id(dbos, id)