beanqueue 1.1.9__tar.gz → 2.0.0rc0__tar.gz

beanqueue-1.1.9/README.md → beanqueue-2.0.0rc0/PKG-INFO

@@ -1,3 +1,22 @@
+Metadata-Version: 2.4
+Name: beanqueue
+Version: 2.0.0rc0
+Summary: BeanQueue or BQ for short, PostgreSQL SKIP LOCK and SQLAlchemy based worker queue library
+Author-email: Fang-Pen Lin <fangpen@launchplatform.com>
+License-Expression: MIT
+License-File: LICENSE
+Requires-Python: <4,>=3.11
+Requires-Dist: blinker<2,>=1.8.2
+Requires-Dist: click<9,>=8.1.7
+Requires-Dist: pydantic-settings<3,>=2.2.1
+Requires-Dist: rich<14,>=13.7.1
+Requires-Dist: sqlalchemy<3,>=2.0.30
+Requires-Dist: venusian<4,>=3.1.0
+Provides-Extra: metrics
+Requires-Dist: starlette<2,>=0.27; extra == 'metrics'
+Requires-Dist: uvicorn<1,>=0.30.0; extra == 'metrics'
+Description-Content-Type: text/markdown
+
 # BeanQueue [![CircleCI](https://dl.circleci.com/status-badge/img/gh/LaunchPlatform/bq/tree/master.svg?style=svg)](https://dl.circleci.com/status-badge/redirect/gh/LaunchPlatform/bq/tree/master)
 
 BeanQueue, a lightweight Python task queue framework based on [SQLAlchemy](https://www.sqlalchemy.org/), PostgreSQL [SKIP LOCKED queries](https://www.2ndquadrant.com/en/blog/what-is-select-skip-locked-for-in-postgresql-9-5/) and [NOTIFY](https://www.postgresql.org/docs/current/sql-notify.html) / [LISTEN](https://www.postgresql.org/docs/current/sql-listen.html) statements.
@@ -14,6 +33,7 @@ BeanQueue, a lightweight Python task queue framework based on [SQLAlchemy](https
 - **Retry**: Built-in and customizable retry policies
 - **Schedule**: Schedule tasks to run later
 - **Worker heartbeat and auto-reschedule**: Each worker keeps updating heartbeat, if one is found dead, the others will reschedule the tasks
+- **Custom health checks**: Optional HTTP `/healthz` endpoint with pluggable checks via Blinker signals
 - **Customizable**: Custom Task, Worker and Event models. Use it as a library and build your own work queue
 
 ## Install
@@ -22,6 +42,20 @@ BeanQueue, a lightweight Python task queue framework based on [SQLAlchemy](https
 pip install beanqueue
 ```
 
+To enable the optional metrics HTTP server (currently `/healthz` only), install with the `metrics` extra:
+
+```bash
+pip install "beanqueue[metrics]"
+```
+
+## Upgrading to 2.0
+
+BeanQueue 2.0 includes breaking changes around the metrics HTTP server and custom health checks:
+
+- **`METRICS_HTTP_SERVER_ENABLED` defaults to `False`** (it was `True` in 1.x). Set `BQ_METRICS_HTTP_SERVER_ENABLED=true` to turn the server back on.
+- **The metrics server requires optional dependencies.** Install `beanqueue[metrics]` (`starlette` and `uvicorn`). Without them, enabling the server raises an error at startup.
+- **Custom health checks use the `healthz_check` event** (`bq.events.healthz_check`) instead of a `healthz_check` argument on `bq.BeanQueue`. Connect sync or async receivers to the signal.
+
 ## Usage
 
 You can define a basic task processor like this
@@ -216,6 +250,91 @@ Or if you prefer to define your own process command, you can also call `process_
 app.process_tasks(channels=("images",))
 ```
 
+### Health check and metrics HTTP server
+
+When enabled, each worker starts a small HTTP server (Starlette + Uvicorn) for operational endpoints.
+Today this only exposes `GET /healthz`, which returns `{"status": "ok"}` by default.
+
+Enable it with the `metrics` extra installed and configuration:
+
+```bash
+pip install "beanqueue[metrics]"
+BQ_METRICS_HTTP_SERVER_ENABLED=true bq process images
+```
+
+Relevant settings (see [bq/config.py](bq/config.py)):
+
+| Setting | Env var | Default |
+| --- | --- | --- |
+| `METRICS_HTTP_SERVER_ENABLED` | `BQ_METRICS_HTTP_SERVER_ENABLED` | `False` |
+| `METRICS_HTTP_SERVER_INTERFACE` | `BQ_METRICS_HTTP_SERVER_INTERFACE` | `""` (all interfaces) |
+| `METRICS_HTTP_SERVER_PORT` | `BQ_METRICS_HTTP_SERVER_PORT` | `8000` |
+| `METRICS_HTTP_SERVER_LOG_LEVEL` | `BQ_METRICS_HTTP_SERVER_LOG_LEVEL` | `30` (`WARNING`) |
+
+Access requests are logged at INFO via `uvicorn.access` (visible even when `METRICS_HTTP_SERVER_LOG_LEVEL` is `WARNING`).
+BeanQueue also uses a `metrics_server` logger for its own messages.
+Override the entire logging setup by passing a [logging.config](https://docs.python.org/3/library/logging.config.html#logging-config-dictschema) dict via `METRICS_HTTP_SERVER_LOG_CONFIG` (or `BQ_METRICS_HTTP_SERVER_LOG_CONFIG` as JSON):
+
+```python
+import bq
+
+config = bq.Config(
+    METRICS_HTTP_SERVER_ENABLED=True,
+    METRICS_HTTP_SERVER_LOG_CONFIG={
+        "version": 1,
+        "disable_existing_loggers": False,
+        "handlers": {
+            "default": {
+                "class": "logging.StreamHandler",
+                "formatter": "default",
+            }
+        },
+        "formatters": {
+            "default": {
+                "format": "%(asctime)s %(name)s %(levelname)s %(message)s",
+            }
+        },
+        "loggers": {
+            "metrics_server": {"handlers": ["default"], "level": "INFO"},
+            "uvicorn.access": {"handlers": ["default"], "level": "INFO"},
+        },
+    },
+)
+app = bq.BeanQueue(config=config)
+```
+
+#### Custom health checks
+
+Register additional checks by connecting receivers to `bq.events.healthz_check`.
+If no receivers are connected, `/healthz` returns OK without touching the database.
+
+With receivers connected, BeanQueue loads the current worker and passes a database `session` to each check.
+Receivers may be synchronous or asynchronous; both can be mixed on the same signal.
+
+```python
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+
+import bq
+from bq import events
+
+app = bq.BeanQueue()
+
+
+@events.healthz_check.connect
+def check_database(sender: bq.BeanQueue, worker, session: Session):
+    session.execute(text("SELECT 1"))
+
+
+@events.healthz_check.connect
+async def check_external_service(sender: bq.BeanQueue, worker, session: Session):
+    # async HTTP call, etc.
+    ...
+```
+
+Receiver signature must accept the keyword arguments you need, for example `(sender, worker, session)`, or use `(sender, **kwargs)`.
+If a check raises an exception, `/healthz` responds with HTTP 500 and a JSON body containing the error message.
+
 ### Define your own tables
 
 BeanQueue is designed to be as customizable as much as possible.

beanqueue-1.1.9/PKG-INFO → beanqueue-2.0.0rc0/README.md

@@ -1,19 +1,3 @@
-Metadata-Version: 2.4
-Name: beanqueue
-Version: 1.1.9
-Summary: BeanQueue or BQ for short, PostgreSQL SKIP LOCK and SQLAlchemy based worker queue library
-Author-email: Fang-Pen Lin <fangpen@launchplatform.com>
-License-Expression: MIT
-License-File: LICENSE
-Requires-Python: ~=3.11
-Requires-Dist: blinker<2,>=1.8.2
-Requires-Dist: click<9,>=8.1.7
-Requires-Dist: pydantic-settings<3,>=2.2.1
-Requires-Dist: rich<14,>=13.7.1
-Requires-Dist: sqlalchemy<3,>=2.0.30
-Requires-Dist: venusian<4,>=3.1.0
-Description-Content-Type: text/markdown
-
 # BeanQueue [![CircleCI](https://dl.circleci.com/status-badge/img/gh/LaunchPlatform/bq/tree/master.svg?style=svg)](https://dl.circleci.com/status-badge/redirect/gh/LaunchPlatform/bq/tree/master)
 
 BeanQueue, a lightweight Python task queue framework based on [SQLAlchemy](https://www.sqlalchemy.org/), PostgreSQL [SKIP LOCKED queries](https://www.2ndquadrant.com/en/blog/what-is-select-skip-locked-for-in-postgresql-9-5/) and [NOTIFY](https://www.postgresql.org/docs/current/sql-notify.html) / [LISTEN](https://www.postgresql.org/docs/current/sql-listen.html) statements.
@@ -30,6 +14,7 @@ BeanQueue, a lightweight Python task queue framework based on [SQLAlchemy](https
 - **Retry**: Built-in and customizable retry policies
 - **Schedule**: Schedule tasks to run later
 - **Worker heartbeat and auto-reschedule**: Each worker keeps updating heartbeat, if one is found dead, the others will reschedule the tasks
+- **Custom health checks**: Optional HTTP `/healthz` endpoint with pluggable checks via Blinker signals
 - **Customizable**: Custom Task, Worker and Event models. Use it as a library and build your own work queue
 
 ## Install
@@ -38,6 +23,20 @@ BeanQueue, a lightweight Python task queue framework based on [SQLAlchemy](https
 pip install beanqueue
 ```
 
+To enable the optional metrics HTTP server (currently `/healthz` only), install with the `metrics` extra:
+
+```bash
+pip install "beanqueue[metrics]"
+```
+
+## Upgrading to 2.0
+
+BeanQueue 2.0 includes breaking changes around the metrics HTTP server and custom health checks:
+
+- **`METRICS_HTTP_SERVER_ENABLED` defaults to `False`** (it was `True` in 1.x). Set `BQ_METRICS_HTTP_SERVER_ENABLED=true` to turn the server back on.
+- **The metrics server requires optional dependencies.** Install `beanqueue[metrics]` (`starlette` and `uvicorn`). Without them, enabling the server raises an error at startup.
+- **Custom health checks use the `healthz_check` event** (`bq.events.healthz_check`) instead of a `healthz_check` argument on `bq.BeanQueue`. Connect sync or async receivers to the signal.
+
 ## Usage
 
 You can define a basic task processor like this
@@ -232,6 +231,91 @@ Or if you prefer to define your own process command, you can also call `process_
 app.process_tasks(channels=("images",))
 ```
 
+### Health check and metrics HTTP server
+
+When enabled, each worker starts a small HTTP server (Starlette + Uvicorn) for operational endpoints.
+Today this only exposes `GET /healthz`, which returns `{"status": "ok"}` by default.
+
+Enable it with the `metrics` extra installed and configuration:
+
+```bash
+pip install "beanqueue[metrics]"
+BQ_METRICS_HTTP_SERVER_ENABLED=true bq process images
+```
+
+Relevant settings (see [bq/config.py](bq/config.py)):
+
+| Setting | Env var | Default |
+| --- | --- | --- |
+| `METRICS_HTTP_SERVER_ENABLED` | `BQ_METRICS_HTTP_SERVER_ENABLED` | `False` |
+| `METRICS_HTTP_SERVER_INTERFACE` | `BQ_METRICS_HTTP_SERVER_INTERFACE` | `""` (all interfaces) |
+| `METRICS_HTTP_SERVER_PORT` | `BQ_METRICS_HTTP_SERVER_PORT` | `8000` |
+| `METRICS_HTTP_SERVER_LOG_LEVEL` | `BQ_METRICS_HTTP_SERVER_LOG_LEVEL` | `30` (`WARNING`) |
+
+Access requests are logged at INFO via `uvicorn.access` (visible even when `METRICS_HTTP_SERVER_LOG_LEVEL` is `WARNING`).
+BeanQueue also uses a `metrics_server` logger for its own messages.
+Override the entire logging setup by passing a [logging.config](https://docs.python.org/3/library/logging.config.html#logging-config-dictschema) dict via `METRICS_HTTP_SERVER_LOG_CONFIG` (or `BQ_METRICS_HTTP_SERVER_LOG_CONFIG` as JSON):
+
+```python
+import bq
+
+config = bq.Config(
+    METRICS_HTTP_SERVER_ENABLED=True,
+    METRICS_HTTP_SERVER_LOG_CONFIG={
+        "version": 1,
+        "disable_existing_loggers": False,
+        "handlers": {
+            "default": {
+                "class": "logging.StreamHandler",
+                "formatter": "default",
+            }
+        },
+        "formatters": {
+            "default": {
+                "format": "%(asctime)s %(name)s %(levelname)s %(message)s",
+            }
+        },
+        "loggers": {
+            "metrics_server": {"handlers": ["default"], "level": "INFO"},
+            "uvicorn.access": {"handlers": ["default"], "level": "INFO"},
+        },
+    },
+)
+app = bq.BeanQueue(config=config)
+```
+
+#### Custom health checks
+
+Register additional checks by connecting receivers to `bq.events.healthz_check`.
+If no receivers are connected, `/healthz` returns OK without touching the database.
+
+With receivers connected, BeanQueue loads the current worker and passes a database `session` to each check.
+Receivers may be synchronous or asynchronous; both can be mixed on the same signal.
+
+```python
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+
+import bq
+from bq import events
+
+app = bq.BeanQueue()
+
+
+@events.healthz_check.connect
+def check_database(sender: bq.BeanQueue, worker, session: Session):
+    session.execute(text("SELECT 1"))
+
+
+@events.healthz_check.connect
+async def check_external_service(sender: bq.BeanQueue, worker, session: Session):
+    # async HTTP call, etc.
+    ...
+```
+
+Receiver signature must accept the keyword arguments you need, for example `(sender, worker, session)`, or use `(sender, **kwargs)`.
+If a check raises an exception, `/healthz` responds with HTTP 500 and a JSON body containing the error message.
+
 ### Define your own tables
 
 BeanQueue is designed to be as customizable as much as possible.

{beanqueue-1.1.9 → beanqueue-2.0.0rc0}/bq/app.py

@@ -1,21 +1,22 @@
 import functools
 import importlib
-import json
 import logging
 import platform
 import sys
 import threading
 import typing
+from concurrent.futures import FIRST_COMPLETED
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import wait as futures_wait
 from importlib.metadata import PackageNotFoundError
 from importlib.metadata import version
-from wsgiref.simple_server import make_server
-from wsgiref.simple_server import WSGIRequestHandler
 
 import venusian
 from sqlalchemy import func
 from sqlalchemy.engine import create_engine
 from sqlalchemy.engine import Engine
 from sqlalchemy.orm import Session as DBSession
+from sqlalchemy.pool import QueuePool
 from sqlalchemy.pool import SingletonThreadPool
 
 from . import constants
@@ -23,6 +24,7 @@ from . import events
 from . import models
 from .config import Config
 from .db.session import SessionMaker
+from .metrics import MetricsServer
 from .processors.processor import Processor
 from .processors.processor import ProcessorHelper
 from .processors.registry import collect
@@ -33,21 +35,6 @@ from .utils import load_module_var
 logger = logging.getLogger(__name__)
 
 
-class WSGIRequestHandlerWithLogger(WSGIRequestHandler):
-    logger = logging.getLogger("metrics_server")
-
-    def log_message(self, format, *args):
-        message = format % args
-        self.logger.info(
-            "%s - - [%s] %s\n"
-            % (
-                self.address_string(),
-                self.log_date_time_string(),
-                message.translate(self._control_char_table),
-            )
-        )
-
-
 class BeanQueue:
     def __init__(
         self,
@@ -63,13 +50,32 @@ class BeanQueue:
         self.dispatch_service_cls = dispatch_service_cls
         self._engine = engine
         self._worker_update_shutdown_event: threading.Event = threading.Event()
-        # noop if metrics thread is not started yet, shutdown if it is started
-        self._metrics_server_shutdown: typing.Callable[[], None] = lambda: None
+        self._metrics_server: MetricsServer | None = None
 
     def create_default_engine(self):
-        return create_engine(
-            str(self.config.DATABASE_URL), poolclass=SingletonThreadPool
-        )
+        # Use thread-safe connection pool when thread pool executor is enabled
+        if self.config.MAX_WORKER_THREADS != 1:
+            # QueuePool is thread-safe and suitable for multi-threaded usage
+            # Configure pool size based on number of worker threads
+            max_workers = (
+                self.config.MAX_WORKER_THREADS
+                if self.config.MAX_WORKER_THREADS > 0
+                else 10
+            )
+            pool_size = (
+                max_workers + 5
+            )  # Extra connections for main thread and worker update thread
+            return create_engine(
+                str(self.config.DATABASE_URL),
+                poolclass=QueuePool,
+                pool_size=pool_size,
+                max_overflow=10,
+            )
+        else:
+            # SingletonThreadPool for single-threaded sequential processing
+            return create_engine(
+                str(self.config.DATABASE_URL), poolclass=SingletonThreadPool
+            )
 
     def make_session(self) -> DBSession:
         return self.session_cls(bind=self.engine)
@@ -196,65 +202,164 @@ class BeanQueue:
             db.add(current_worker)
             db.commit()
 
-    def _serve_http_request(
-        self, worker_id: typing.Any, environ: dict, start_response: typing.Callable
-    ) -> list[bytes]:
-        path = environ["PATH_INFO"]
-        if path == "/healthz":
-            db = self.make_session()
-            worker_service = self._make_worker_service(db)
-            worker = worker_service.get_worker(worker_id)
-            if worker is not None and worker.state == models.WorkerState.RUNNING:
-                start_response(
-                    "200 OK",
-                    [
-                        ("Content-Type", "application/json"),
-                    ],
-                )
-                return [
-                    json.dumps(dict(status="ok", worker_id=str(worker_id))).encode(
-                        "utf8"
+    def _process_task_in_thread(
+        self,
+        task_id: typing.Any,
+        registry: typing.Any,
+    ):
+        """Process a single task in a thread-safe manner with its own database session.
+
+        This method is called from worker threads in the thread pool. It creates its own
+        database session to avoid SQLAlchemy session conflicts between threads.
+        """
+        db = self.make_session()
+        try:
+            # Reload the task in this thread's session to avoid SQLAlchemy context issues
+            task = db.query(self.task_model).filter(self.task_model.id == task_id).one()
+
+            logger.info(
+                "Processing task %s, channel=%s, module=%s, func=%s",
+                task.id,
+                task.channel,
+                task.module,
+                task.func_name,
+            )
+            registry.process(task, event_cls=self.event_model)
+            db.commit()
+        except Exception as e:
+            logger.exception("Error processing task %s: %s", task_id, e)
+            db.rollback()
+            raise
+        finally:
+            db.close()
+
+    def _process_tasks_sequential(
+        self,
+        db: DBSession,
+        dispatch_service: DispatchService,
+        registry: typing.Any,
+        channels: tuple[str, ...],
+        worker_id: typing.Any,
+    ):
+        """Process tasks sequentially (original behavior for MAX_WORKER_THREADS=1)."""
+        while True:
+            while True:
+                tasks = dispatch_service.dispatch(
+                    channels,
+                    worker_id=worker_id,
+                    limit=self.config.BATCH_SIZE,
+                ).all()
+
+                for task in tasks:
+                    logger.info(
+                        "Processing task %s, channel=%s, module=%s, func=%s",
+                        task.id,
+                        task.channel,
+                        task.module,
+                        task.func_name,
                     )
-                ]
-            else:
-                logger.warning("Bad worker %s state %s", worker_id, worker.state)
-                start_response(
-                    "500 Internal Server Error",
-                    [
-                        ("Content-Type", "application/json"),
-                    ],
+                    registry.process(task, event_cls=self.event_model)
+                if tasks:
+                    db.commit()
+
+                if not tasks:
+                    break
+
+            db.close()
+            try:
+                for notification in dispatch_service.poll(
+                    timeout=self.config.POLL_TIMEOUT
+                ):
+                    logger.debug("Receive notification %s", notification)
+            except TimeoutError:
+                logger.debug("Poll timeout, try again")
+                continue
+
+    def _process_tasks_threaded(
+        self,
+        db: DBSession,
+        executor: ThreadPoolExecutor,
+        dispatch_service: DispatchService,
+        registry: typing.Any,
+        channels: tuple[str, ...],
+        worker_id: typing.Any,
+    ):
+        """Process tasks using thread pool with continuous task feeding.
+
+        This implementation continuously checks for completed futures and fetches new tasks
+        when there's capacity in the thread pool. It uses concurrent.futures.wait() to
+        properly detect ANY completed future, not just the first one submitted.
+        """
+        max_workers = self.config.MAX_WORKER_THREADS
+        if max_workers == 0:
+            max_workers = 10  # Default when set to auto
+
+        running_futures: set = set()
+
+        while True:
+            # Clean up ANY completed futures using wait() with zero timeout
+            if running_futures:
+                done, running_futures = futures_wait(
+                    running_futures, timeout=0, return_when=FIRST_COMPLETED
                 )
-                return [
-                    json.dumps(
-                        dict(
-                            status="internal error",
-                            worker_id=str(worker_id),
-                            state=str(worker.state),
+                for f in done:
+                    try:
+                        f.result()
+                    except Exception as e:
+                        logger.error("Task processing failed: %s", e)
+
+            # If we have capacity, fetch and submit more tasks
+            capacity = max_workers - len(running_futures)
+            if capacity > 0:
+                tasks = dispatch_service.dispatch(
+                    channels,
+                    worker_id=worker_id,
+                    limit=min(capacity, self.config.BATCH_SIZE),
+                ).all()
+
+                # Always commit to close the transaction and refresh the snapshot,
+                # so subsequent dispatch calls can see newly committed tasks
+                db.commit()
+
+                if tasks:
+                    logger.debug(
+                        "Dispatching %d tasks (running=%d, capacity=%d)",
+                        len(tasks),
+                        len(running_futures),
+                        capacity,
+                    )
+
+                    for task in tasks:
+                        future = executor.submit(
+                            self._process_task_in_thread,
+                            task.id,
+                            registry,
                         )
-                    ).encode("utf8")
-                ]
-        # TODO: add other metrics endpoints
-        start_response(
-            "404 NOT FOUND",
-            [
-                ("Content-Type", "application/json"),
-            ],
-        )
-        return [json.dumps(dict(status="not found")).encode("utf8")]
-
-    def run_metrics_http_server(self, worker_id: typing.Any):
-        host = self.config.METRICS_HTTP_SERVER_INTERFACE
-        port = self.config.METRICS_HTTP_SERVER_PORT
-        with make_server(
-            host,
-            port,
-            functools.partial(self._serve_http_request, worker_id),
-            handler_class=WSGIRequestHandlerWithLogger,
-        ) as httpd:
-            # expose graceful shutdown to the main thread
-            self._metrics_server_shutdown = httpd.shutdown
-            logger.info("Run metrics HTTP server on %s:%s", host, port)
-            httpd.serve_forever()
+                        running_futures.add(future)
+
+            # If we have running tasks, wait briefly for any to complete then check for new tasks
+            if running_futures:
+                # Short wait - allows checking for new tasks frequently
+                done, running_futures = futures_wait(
+                    running_futures, timeout=0.05, return_when=FIRST_COMPLETED
+                )
+                for f in done:
+                    try:
+                        f.result()
+                    except Exception as e:
+                        logger.error("Task processing failed: %s", e)
+                continue
+
+            # No running tasks and no new tasks found - poll for notifications
+            db.close()
+            try:
+                for notification in dispatch_service.poll(
+                    timeout=self.config.POLL_TIMEOUT
+                ):
+                    logger.debug("Receive notification %s", notification)
+            except TimeoutError:
+                logger.debug("Poll timeout, try again")
+                continue
 
     def process_tasks(
         self,
@@ -300,17 +405,9 @@ class BeanQueue:
         dispatch_service.listen(channels)
         db.commit()
 
-        metrics_server_thread = None
         if self.config.METRICS_HTTP_SERVER_ENABLED:
-            WSGIRequestHandlerWithLogger.logger.setLevel(
-                self.config.METRICS_HTTP_SERVER_LOG_LEVEL
-            )
-            metrics_server_thread = threading.Thread(
-                target=self.run_metrics_http_server,
-                args=(worker.id,),
-            )
-            metrics_server_thread.daemon = True
-            metrics_server_thread.start()
+            self._metrics_server = MetricsServer(self, worker.id)
+            self._metrics_server.start()
 
         logger.info("Created worker %s, name=%s", worker.id, worker.name)
         events.worker_init.send(self, worker=worker)
@@ -329,50 +426,53 @@ class BeanQueue:
 
         worker_id = worker.id
 
+        # Determine the number of worker threads
+        max_workers = self.config.MAX_WORKER_THREADS
+        if max_workers == 0:
+            max_workers = None  # Default to (num_cpus * 5)
+
+        # Create thread pool executor for concurrent task processing
+        executor = None
+        if max_workers != 1:
+            executor = ThreadPoolExecutor(
+                max_workers=max_workers, thread_name_prefix="task_worker"
+            )
+            logger.info("Created thread pool executor with max_workers=%s", max_workers)
+
         try:
-            while True:
-                while True:
-                    tasks = dispatch_service.dispatch(
-                        channels,
-                        worker_id=worker_id,
-                        limit=self.config.BATCH_SIZE,
-                    ).all()
-                    for task in tasks:
-                        logger.info(
-                            "Processing task %s, channel=%s, module=%s, func=%s",
-                            task.id,
-                            task.channel,
-                            task.module,
-                            task.func_name,
-                        )
-                        # TODO: support processor pool and other approaches to dispatch the workload
-                        registry.process(task, event_cls=self.event_model)
-                    if not tasks:
-                        # we should try to keep dispatching until we cannot find tasks
-                        break
-                    else:
-                        db.commit()
-                # we will not see notifications in a transaction, need to close the transaction first before entering
-                # polling
-                db.close()
-                try:
-                    for notification in dispatch_service.poll(
-                        timeout=self.config.POLL_TIMEOUT
-                    ):
-                        logger.debug("Receive notification %s", notification)
-                except TimeoutError:
-                    logger.debug("Poll timeout, try again")
-                    continue
+            if executor is not None:
+                # Threaded processing with continuous task feeding
+                self._process_tasks_threaded(
+                    db=db,
+                    executor=executor,
+                    dispatch_service=dispatch_service,
+                    registry=registry,
+                    channels=channels,
+                    worker_id=worker_id,
+                )
+            else:
+                # Sequential processing (original behavior)
+                self._process_tasks_sequential(
+                    db=db,
+                    dispatch_service=dispatch_service,
+                    registry=registry,
+                    channels=channels,
+                    worker_id=worker_id,
+                )
         except (SystemExit, KeyboardInterrupt):
             db.rollback()
             logger.info("Shutting down ...")
+
+            # Shutdown the executor if it was created
+            if executor is not None:
+                logger.info("Shutting down thread pool executor...")
+                executor.shutdown(wait=True, cancel_futures=False)
+                logger.info("Thread pool executor shutdown complete")
+
             self._worker_update_shutdown_event.set()
             worker_update_thread.join(5)
-            if metrics_server_thread is not None:
-                # set a threading event, waits until server is shutdown
-                # serve the ongoing requests
-                self._metrics_server_shutdown()
-                metrics_server_thread.join(1)
+            if self._metrics_server is not None:
+                self._metrics_server.shutdown()
 
         worker.state = models.WorkerState.SHUTDOWN
         db.add(worker)

{beanqueue-1.1.9 → beanqueue-2.0.0rc0}/bq/config.py

@@ -1,3 +1,4 @@
+import json
 import typing
 
 from pydantic import Field
@@ -16,6 +17,11 @@ class Config(BaseSettings):
     # Size of tasks batch to fetch each time from the database
     BATCH_SIZE: int = 1
 
+    # Maximum number of worker threads for concurrent task processing
+    # Set to 1 to disable thread pool and process tasks sequentially
+    # Set to 0 to use the default (number of CPUs * 5)
+    MAX_WORKER_THREADS: int = 1
+
     # How long we should poll before timeout in seconds
     POLL_TIMEOUT: int = 60
 
@@ -35,7 +41,7 @@ class Config(BaseSettings):
     EVENT_MODEL: str | None = "bq.Event"
 
     # Enable metrics HTTP server
-    METRICS_HTTP_SERVER_ENABLED: bool = True
+    METRICS_HTTP_SERVER_ENABLED: bool = False
 
     # the metrics http server interface to listen
     METRICS_HTTP_SERVER_INTERFACE: str = ""
@@ -46,6 +52,11 @@ class Config(BaseSettings):
     # default log level for metrics http server
     METRICS_HTTP_SERVER_LOG_LEVEL: int = 30
 
+    # Optional logging.config dict for the metrics HTTP server (uvicorn).
+    # When unset, a default config is used. Pass a dict programmatically or
+    # JSON via BQ_METRICS_HTTP_SERVER_LOG_CONFIG.
+    METRICS_HTTP_SERVER_LOG_CONFIG: dict[str, typing.Any] | None = None
+
     POSTGRES_SERVER: str = "localhost"
     POSTGRES_USER: str = "bq"
     POSTGRES_PASSWORD: str = ""
@@ -53,6 +64,15 @@ class Config(BaseSettings):
     # The URL of postgresql database to connect
     DATABASE_URL: typing.Optional[PostgresDsn] = None
 
+    @field_validator("METRICS_HTTP_SERVER_LOG_CONFIG", mode="before")
+    @classmethod
+    def parse_metrics_log_config(cls, v: typing.Any) -> typing.Any:
+        if v is None or isinstance(v, dict):
+            return v
+        if isinstance(v, str):
+            return json.loads(v)
+        raise ValueError("Unexpected METRICS_HTTP_SERVER_LOG_CONFIG type")
+
     @field_validator("DATABASE_URL", mode="before")
     def assemble_db_connection(
         cls, v: typing.Optional[str], info: ValidationInfo

{beanqueue-1.1.9 → beanqueue-2.0.0rc0}/bq/events.py

@@ -3,3 +3,5 @@ import blinker
 worker_init = blinker.signal("worker-init")
 
 task_failure = blinker.signal("task-failure")
+
+healthz_check = blinker.signal("healthz-check")

beanqueue-2.0.0rc0/bq/metrics.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+import copy
+import logging.config
+import threading
+import typing
+from collections.abc import Callable
+from collections.abc import Coroutine
+from importlib.util import find_spec
+
+from sqlalchemy.orm import Session as DBSession
+
+from . import events
+from . import models
+
+if typing.TYPE_CHECKING:
+    from .app import BeanQueue
+    from .config import Config
+
+logger = logging.getLogger(__name__)
+
+METRICS_EXTRA = "metrics"
+METRICS_SERVER_LOGGER = "metrics_server"
+
+
+def _healthz_sync_wrapper(
+    func: Callable[..., typing.Any],
+) -> Callable[..., Coroutine[typing.Any, typing.Any, typing.Any]]:
+    async def wrapper(sender: typing.Any, **kwargs: typing.Any) -> typing.Any:
+        return func(sender, **kwargs)
+
+    return wrapper
+
+
+class MetricsExtrasNotInstalledError(ImportError):
+    """Raised when metrics optional dependencies are not installed."""
+
+
+def require_metrics_extras() -> None:
+    missing = []
+    if find_spec("starlette") is None:
+        missing.append("starlette")
+    if find_spec("uvicorn") is None:
+        missing.append("uvicorn")
+    if missing:
+        raise MetricsExtrasNotInstalledError(
+            "Health check and metrics HTTP server require optional dependencies "
+            f"({', '.join(missing)}). "
+            f"Install them with: pip install beanqueue[{METRICS_EXTRA}]"
+        )
+
+
+def default_metrics_log_config(log_level: int) -> dict[str, typing.Any]:
+    from uvicorn.config import LOGGING_CONFIG
+
+    level_name = logging.getLevelName(log_level)
+    info_level = logging.getLevelName(logging.INFO)
+    log_config = copy.deepcopy(LOGGING_CONFIG)
+    log_config["handlers"]["access"]["stream"] = "ext://sys.stderr"
+    log_config["loggers"]["uvicorn"]["level"] = level_name
+    log_config["loggers"]["uvicorn.error"]["level"] = level_name
+    # Access lines are logged at INFO by uvicorn regardless of server log level.
+    log_config["loggers"]["uvicorn.access"]["level"] = info_level
+    log_config["loggers"][METRICS_SERVER_LOGGER] = {
+        "handlers": ["default"],
+        "level": info_level,
+        "propagate": False,
+    }
+    return log_config
+
+
+def resolve_metrics_log_config(config: Config) -> dict[str, typing.Any]:
+    if config.METRICS_HTTP_SERVER_LOG_CONFIG is not None:
+        return config.METRICS_HTTP_SERVER_LOG_CONFIG
+    return default_metrics_log_config(config.METRICS_HTTP_SERVER_LOG_LEVEL)
+
+
+class MetricsServer:
+    def __init__(self, bq: BeanQueue, worker_id: typing.Any):
+        require_metrics_extras()
+        self._bq = bq
+        self._worker_id = worker_id
+        self._server = None
+        self._thread: threading.Thread | None = None
+
+    def _has_custom_health_checks(self) -> bool:
+        return bool(events.healthz_check.receivers)
+
+    async def _run_healthz_checks(
+        self, worker: models.Worker, session: DBSession, body: dict[str, typing.Any]
+    ) -> bool:
+        try:
+            await events.healthz_check.send_async(
+                self._bq,
+                _sync_wrapper=_healthz_sync_wrapper,
+                worker=worker,
+                session=session,
+            )
+        except Exception as exc:
+            logger.exception("Custom healthz check failed")
+            body["error"] = str(exc)
+            return False
+        return True
+
+    async def check_healthz(self) -> tuple[bool, dict[str, typing.Any]]:
+        body: dict[str, typing.Any] = {"status": "ok"}
+
+        if not self._has_custom_health_checks():
+            return True, body
+
+        with self._bq.make_session() as db:
+            worker_service = self._bq._make_worker_service(db)
+            worker = worker_service.get_worker(self._worker_id)
+            body["worker_id"] = str(self._worker_id)
+
+            if not await self._run_healthz_checks(worker, db, body):
+                body["status"] = "internal error"
+                return False, body
+        return True, body
+
+    def create_app(self):
+        from starlette.applications import Starlette
+        from starlette.responses import JSONResponse
+        from starlette.routing import Route
+
+        async def healthz(_request):
+            ok, body = await self.check_healthz()
+            return JSONResponse(body, status_code=200 if ok else 500)
+
+        return Starlette(
+            routes=[
+                Route("/healthz", healthz),
+            ]
+        )
+
+    def start(self) -> None:
+        import uvicorn
+
+        require_metrics_extras()
+        host = self._bq.config.METRICS_HTTP_SERVER_INTERFACE
+        port = self._bq.config.METRICS_HTTP_SERVER_PORT
+        log_config = resolve_metrics_log_config(self._bq.config)
+        logging.config.dictConfig(log_config)
+
+        app = self.create_app()
+        # log_level is intentionally omitted: uvicorn would override logger levels
+        # from log_config (including uvicorn.access) after configure_logging().
+        config = uvicorn.Config(
+            app,
+            host=host,
+            port=port,
+            log_config=log_config,
+            access_log=True,
+        )
+        self._server = uvicorn.Server(config)
+
+        def run() -> None:
+            logging.getLogger(METRICS_SERVER_LOGGER).info(
+                "Run metrics HTTP server on %s:%s", host, port
+            )
+            self._server.run()
+
+        self._thread = threading.Thread(target=run, name="metrics_server")
+        self._thread.daemon = True
+        self._thread.start()
+
+    def shutdown(self) -> None:
+        if self._server is not None:
+            self._server.should_exit = True
+        if self._thread is not None:
+            self._thread.join(1)

{beanqueue-1.1.9 → beanqueue-2.0.0rc0}/pyproject.toml

@@ -1,9 +1,9 @@
 [project]
 name = "beanqueue"
-version = "1.1.9"
+version = "2.0.0rc0"
 description = "BeanQueue or BQ for short, PostgreSQL SKIP LOCK and SQLAlchemy based worker queue library"
 authors = [{ name = "Fang-Pen Lin", email = "fangpen@launchplatform.com" }]
-requires-python = "~=3.11"
+requires-python = ">=3.11,<4"
 readme = "README.md"
 license = "MIT"
 dependencies = [
@@ -18,10 +18,19 @@ dependencies = [
 [project.scripts]
 bq = "bq.cmds.main:cli"
 
+[project.optional-dependencies]
+metrics = [
+    "starlette>=0.27,<2",
+    "uvicorn>=0.30.0,<1",
+]
+
 [dependency-groups]
 dev = [
-    "psycopg2-binary>=2.9.9,<3",
+    "psycopg2-binary>=2.9.10,<3",
     "pytest-factoryboy>=2.7.0,<3",
+    "starlette>=0.27,<2",
+    "uvicorn>=0.30.0,<1",
+    "httpx>=0.27.0,<1",
 ]
 
 [tool.hatch.build.targets.sdist]
@@ -33,3 +42,6 @@ include = ["bq"]
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]