dory-processor-sdk 0.0.1__py3-none-any.whl

dory/migration/transfer.py

@@ -0,0 +1,382 @@
+"""
+State transfer utilities with timeout and size validation.
+
+Provides utilities for safe state capture and restore operations
+that align with Orchestrator timeout expectations.
+"""
+
+import asyncio
+import functools
+import logging
+import time
+from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
+from dataclasses import dataclass
+from typing import Any, Callable, TypeVar
+
+from dory.utils.errors import DoryStateError
+
+logger = logging.getLogger(__name__)
+
+# Type variable for generic return type
+T = TypeVar("T")
+
+# Orchestrator constants (from transfer.go)
+ORCHESTRATOR_STATE_TIMEOUT_SEC = 30  # DefaultHTTPTimeout in transfer.go
+ORCHESTRATOR_MAX_STATE_SIZE = 10 * 1024 * 1024  # MaxResponseBodySize in transfer.go (10MB)
+
+# Default SDK limits (with safety margin)
+DEFAULT_CAPTURE_TIMEOUT_SEC = 25  # 5s buffer before Orchestrator timeout
+DEFAULT_RESTORE_TIMEOUT_SEC = 25
+DEFAULT_MAX_STATE_SIZE = 8 * 1024 * 1024  # 8MB, 2MB buffer before Orchestrator limit
+DEFAULT_SIZE_WARN_THRESHOLD = 0.75  # Warn at 75% of max
+
+
+@dataclass
+class TransferConfig:
+    """Configuration for state transfer operations."""
+
+    capture_timeout_sec: float = DEFAULT_CAPTURE_TIMEOUT_SEC
+    restore_timeout_sec: float = DEFAULT_RESTORE_TIMEOUT_SEC
+    max_size_bytes: int = DEFAULT_MAX_STATE_SIZE
+    size_warn_threshold: float = DEFAULT_SIZE_WARN_THRESHOLD
+
+    def __post_init__(self):
+        """Validate configuration against Orchestrator limits."""
+        if self.capture_timeout_sec >= ORCHESTRATOR_STATE_TIMEOUT_SEC:
+            logger.warning(
+                f"state_capture_timeout_sec ({self.capture_timeout_sec}s) >= "
+                f"Orchestrator timeout ({ORCHESTRATOR_STATE_TIMEOUT_SEC}s). "
+                "Reducing to {ORCHESTRATOR_STATE_TIMEOUT_SEC - 5}s."
+            )
+            self.capture_timeout_sec = ORCHESTRATOR_STATE_TIMEOUT_SEC - 5
+
+        if self.restore_timeout_sec >= ORCHESTRATOR_STATE_TIMEOUT_SEC:
+            logger.warning(
+                f"state_restore_timeout_sec ({self.restore_timeout_sec}s) >= "
+                f"Orchestrator timeout ({ORCHESTRATOR_STATE_TIMEOUT_SEC}s). "
+                f"Reducing to {ORCHESTRATOR_STATE_TIMEOUT_SEC - 5}s."
+            )
+            self.restore_timeout_sec = ORCHESTRATOR_STATE_TIMEOUT_SEC - 5
+
+        if self.max_size_bytes > ORCHESTRATOR_MAX_STATE_SIZE:
+            logger.warning(
+                f"state_max_size_bytes ({self.max_size_bytes}) > "
+                f"Orchestrator limit ({ORCHESTRATOR_MAX_STATE_SIZE}). "
+                f"Reducing to {ORCHESTRATOR_MAX_STATE_SIZE}."
+            )
+            self.max_size_bytes = ORCHESTRATOR_MAX_STATE_SIZE
+
+
+@dataclass
+class TransferMetrics:
+    """Metrics from a state transfer operation."""
+
+    duration_sec: float
+    size_bytes: int
+    size_ratio: float  # size / max_size
+    timed_out: bool
+    size_exceeded: bool
+
+
+class StateTransferError(DoryStateError):
+    """Error during state transfer operation."""
+
+    def __init__(
+        self,
+        message: str,
+        metrics: TransferMetrics | None = None,
+        cause: Exception | None = None,
+    ):
+        super().__init__(message, cause=cause)
+        self.metrics = metrics
+
+
+class StateTransferTimeout(StateTransferError):
+    """State transfer operation timed out."""
+    pass
+
+
+class StateSizeExceeded(StateTransferError):
+    """State size exceeds configured maximum."""
+    pass
+
+
+def validate_state_size(
+    state_json: str,
+    max_size: int = DEFAULT_MAX_STATE_SIZE,
+    warn_threshold: float = DEFAULT_SIZE_WARN_THRESHOLD,
+) -> TransferMetrics:
+    """
+    Validate state size against limits.
+
+    Args:
+        state_json: Serialized state JSON string
+        max_size: Maximum allowed size in bytes
+        warn_threshold: Fraction of max_size to trigger warning
+
+    Returns:
+        TransferMetrics with size information
+
+    Raises:
+        StateSizeExceeded: If state exceeds max_size
+    """
+    size_bytes = len(state_json.encode("utf-8"))
+    size_ratio = size_bytes / max_size if max_size > 0 else 0
+
+    metrics = TransferMetrics(
+        duration_sec=0,
+        size_bytes=size_bytes,
+        size_ratio=size_ratio,
+        timed_out=False,
+        size_exceeded=size_bytes > max_size,
+    )
+
+    if size_bytes > max_size:
+        raise StateSizeExceeded(
+            f"State size ({size_bytes:,} bytes) exceeds maximum "
+            f"({max_size:,} bytes). Orchestrator will reject this state.",
+            metrics=metrics,
+        )
+
+    if size_ratio >= warn_threshold:
+        logger.warning(
+            f"State size ({size_bytes:,} bytes) is {size_ratio:.1%} of maximum "
+            f"({max_size:,} bytes). Consider reducing state size to avoid "
+            "transfer failures."
+        )
+
+    return metrics
+
+
+def with_timeout(
+    timeout_sec: float,
+    operation_name: str = "operation",
+) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    """
+    Decorator to add timeout to synchronous functions.
+
+    Runs the function in a thread pool executor with timeout.
+
+    Args:
+        timeout_sec: Timeout in seconds
+        operation_name: Name for error messages
+
+    Returns:
+        Decorated function with timeout
+    """
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> T:
+            start_time = time.monotonic()
+
+            with ThreadPoolExecutor(max_workers=1) as executor:
+                future = executor.submit(func, *args, **kwargs)
+
+                try:
+                    result = future.result(timeout=timeout_sec)
+                    duration = time.monotonic() - start_time
+
+                    # Log if operation took significant time
+                    if duration > timeout_sec * 0.5:
+                        logger.warning(
+                            f"{operation_name} took {duration:.2f}s "
+                            f"({duration/timeout_sec:.1%} of {timeout_sec}s timeout)"
+                        )
+
+                    return result
+
+                except FuturesTimeoutError:
+                    duration = time.monotonic() - start_time
+                    metrics = TransferMetrics(
+                        duration_sec=duration,
+                        size_bytes=0,
+                        size_ratio=0,
+                        timed_out=True,
+                        size_exceeded=False,
+                    )
+                    raise StateTransferTimeout(
+                        f"{operation_name} timed out after {timeout_sec}s. "
+                        "Consider reducing state size or optimizing get_state().",
+                        metrics=metrics,
+                    )
+
+        return wrapper
+    return decorator
+
+
+async def async_with_timeout(
+    coro: Any,
+    timeout_sec: float,
+    operation_name: str = "operation",
+) -> Any:
+    """
+    Execute coroutine with timeout.
+
+    Args:
+        coro: Coroutine to execute
+        timeout_sec: Timeout in seconds
+        operation_name: Name for error messages
+
+    Returns:
+        Result of the coroutine
+
+    Raises:
+        StateTransferTimeout: If operation times out
+    """
+    start_time = time.monotonic()
+
+    try:
+        result = await asyncio.wait_for(coro, timeout=timeout_sec)
+        duration = time.monotonic() - start_time
+
+        # Log if operation took significant time
+        if duration > timeout_sec * 0.5:
+            logger.warning(
+                f"{operation_name} took {duration:.2f}s "
+                f"({duration/timeout_sec:.1%} of {timeout_sec}s timeout)"
+            )
+
+        return result
+
+    except asyncio.TimeoutError:
+        duration = time.monotonic() - start_time
+        metrics = TransferMetrics(
+            duration_sec=duration,
+            size_bytes=0,
+            size_ratio=0,
+            timed_out=True,
+            size_exceeded=False,
+        )
+        raise StateTransferTimeout(
+            f"{operation_name} timed out after {timeout_sec}s. "
+            "Consider reducing state size or optimizing the operation.",
+            metrics=metrics,
+        )
+
+
+class StateCaptureGuard:
+    """
+    Context manager for safe state capture with timeout and size validation.
+
+    Usage:
+        config = TransferConfig(capture_timeout_sec=25, max_size_bytes=8*1024*1024)
+
+        with StateCaptureGuard(config) as guard:
+            state = processor.get_state()
+            state_json = json.dumps(state)
+            guard.validate(state_json)
+    """
+
+    def __init__(self, config: TransferConfig | None = None):
+        """
+        Initialize capture guard.
+
+        Args:
+            config: Transfer configuration
+        """
+        self._config = config or TransferConfig()
+        self._start_time: float = 0
+        self._metrics: TransferMetrics | None = None
+
+    def __enter__(self) -> "StateCaptureGuard":
+        self._start_time = time.monotonic()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
+        duration = time.monotonic() - self._start_time
+
+        if duration > self._config.capture_timeout_sec:
+            logger.error(
+                f"State capture took {duration:.2f}s, exceeding "
+                f"{self._config.capture_timeout_sec}s timeout"
+            )
+
+        return False  # Don't suppress exceptions
+
+    def validate(self, state_json: str) -> TransferMetrics:
+        """
+        Validate captured state.
+
+        Args:
+            state_json: Serialized state JSON
+
+        Returns:
+            TransferMetrics with capture information
+
+        Raises:
+            StateSizeExceeded: If state exceeds max size
+            StateTransferTimeout: If capture exceeded timeout
+        """
+        duration = time.monotonic() - self._start_time
+
+        # Check timeout
+        if duration > self._config.capture_timeout_sec:
+            self._metrics = TransferMetrics(
+                duration_sec=duration,
+                size_bytes=len(state_json.encode("utf-8")),
+                size_ratio=0,
+                timed_out=True,
+                size_exceeded=False,
+            )
+            raise StateTransferTimeout(
+                f"State capture took {duration:.2f}s, exceeding "
+                f"{self._config.capture_timeout_sec}s timeout",
+                metrics=self._metrics,
+            )
+
+        # Check size
+        size_metrics = validate_state_size(
+            state_json,
+            max_size=self._config.max_size_bytes,
+            warn_threshold=self._config.size_warn_threshold,
+        )
+
+        self._metrics = TransferMetrics(
+            duration_sec=duration,
+            size_bytes=size_metrics.size_bytes,
+            size_ratio=size_metrics.size_ratio,
+            timed_out=False,
+            size_exceeded=size_metrics.size_exceeded,
+        )
+
+        return self._metrics
+
+    @property
+    def metrics(self) -> TransferMetrics | None:
+        """Get capture metrics."""
+        return self._metrics
+
+
+def log_transfer_summary(
+    operation: str,
+    metrics: TransferMetrics,
+    config: TransferConfig,
+) -> None:
+    """
+    Log a summary of the transfer operation.
+
+    Args:
+        operation: Operation name (e.g., "capture", "restore")
+        metrics: Transfer metrics
+        config: Transfer configuration
+    """
+    level = logging.INFO
+    status = "completed"
+
+    if metrics.timed_out:
+        level = logging.ERROR
+        status = "TIMED OUT"
+    elif metrics.size_exceeded:
+        level = logging.ERROR
+        status = "SIZE EXCEEDED"
+    elif metrics.size_ratio >= config.size_warn_threshold:
+        level = logging.WARNING
+        status = "completed (size warning)"
+
+    logger.log(
+        level,
+        f"State {operation} {status}: "
+        f"duration={metrics.duration_sec:.2f}s/{config.capture_timeout_sec}s, "
+        f"size={metrics.size_bytes:,}B/{config.max_size_bytes:,}B "
+        f"({metrics.size_ratio:.1%})",
+    )

dory/monitoring/__init__.py

@@ -0,0 +1,29 @@
+"""
+Dory Monitoring and Observability
+
+OpenTelemetry integration for distributed tracing and metrics.
+"""
+
+from dory.monitoring.opentelemetry import (
+    OpenTelemetryManager,
+    trace_function,
+    create_span,
+    add_span_attributes,
+    record_exception,
+    get_tracer,
+    initialize_otel,
+    get_global_otel,
+    OTEL_AVAILABLE,
+)
+
+__all__ = [
+    "OpenTelemetryManager",
+    "trace_function",
+    "create_span",
+    "add_span_attributes",
+    "record_exception",
+    "get_tracer",
+    "initialize_otel",
+    "get_global_otel",
+    "OTEL_AVAILABLE",
+]