Custom Workers¶

# worker/my_new_worker/config.py
import os


class Config:
    # Database
    DB_HOST = os.environ.get("DB_HOST", "mysql")
    DB_PORT = int(os.environ.get("DB_PORT", 3306))
    DB_NAME = os.environ.get("DB_NAME", "mailserver")
    DB_USER = os.environ.get("DB_USER", "mailuser")
    DB_PASSWORD = os.environ.get("DB_PASSWORD", "")

    # Redis
    REDIS_HOST = os.environ.get("REDIS_HOST", "redis")
    REDIS_PORT = int(os.environ.get("REDIS_PORT", 6379))

    # Worker-specific
    SERVICE_PORT = int(os.environ.get("PORT", 8090))
    POLL_INTERVAL = int(os.environ.get("POLL_INTERVAL", 60))
    WORKER_THREADS = int(os.environ.get("WORKER_THREADS", 2))

# worker/my_new_worker/service.py
import logging
import threading
import time

logger = logging.getLogger(__name__)


class MyNewWorkerService:
    def __init__(self, config, db, redis_client):
        self.config = config
        self.db = db
        self.redis = redis_client
        self._running = False
        self._thread = None

    def start(self):
        """Start the worker loop in a background thread."""
        self._running = True
        self._thread = threading.Thread(target=self._run_loop, daemon=True)
        self._thread.start()
        logger.info("MyNewWorker started (poll interval: %ds)", self.config.POLL_INTERVAL)

    def stop(self):
        """Stop the worker loop."""
        self._running = False
        if self._thread:
            self._thread.join(timeout=10)
        logger.info("MyNewWorker stopped")

    def _run_loop(self):
        """Main worker loop."""
        while self._running:
            try:
                self.process()
            except Exception as e:
                logger.error("Worker loop error: %s", e)
            time.sleep(self.config.POLL_INTERVAL)

    def process(self):
        """Override this with your actual work."""
        logger.debug("Processing...")
        # Your logic here:
        # - Query the database for items to process
        # - Do the work
        # - Update the database with results
        # - Emit metrics

    def health_check(self) -> dict:
        """Return worker health status."""
        return {
            "status": "healthy" if self._running else "unhealthy",
            "service": "my_new_worker",
            "thread_alive": self._thread.is_alive() if self._thread else False,
        }

# worker/my_new_worker/routes.py
from fastapi import FastAPI
from prometheus_client import generate_latest, CONTENT_TYPE_LATEST
from starlette.responses import Response


def create_app(service) -> FastAPI:
    app = FastAPI(title="My New Worker")

    @app.get("/health")
    async def health():
        return service.health_check()

    @app.get("/metrics")
    async def metrics():
        return Response(
            content=generate_latest(),
            media_type=CONTENT_TYPE_LATEST,
        )

    return app

# worker/my_new_worker/main.py
import logging
import uvicorn

from config import Config
from service import MyNewWorkerService
from routes import create_app

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)


def main():
    config = Config()

    # Initialize connections
    from shared.database import get_connection
    from shared.redis_client import get_redis

    db = get_connection(config)
    redis_client = get_redis(config)

    # Create and start the worker
    service = MyNewWorkerService(config, db, redis_client)
    service.start()

    # Create the HTTP app (health + metrics)
    app = create_app(service)

    # Run the HTTP server
    logger.info("Starting HTTP server on port %d", config.SERVICE_PORT)
    uvicorn.run(app, host="0.0.0.0", port=config.SERVICE_PORT)


if __name__ == "__main__":
    main()

# In service.py
from prometheus_client import Counter, Histogram, Gauge

# Define metrics
ITEMS_PROCESSED = Counter(
    "mailyte_my_worker_items_processed_total",
    "Total items processed",
    ["status"],
)
PROCESSING_DURATION = Histogram(
    "mailyte_my_worker_processing_seconds",
    "Time to process each item",
)
QUEUE_SIZE = Gauge(
    "mailyte_my_worker_queue_size",
    "Number of items waiting to be processed",
)


class MyNewWorkerService:
    def process(self):
        # Update queue size
        pending = self.db.execute("SELECT COUNT(*) FROM my_table WHERE status = 'pending'")
        QUEUE_SIZE.set(pending)

        # Process items
        with PROCESSING_DURATION.time():
            result = self.do_work()

        if result.success:
            ITEMS_PROCESSED.labels(status="success").inc()
        else:
            ITEMS_PROCESSED.labels(status="failure").inc()