diff --git a/src/omnibase_infra/handlers/handler_db.py b/src/omnibase_infra/handlers/handler_db.py index a4e07ea60f..3bd7930cd5 100644 --- a/src/omnibase_infra/handlers/handler_db.py +++ b/src/omnibase_infra/handlers/handler_db.py @@ -70,6 +70,7 @@ from __future__ import annotations +import asyncio import logging from uuid import UUID, uuid4 @@ -126,6 +127,11 @@ ) _SUPPORTED_OPERATIONS: frozenset[str] = frozenset({"db.query", "db.execute"}) +# Per-handler pool close timeout (OMN-882) +# asyncpg pool.close() can take 30s timeout x pool_size connections worst case. +# This timeout prevents unbounded shutdown waits. +_POOL_CLOSE_TIMEOUT_SECONDS: float = 5.0 + # Error message prefixes for PostgreSQL errors # Used by _map_postgres_error to build descriptive error messages _POSTGRES_ERROR_PREFIXES: dict[type[asyncpg.PostgresError], str] = { @@ -396,7 +402,14 @@ async def initialize(self, config: dict[str, object]) -> None: ) from e async def shutdown(self) -> None: - """Close database connection pool and release resources.""" + """Close database connection pool and release resources. + + The pool close is wrapped with a timeout to prevent unbounded shutdown + waits. asyncpg pool.close() can take up to 30s per connection x pool_size + in the worst case (e.g., 150s for 5 connections). The timeout ensures + predictable shutdown behavior. If the pool close times out, the pool + reference is still cleared to allow garbage collection. + """ # Reset circuit breaker state if self._circuit_breaker_initialized: async with self._circuit_breaker_lock: @@ -404,8 +417,22 @@ async def shutdown(self) -> None: self._circuit_breaker_initialized = False if self._pool is not None: - await self._pool.close() - self._pool = None + try: + await asyncio.wait_for( + self._pool.close(), timeout=_POOL_CLOSE_TIMEOUT_SECONDS + ) + except TimeoutError: + logger.warning( + "Database pool close timed out, terminating connections", + extra={ + "timeout_seconds": _POOL_CLOSE_TIMEOUT_SECONDS, + "pool_size": self._pool_size, + }, + ) + # Terminate remaining connections forcefully + self._pool.terminate() + finally: + self._pool = None self._initialized = False logger.info("HandlerDb shutdown complete") diff --git a/src/omnibase_infra/handlers/handler_http.py b/src/omnibase_infra/handlers/handler_http.py index ed1325f5c0..9f5560fc5e 100644 --- a/src/omnibase_infra/handlers/handler_http.py +++ b/src/omnibase_infra/handlers/handler_http.py @@ -18,6 +18,7 @@ from __future__ import annotations +import asyncio import json import logging from uuid import UUID, uuid4 @@ -75,6 +76,9 @@ # Streaming chunk size for responses without Content-Length header _STREAMING_CHUNK_SIZE: int = 8192 # 8 KB chunks +# Per-handler client close timeout (OMN-882) +_CLIENT_CLOSE_TIMEOUT_SECONDS: float = 5.0 + # Handler ID for ModelHandlerOutput HANDLER_ID_HTTP: str = "http-handler" @@ -296,10 +300,24 @@ async def initialize(self, config: dict[str, object]) -> None: ) from e async def shutdown(self) -> None: - """Close HTTP client and release resources.""" + """Close HTTP client and release resources. + + The client close is wrapped with a timeout to prevent unbounded shutdown + waits if the underlying connection pool has hung connections. + """ if self._client is not None: - await self._client.aclose() - self._client = None + try: + await asyncio.wait_for( + self._client.aclose(), + timeout=_CLIENT_CLOSE_TIMEOUT_SECONDS, + ) + except TimeoutError: + logger.warning( + "HTTP client close timed out", + extra={"timeout_seconds": _CLIENT_CLOSE_TIMEOUT_SECONDS}, + ) + finally: + self._client = None self._initialized = False logger.info("HandlerHttpRest shutdown complete") diff --git a/src/omnibase_infra/runtime/models/model_shutdown_config.py b/src/omnibase_infra/runtime/models/model_shutdown_config.py index a738026bd5..ff75901e88 100644 --- a/src/omnibase_infra/runtime/models/model_shutdown_config.py +++ b/src/omnibase_infra/runtime/models/model_shutdown_config.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: MIT # Copyright (c) 2025 OmniNode Team -"""Shutdown Configuration Model. - -The Pydantic model for shutdown configuration. -""" +"""Shutdown configuration for ONEX runtime graceful shutdown.""" from __future__ import annotations @@ -56,23 +53,30 @@ def _coerce_grace_period(v: object) -> int: class ModelShutdownConfig(BaseModel): """Shutdown configuration model. - Defines graceful shutdown parameters. + Defines graceful shutdown parameters including per-handler timeout controls. Attributes: - grace_period_seconds: Time in seconds to wait for graceful shutdown. + grace_period_seconds: Time in seconds to wait for overall graceful shutdown. Must be >= 0. A value of 0 means immediate shutdown with no grace period (use with caution as in-flight operations may be interrupted). + handler_shutdown_timeout_seconds: Maximum time in seconds to wait for each + individual handler's shutdown method to complete. If a handler exceeds + this timeout, the shutdown is considered failed for that handler but + other handlers continue their shutdown. This prevents a single slow + handler (e.g., asyncpg pool close with 30s per-connection timeout) + from blocking the entire shutdown sequence. Edge Cases: - - 0: Immediate shutdown, no waiting for in-flight operations - - Values > 3600: Rejected by Pydantic validation (le=3600 constraint); - consider using reasonable timeouts (30-120 seconds recommended) + - grace_period_seconds=0: Immediate shutdown, no waiting for in-flight operations + - Values > 3600: Rejected by Pydantic validation (le=3600 constraint) - Negative values: Rejected by Pydantic validation (ge=0 constraint) + - handler_shutdown_timeout_seconds=0: Not allowed (ge=1 constraint) Production Recommendation: - Set grace_period_seconds between 30-120 seconds for production deployments - to allow sufficient time for in-flight operations while preventing excessive - delays during shutdown sequences. + Set grace_period_seconds between 30-120 seconds for production deployments. + Set handler_shutdown_timeout_seconds to 10-15 seconds to prevent any single + handler from consuming the entire grace period. The handler timeout should + always be less than the overall grace period. """ model_config = ConfigDict( @@ -90,5 +94,16 @@ class ModelShutdownConfig(BaseModel): description="Time in seconds to wait for graceful shutdown (0-3600)", ) + handler_shutdown_timeout_seconds: float = Field( + default=10.0, + ge=1.0, + le=300.0, # Max 5 minutes per handler + strict=False, + description=( + "Maximum time in seconds to wait for each individual handler shutdown. " + "Prevents a single slow handler from blocking the entire shutdown sequence." + ), + ) + __all__: list[str] = ["ModelShutdownConfig"] diff --git a/src/omnibase_infra/runtime/protocol_lifecycle_executor.py b/src/omnibase_infra/runtime/protocol_lifecycle_executor.py index cf25e8f5c7..333cd28882 100644 --- a/src/omnibase_infra/runtime/protocol_lifecycle_executor.py +++ b/src/omnibase_infra/runtime/protocol_lifecycle_executor.py @@ -2,13 +2,9 @@ # Copyright (c) 2025 OmniNode Team """Protocol lifecycle executor for ONEX Infrastructure. -The ProtocolLifecycleExecutor class for managing protocol -handler lifecycle operations including shutdown, health checks, and priority -ordering. - -The ProtocolLifecycleExecutor is a reusable utility extracted from -RuntimeHostProcess to enable lifecycle management in other contexts (e.g., test -fixtures, standalone management, future node implementations). +Reusable utility extracted from RuntimeHostProcess to manage protocol handler +lifecycle operations (shutdown, health checks, priority ordering) in multiple +contexts: runtime host, test fixtures, standalone management. Example Usage: ```python @@ -58,6 +54,11 @@ MAX_HEALTH_CHECK_TIMEOUT = 60.0 DEFAULT_HEALTH_CHECK_TIMEOUT = 5.0 +# Handler shutdown timeout bounds (per OMN-882) +MIN_HANDLER_SHUTDOWN_TIMEOUT = 1.0 +MAX_HANDLER_SHUTDOWN_TIMEOUT = 300.0 +DEFAULT_HANDLER_SHUTDOWN_TIMEOUT = 10.0 + class ProtocolLifecycleExecutor: """Executes protocol lifecycle operations (shutdown, health check, priority). @@ -65,10 +66,9 @@ class ProtocolLifecycleExecutor: Reusable utility for managing protocol lifecycles. Used by RuntimeHostProcess and can be used by future node implementations. - Exports: - - Priority-based shutdown ordering (higher priority handlers shutdown first) - - Health check execution with configurable timeouts - - Parallel shutdown within priority groups for performance + Capabilities: priority-based shutdown ordering (higher priority first), + health check execution with configurable timeouts, parallel shutdown + within priority groups. Thread Safety: This class is thread-safe for concurrent method calls. @@ -106,6 +106,7 @@ class ProtocolLifecycleExecutor: Attributes: health_check_timeout_seconds: Default timeout for health checks (1-60 seconds). + handler_shutdown_timeout_seconds: Per-handler shutdown timeout (1-300 seconds). Example: ```python @@ -129,7 +130,9 @@ class ProtocolLifecycleExecutor: """ def __init__( - self, health_check_timeout_seconds: float = DEFAULT_HEALTH_CHECK_TIMEOUT + self, + health_check_timeout_seconds: float = DEFAULT_HEALTH_CHECK_TIMEOUT, + handler_shutdown_timeout_seconds: float = DEFAULT_HANDLER_SHUTDOWN_TIMEOUT, ) -> None: """Initialize the protocol lifecycle executor. @@ -137,8 +140,14 @@ def __init__( health_check_timeout_seconds: Default timeout for health checks. Valid range: 1-60 seconds per ModelLifecycleSubcontract. Values outside this range are clamped with a warning. + handler_shutdown_timeout_seconds: Per-handler shutdown timeout. + Valid range: 1-300 seconds. Each individual handler's shutdown() + call is wrapped with asyncio.wait_for() using this timeout. + Prevents a single slow handler (e.g., asyncpg pool close) from + blocking the entire shutdown sequence. + Values outside this range are clamped with a warning. """ - # Validate and clamp timeout to valid bounds + # Validate and clamp health check timeout to valid bounds if ( health_check_timeout_seconds < MIN_HEALTH_CHECK_TIMEOUT or health_check_timeout_seconds > MAX_HEALTH_CHECK_TIMEOUT @@ -160,7 +169,33 @@ def __init__( min(health_check_timeout_seconds, MAX_HEALTH_CHECK_TIMEOUT), ) + # Validate and clamp handler shutdown timeout to valid bounds + if ( + handler_shutdown_timeout_seconds < MIN_HANDLER_SHUTDOWN_TIMEOUT + or handler_shutdown_timeout_seconds > MAX_HANDLER_SHUTDOWN_TIMEOUT + ): + logger.warning( + "handler_shutdown_timeout_seconds out of valid range, clamping", + extra={ + "original_value": handler_shutdown_timeout_seconds, + "min_value": MIN_HANDLER_SHUTDOWN_TIMEOUT, + "max_value": MAX_HANDLER_SHUTDOWN_TIMEOUT, + "clamped_value": max( + MIN_HANDLER_SHUTDOWN_TIMEOUT, + min( + handler_shutdown_timeout_seconds, + MAX_HANDLER_SHUTDOWN_TIMEOUT, + ), + ), + }, + ) + handler_shutdown_timeout_seconds = max( + MIN_HANDLER_SHUTDOWN_TIMEOUT, + min(handler_shutdown_timeout_seconds, MAX_HANDLER_SHUTDOWN_TIMEOUT), + ) + self._health_check_timeout_seconds: float = health_check_timeout_seconds + self._handler_shutdown_timeout_seconds: float = handler_shutdown_timeout_seconds @property def health_check_timeout_seconds(self) -> float: @@ -171,6 +206,15 @@ def health_check_timeout_seconds(self) -> float: """ return self._health_check_timeout_seconds + @property + def handler_shutdown_timeout_seconds(self) -> float: + """Return the configured per-handler shutdown timeout. + + Returns: + The timeout in seconds for individual handler shutdowns. + """ + return self._handler_shutdown_timeout_seconds + @staticmethod def get_shutdown_priority(handler: ProtocolContainerAware) -> int: """Get shutdown priority for a handler. @@ -226,26 +270,42 @@ async def shutdown_handler( self, handler_type: str, handler: ProtocolContainerAware, + timeout_seconds: float = -1.0, ) -> ModelLifecycleResult: - """Shutdown a single handler with error handling. + """Shutdown a single handler with per-handler timeout and error handling. - This method performs individual handler shutdown with comprehensive error - handling to ensure all handlers get a chance to cleanup even if one fails. + This method performs individual handler shutdown wrapped with + asyncio.wait_for() to enforce a per-handler timeout. This prevents a + single slow handler (e.g., asyncpg pool close with 30s per-connection + timeout x 5 connections = 150s worst case) from blocking the entire + shutdown sequence. Args: handler_type: The handler type identifier. handler: The handler instance to shutdown. + timeout_seconds: Override timeout for this specific shutdown. If + negative (default: -1.0), uses the configured + handler_shutdown_timeout_seconds. Returns: ModelLifecycleResult with handler_type, success status, and optional - error_message if the shutdown failed. + error_message if the shutdown failed or timed out. """ + effective_timeout = ( + timeout_seconds + if timeout_seconds > 0 + else self._handler_shutdown_timeout_seconds + ) + try: if hasattr(handler, "shutdown"): - await handler.shutdown() + await asyncio.wait_for(handler.shutdown(), timeout=effective_timeout) logger.debug( "Handler shutdown completed", - extra={"handler_type": handler_type}, + extra={ + "handler_type": handler_type, + "timeout_seconds": effective_timeout, + }, ) return ModelLifecycleResult.succeeded(handler_type) else: @@ -255,6 +315,19 @@ async def shutdown_handler( extra={"handler_type": handler_type}, ) return ModelLifecycleResult.succeeded(handler_type) + except TimeoutError: + # Handler shutdown exceeded the per-handler timeout + logger.warning( + "Handler shutdown timed out", + extra={ + "handler_type": handler_type, + "timeout_seconds": effective_timeout, + }, + ) + return ModelLifecycleResult.failed( + handler_type, + f"Shutdown timed out after {effective_timeout}s", + ) except Exception as e: # Log exception but return failure status instead of raising # This ensures all handlers get a chance to cleanup even if one fails @@ -428,8 +501,11 @@ async def shutdown_handlers_by_priority( __all__: list[str] = [ + "DEFAULT_HANDLER_SHUTDOWN_TIMEOUT", "DEFAULT_HEALTH_CHECK_TIMEOUT", + "MAX_HANDLER_SHUTDOWN_TIMEOUT", "MAX_HEALTH_CHECK_TIMEOUT", + "MIN_HANDLER_SHUTDOWN_TIMEOUT", "MIN_HEALTH_CHECK_TIMEOUT", "ProtocolLifecycleExecutor", ] diff --git a/src/omnibase_infra/runtime/service_runtime_host_process.py b/src/omnibase_infra/runtime/service_runtime_host_process.py index 49009e675c..c417a77d75 100644 --- a/src/omnibase_infra/runtime/service_runtime_host_process.py +++ b/src/omnibase_infra/runtime/service_runtime_host_process.py @@ -107,7 +107,10 @@ async def main() -> None: from omnibase_infra.runtime.models.model_materialized_resources import ( ModelMaterializedResources, ) -from omnibase_infra.runtime.protocol_lifecycle_executor import ProtocolLifecycleExecutor +from omnibase_infra.runtime.protocol_lifecycle_executor import ( + DEFAULT_HANDLER_SHUTDOWN_TIMEOUT, + ProtocolLifecycleExecutor, +) from omnibase_infra.runtime.runtime_contract_config_loader import ( RuntimeContractConfigLoader, ) @@ -165,7 +168,12 @@ async def main() -> None: from omnibase_infra.runtime.handler_plugin_loader import HandlerPluginLoader from omnibase_infra.runtime.kafka_contract_source import KafkaContractSource from omnibase_infra.runtime.protocol_contract_source import ProtocolContractSource -from omnibase_infra.topics import TopicResolutionError, TopicResolver +from omnibase_infra.topics import ( + SUFFIX_CONTRACT_DEREGISTERED, + SUFFIX_CONTRACT_REGISTERED, + TopicResolutionError, + TopicResolver, +) # Expose wire_default_handlers as wire_handlers for test patching compatibility # Tests patch "omnibase_infra.runtime.service_runtime_host_process.wire_handlers" @@ -887,9 +895,19 @@ async def create_runtime() -> RuntimeHostProcess: max_concurrent_value ) + # Per-handler shutdown timeout (OMN-882) + # Prevents a single slow handler from blocking entire shutdown sequence + _handler_shutdown_raw = config.get("handler_shutdown_timeout_seconds") + handler_shutdown_value: float = DEFAULT_HANDLER_SHUTDOWN_TIMEOUT + if isinstance(_handler_shutdown_raw, int | float): + handler_shutdown_value = float(_handler_shutdown_raw) + + self._handler_shutdown_timeout_seconds: float = handler_shutdown_value + # Handler executor for lifecycle operations (shutdown, health check) self._lifecycle_executor = ProtocolLifecycleExecutor( - health_check_timeout_seconds=self._health_check_timeout_seconds + health_check_timeout_seconds=self._health_check_timeout_seconds, + handler_shutdown_timeout_seconds=self._handler_shutdown_timeout_seconds, ) # Store full config for handler initialization diff --git a/tests/unit/runtime/test_protocol_lifecycle_executor.py b/tests/unit/runtime/test_protocol_lifecycle_executor.py index df8d3fe578..c095413182 100644 --- a/tests/unit/runtime/test_protocol_lifecycle_executor.py +++ b/tests/unit/runtime/test_protocol_lifecycle_executor.py @@ -213,6 +213,30 @@ async def shutdown(self) -> None: raise RuntimeError("Shutdown failed") +class MockHandlerWithSlowShutdown: + """Mock handler with slow shutdown for timeout testing.""" + + def __init__( + self, + handler_type: str = "slow_shutdown", + delay_seconds: float = 10.0, + ) -> None: + """Initialize mock handler. + + Args: + handler_type: The type identifier for this handler. + delay_seconds: How long shutdown should take. + """ + self.handler_type = handler_type + self.delay_seconds = delay_seconds + self.shutdown_called = False + + async def shutdown(self) -> None: + """Simulate slow shutdown.""" + self.shutdown_called = True + await asyncio.sleep(self.delay_seconds) + + # ============================================================================= # Test Fixtures # ============================================================================= @@ -757,6 +781,141 @@ async def test_shutdown_handlers_logs_priority_groups( ) +# ============================================================================= +# TestProtocolLifecycleExecutorShutdownTimeout (OMN-882) +# ============================================================================= + + +class TestProtocolLifecycleExecutorShutdownTimeout: + """Test per-handler shutdown timeout behavior (OMN-882).""" + + def test_default_handler_shutdown_timeout(self) -> None: + """Test that ProtocolLifecycleExecutor initializes with default handler shutdown timeout.""" + executor = ProtocolLifecycleExecutor() + + assert executor.handler_shutdown_timeout_seconds == 10.0 + + def test_custom_handler_shutdown_timeout(self) -> None: + """Test that ProtocolLifecycleExecutor accepts custom handler shutdown timeout.""" + executor = ProtocolLifecycleExecutor(handler_shutdown_timeout_seconds=15.0) + + assert executor.handler_shutdown_timeout_seconds == 15.0 + + def test_handler_shutdown_timeout_clamped_below_minimum( + self, + caplog: pytest.LogCaptureFixture, + ) -> None: + """Test that handler shutdown timeout below minimum is clamped.""" + with caplog.at_level(logging.WARNING): + executor = ProtocolLifecycleExecutor( + handler_shutdown_timeout_seconds=0.1, + ) + + assert executor.handler_shutdown_timeout_seconds == 1.0 + assert any( + "handler_shutdown_timeout_seconds out of valid range" in r.message + for r in caplog.records + ) + + def test_handler_shutdown_timeout_clamped_above_maximum( + self, + caplog: pytest.LogCaptureFixture, + ) -> None: + """Test that handler shutdown timeout above maximum is clamped.""" + with caplog.at_level(logging.WARNING): + executor = ProtocolLifecycleExecutor( + handler_shutdown_timeout_seconds=500.0, + ) + + assert executor.handler_shutdown_timeout_seconds == 300.0 + assert any( + "handler_shutdown_timeout_seconds out of valid range" in r.message + for r in caplog.records + ) + + @pytest.mark.asyncio + async def test_shutdown_handler_times_out_slow_handler(self) -> None: + """Test that slow handler shutdown is terminated by per-handler timeout.""" + executor = ProtocolLifecycleExecutor(handler_shutdown_timeout_seconds=0.5) + handler = MockHandlerWithSlowShutdown(handler_type="slow", delay_seconds=10.0) + + start = time.monotonic() + result = await executor.shutdown_handler("slow", handler) # type: ignore[arg-type] + elapsed = time.monotonic() - start + + assert result.handler_type == "slow" + assert result.success is False + assert "timed out" in result.error_message.lower() + # Should complete near the timeout, not the full delay + assert elapsed < 2.0, f"Expected timeout around 0.5s, took {elapsed}s" + + @pytest.mark.asyncio + async def test_shutdown_handler_with_override_timeout(self) -> None: + """Test that per-call timeout override works.""" + executor = ProtocolLifecycleExecutor(handler_shutdown_timeout_seconds=30.0) + handler = MockHandlerWithSlowShutdown(handler_type="slow", delay_seconds=10.0) + + start = time.monotonic() + # Override with very short timeout + result = await executor.shutdown_handler( + "slow", + handler, + timeout_seconds=0.5, # type: ignore[arg-type] + ) + elapsed = time.monotonic() - start + + assert result.success is False + assert "timed out" in result.error_message.lower() + assert elapsed < 2.0 + + @pytest.mark.asyncio + async def test_shutdown_handler_succeeds_within_timeout(self) -> None: + """Test that handler completing within timeout succeeds normally.""" + executor = ProtocolLifecycleExecutor(handler_shutdown_timeout_seconds=5.0) + handler = MockHandlerWithSlowShutdown(handler_type="quick", delay_seconds=0.01) + + result = await executor.shutdown_handler("quick", handler) # type: ignore[arg-type] + + assert result.handler_type == "quick" + assert result.success is True + assert handler.shutdown_called is True + + @pytest.mark.asyncio + async def test_shutdown_by_priority_with_timeout_slow_handler(self) -> None: + """Test that slow handler in priority group doesn't block other groups.""" + executor = ProtocolLifecycleExecutor(handler_shutdown_timeout_seconds=0.5) + + # Slow handler at high priority + slow_handler = MockHandlerWithSlowShutdown( + handler_type="slow", delay_seconds=10.0 + ) + + # Normal handler at low priority + normal_handler = MockHandler(handler_type="normal") + + handlers: dict[str, ProtocolHandler] = { + "slow": slow_handler, # type: ignore[dict-item] + "normal": normal_handler, # type: ignore[dict-item] + } + + start = time.monotonic() + result = await executor.shutdown_handlers_by_priority(handlers) + elapsed = time.monotonic() - start + + # Both handlers should be attempted + assert normal_handler.shutdown_called is True + assert slow_handler.shutdown_called is True + + # Slow handler should have timed out + assert result.has_failures + assert any( + "timed out" in f.error_message.lower() for f in result.failed_handlers + ) + + # Total time should be bounded by timeouts, not handler delays + assert elapsed < 3.0, f"Expected bounded shutdown, took {elapsed}s" + + # ============================================================================= # Additional Edge Case Tests # ============================================================================= @@ -861,6 +1020,7 @@ async def tracked_shutdown() -> None: "TestProtocolLifecycleExecutorInit", "TestProtocolLifecycleExecutorShutdownPriority", "TestProtocolLifecycleExecutorShutdown", + "TestProtocolLifecycleExecutorShutdownTimeout", "TestProtocolLifecycleExecutorHealthCheck", "TestProtocolLifecycleExecutorShutdownByPriority", "TestProtocolLifecycleExecutorEdgeCases", @@ -871,6 +1031,7 @@ async def tracked_shutdown() -> None: "MockHandlerWithoutShutdown", "MockHandlerWithoutHealthCheck", "MockHandlerWithSlowHealthCheck", + "MockHandlerWithSlowShutdown", "MockHandlerWithFailingHealthCheck", "MockHandlerWithFailingShutdown", ]