ccproxy-api 0.1.0__py3-none-any.whl

ccproxy/scheduler/exceptions.py

@@ -0,0 +1,34 @@
+"""Scheduler-specific exceptions."""
+
+
+class SchedulerError(Exception):
+    """Base exception for scheduler-related errors."""
+
+    pass
+
+
+class TaskRegistrationError(SchedulerError):
+    """Raised when task registration fails."""
+
+    pass
+
+
+class TaskNotFoundError(SchedulerError):
+    """Raised when attempting to access a task that doesn't exist."""
+
+    pass
+
+
+class TaskExecutionError(SchedulerError):
+    """Raised when task execution encounters an error."""
+
+    def __init__(self, task_name: str, original_error: Exception):
+        self.task_name = task_name
+        self.original_error = original_error
+        super().__init__(f"Task '{task_name}' execution failed: {original_error}")
+
+
+class SchedulerShutdownError(SchedulerError):
+    """Raised when scheduler shutdown encounters an error."""
+
+    pass

ccproxy/scheduler/manager.py

@@ -0,0 +1,186 @@
+"""Scheduler management for FastAPI integration."""
+
+import asyncio
+from typing import Any
+
+import structlog
+
+from ccproxy.config.settings import Settings
+
+from .core import Scheduler, get_scheduler
+from .registry import register_task
+from .tasks import PricingCacheUpdateTask, PushgatewayTask, StatsPrintingTask
+
+
+logger = structlog.get_logger(__name__)
+
+
+async def setup_scheduler_tasks(scheduler: Scheduler, settings: Settings) -> None:
+    """
+    Setup and configure all scheduler tasks based on settings.
+
+    Args:
+        scheduler: Scheduler instance
+        settings: Application settings
+    """
+    scheduler_config = settings.scheduler
+
+    if not scheduler_config.enabled:
+        logger.info("scheduler_disabled")
+        return
+
+    # Add pushgateway task if enabled
+    if scheduler_config.pushgateway_enabled:
+        try:
+            await scheduler.add_task(
+                task_name="pushgateway",
+                task_type="pushgateway",
+                interval_seconds=scheduler_config.pushgateway_interval_seconds,
+                enabled=True,
+                max_backoff_seconds=scheduler_config.pushgateway_max_backoff_seconds,
+            )
+            logger.info(
+                "pushgateway_task_added",
+                interval_seconds=scheduler_config.pushgateway_interval_seconds,
+            )
+        except Exception as e:
+            logger.error(
+                "pushgateway_task_add_failed",
+                error=str(e),
+                error_type=type(e).__name__,
+            )
+
+    # Add stats printing task if enabled
+    if scheduler_config.stats_printing_enabled:
+        try:
+            await scheduler.add_task(
+                task_name="stats_printing",
+                task_type="stats_printing",
+                interval_seconds=scheduler_config.stats_printing_interval_seconds,
+                enabled=True,
+            )
+            logger.info(
+                "stats_printing_task_added",
+                interval_seconds=scheduler_config.stats_printing_interval_seconds,
+            )
+        except Exception as e:
+            logger.error(
+                "stats_printing_task_add_failed",
+                error=str(e),
+                error_type=type(e).__name__,
+            )
+
+    # Add pricing cache update task if enabled
+    if scheduler_config.pricing_update_enabled:
+        try:
+            # Convert hours to seconds
+            interval_seconds = scheduler_config.pricing_update_interval_hours * 3600
+
+            await scheduler.add_task(
+                task_name="pricing_cache_update",
+                task_type="pricing_cache_update",
+                interval_seconds=interval_seconds,
+                enabled=True,
+                force_refresh_on_startup=scheduler_config.pricing_force_refresh_on_startup,
+            )
+            logger.info(
+                "pricing_update_task_added",
+                interval_hours=scheduler_config.pricing_update_interval_hours,
+                force_refresh_on_startup=scheduler_config.pricing_force_refresh_on_startup,
+            )
+        except Exception as e:
+            logger.error(
+                "pricing_update_task_add_failed",
+                error=str(e),
+                error_type=type(e).__name__,
+            )
+
+
+def _register_default_tasks() -> None:
+    """Register default task types in the global registry."""
+    from .registry import get_task_registry
+
+    registry = get_task_registry()
+
+    # Only register if not already registered
+    if not registry.is_registered("pushgateway"):
+        register_task("pushgateway", PushgatewayTask)
+    if not registry.is_registered("stats_printing"):
+        register_task("stats_printing", StatsPrintingTask)
+    if not registry.is_registered("pricing_cache_update"):
+        register_task("pricing_cache_update", PricingCacheUpdateTask)
+
+
+async def start_scheduler(settings: Settings) -> Scheduler | None:
+    """
+    Start the scheduler with configured tasks.
+
+    Args:
+        settings: Application settings
+
+    Returns:
+        Scheduler instance if successful, None otherwise
+    """
+    try:
+        if not settings.scheduler.enabled:
+            logger.info("scheduler_disabled")
+            return None
+
+        # Register task types (only when actually starting scheduler)
+        _register_default_tasks()
+
+        # Create scheduler with settings
+        scheduler = Scheduler(
+            max_concurrent_tasks=settings.scheduler.max_concurrent_tasks,
+            graceful_shutdown_timeout=settings.scheduler.graceful_shutdown_timeout,
+        )
+
+        # Start the scheduler
+        await scheduler.start()
+
+        # Setup tasks based on configuration
+        await setup_scheduler_tasks(scheduler, settings)
+
+        logger.info(
+            "scheduler_started",
+            max_concurrent_tasks=settings.scheduler.max_concurrent_tasks,
+            active_tasks=scheduler.task_count,
+            running_tasks=len(
+                [
+                    name
+                    for name in scheduler.list_tasks()
+                    if scheduler.get_task(name).is_running
+                ]
+            ),
+        )
+
+        return scheduler
+
+    except Exception as e:
+        logger.error(
+            "scheduler_start_failed",
+            error=str(e),
+            error_type=type(e).__name__,
+        )
+        return None
+
+
+async def stop_scheduler(scheduler: Scheduler | None) -> None:
+    """
+    Stop the scheduler gracefully.
+
+    Args:
+        scheduler: Scheduler instance to stop
+    """
+    if scheduler is None:
+        return
+
+    try:
+        await scheduler.stop()
+        logger.info("scheduler_stopped")
+    except Exception as e:
+        logger.error(
+            "scheduler_stop_failed",
+            error=str(e),
+            error_type=type(e).__name__,
+        )

ccproxy/scheduler/registry.py

@@ -0,0 +1,150 @@
+"""Task registry for dynamic task registration and discovery."""
+
+from typing import Any
+
+import structlog
+
+from .exceptions import TaskRegistrationError
+from .tasks import BaseScheduledTask
+
+
+logger = structlog.get_logger(__name__)
+
+
+class TaskRegistry:
+    """
+    Registry for managing scheduled task registration and discovery.
+
+    Provides a centralized way to register and retrieve scheduled tasks,
+    enabling dynamic task management and configuration.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the task registry."""
+        self._tasks: dict[str, type[BaseScheduledTask]] = {}
+
+    def register(self, name: str, task_class: type[BaseScheduledTask]) -> None:
+        """
+        Register a scheduled task class.
+
+        Args:
+            name: Unique name for the task
+            task_class: Task class that inherits from BaseScheduledTask
+
+        Raises:
+            TaskRegistrationError: If task name is already registered or invalid
+        """
+        if name in self._tasks:
+            raise TaskRegistrationError(f"Task '{name}' is already registered")
+
+        if not issubclass(task_class, BaseScheduledTask):
+            raise TaskRegistrationError(
+                f"Task class for '{name}' must inherit from BaseScheduledTask"
+            )
+
+        self._tasks[name] = task_class
+        logger.debug("task_registered", task_name=name, task_class=task_class.__name__)
+
+    def unregister(self, name: str) -> None:
+        """
+        Unregister a scheduled task.
+
+        Args:
+            name: Name of the task to unregister
+
+        Raises:
+            TaskRegistrationError: If task is not registered
+        """
+        if name not in self._tasks:
+            raise TaskRegistrationError(f"Task '{name}' is not registered")
+
+        del self._tasks[name]
+        logger.debug("task_unregistered", task_name=name)
+
+    def get(self, name: str) -> type[BaseScheduledTask]:
+        """
+        Get a registered task class by name.
+
+        Args:
+            name: Name of the task to retrieve
+
+        Returns:
+            Task class
+
+        Raises:
+            TaskRegistrationError: If task is not registered
+        """
+        if name not in self._tasks:
+            raise TaskRegistrationError(f"Task '{name}' is not registered")
+
+        return self._tasks[name]
+
+    def list_tasks(self) -> list[str]:
+        """
+        Get list of all registered task names.
+
+        Returns:
+            List of registered task names
+        """
+        return list(self._tasks.keys())
+
+    def is_registered(self, name: str) -> bool:
+        """
+        Check if a task is registered.
+
+        Args:
+            name: Task name to check
+
+        Returns:
+            True if task is registered, False otherwise
+        """
+        return name in self._tasks
+
+    def clear(self) -> None:
+        """Clear all registered tasks."""
+        self._tasks.clear()
+        logger.debug("task_registry_cleared")
+
+    def get_registry_info(self) -> dict[str, Any]:
+        """
+        Get information about the current registry state.
+
+        Returns:
+            Dictionary with registry information
+        """
+        return {
+            "total_tasks": len(self._tasks),
+            "registered_tasks": list(self._tasks.keys()),
+            "task_classes": {name: cls.__name__ for name, cls in self._tasks.items()},
+        }
+
+
+# Global task registry instance
+_global_registry: TaskRegistry | None = None
+
+
+def get_task_registry() -> TaskRegistry:
+    """
+    Get the global task registry instance.
+
+    Returns:
+        Global TaskRegistry instance
+    """
+    global _global_registry
+
+    if _global_registry is None:
+        _global_registry = TaskRegistry()
+
+    return _global_registry
+
+
+def register_task(name: str, task_class: type[BaseScheduledTask]) -> None:
+    """
+    Register a task in the global registry.
+
+    Args:
+        name: Unique name for the task
+        task_class: Task class that inherits from BaseScheduledTask
+    """
+    registry = get_task_registry()
+    registry.register(name, task_class)