state.py

"""Model for execution state."""

from __future__ import annotations

import json
import logging
import queue
import threading
import time
from dataclasses import dataclass
from enum import Enum
from threading import Lock
from typing import TYPE_CHECKING

from aws_durable_execution_sdk_python.exceptions import (
    BackgroundThreadError,
    CallableRuntimeError,
    DurableExecutionsError,
    OrphanedChildException,
)
from aws_durable_execution_sdk_python.lambda_service import (
    CheckpointOutput,
    DurableServiceClient,
    ErrorObject,
    Operation,
    OperationAction,
    OperationStatus,
    OperationType,
    OperationUpdate,
    StateOutput,
)
from aws_durable_execution_sdk_python.threading import CompletionEvent, OrderedLock

if TYPE_CHECKING:
    import datetime
    from collections.abc import MutableMapping

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class CheckpointBatcherConfig:
    """Configuration for checkpoint batching behavior.

    Attributes:
        max_batch_size_bytes: Maximum batch size in bytes (default: 750KB)
        max_batch_time_seconds: Maximum time to wait before flushing batch (default: 1.0 second)
        max_batch_operations: Maximum number of operations per batch (default: 250)
    """

    max_batch_size_bytes: int = 750 * 1024  # 750KB
    max_batch_time_seconds: float = 1.0
    max_batch_operations: int = 250


@dataclass(frozen=True)
class QueuedOperation:
    """Wrapper for operations in the checkpoint queue.

    Attributes:
        operation_update: The operation update to be checkpointed, or None for empty checkpoints
        completion_event: CompletionEvent for synchronous operations, or None for async operations
    """

    operation_update: OperationUpdate | None
    completion_event: CompletionEvent | None = None


@dataclass(frozen=True)
class CheckpointedResult:
    """Result of a checkpointed operation.

    Set by ExecutionState.get_checkpoint_result. This is a convenience wrapper around
    Operation.

    Attributes:
        operation (Operation): The wrapped operation for the checkpoint result.
        status (OperationStatus): The status of the operation.
        result (str): the result of the operation.
        error (ErrorObject): the error of the operation.
    """

    operation: Operation | None = None
    status: OperationStatus | None = None
    result: str | None = None
    error: ErrorObject | None = None

    @classmethod
    def create_from_operation(cls, operation: Operation) -> CheckpointedResult:
        """Create a result from an operation."""
        result: str | None = None
        error: ErrorObject | None = None
        match operation.operation_type:
            case OperationType.STEP:
                step_details = operation.step_details
                result = step_details.result if step_details else None
                error = step_details.error if step_details else None

            case OperationType.CALLBACK:
                callback_details = operation.callback_details
                result = callback_details.result if callback_details else None
                error = callback_details.error if callback_details else None

            case OperationType.CHAINED_INVOKE:
                invoke_details = operation.chained_invoke_details
                result = invoke_details.result if invoke_details else None
                error = invoke_details.error if invoke_details else None

            case OperationType.CONTEXT:
                context_details = operation.context_details
                result = context_details.result if context_details else None
                error = context_details.error if context_details else None

        return cls(
            operation=operation, status=operation.status, result=result, error=error
        )

    @classmethod
    def create_not_found(cls) -> CheckpointedResult:
        """Create a result when the checkpoint was not found."""
        return cls(operation=None)

    def is_existent(self) -> bool:
        """Return true if a checkpoint of any type exists."""
        return self.operation is not None

    def is_succeeded(self) -> bool:
        """Return True if the checkpointed operation is SUCCEEDED."""
        op = self.operation
        if not op:
            return False

        return op.status is OperationStatus.SUCCEEDED

    def is_cancelled(self) -> bool:
        if op := self.operation:
            return op.status is OperationStatus.CANCELLED
        return False

    def is_failed(self) -> bool:
        """Return True if the checkpointed operation is FAILED."""
        op = self.operation
        if not op:
            return False

        return op.status is OperationStatus.FAILED

    def is_stopped(self) -> bool:
        """Return True if the checkpointed operation is STOPPED"""
        op = self.operation
        if not op:
            return False

        return op.status is OperationStatus.STOPPED

    def is_started(self) -> bool:
        """Return True if the checkpointed operation is STARTED."""
        op = self.operation
        if not op:
            return False
        return op.status is OperationStatus.STARTED

    def is_started_or_ready(self) -> bool:
        """Return True if the checkpointed operation is STARTED or READY."""
        op = self.operation
        if not op:
            return False
        return op.status in {OperationStatus.STARTED, OperationStatus.READY}

    def is_pending(self) -> bool:
        """Return True if the checkpointed operation is PENDING."""
        op = self.operation
        if not op:
            return False
        return op.status is OperationStatus.PENDING

    def is_timed_out(self) -> bool:
        """Return True if the checkpointed operation is TIMED_OUT."""
        op = self.operation
        if not op:
            return False
        return op.status is OperationStatus.TIMED_OUT

    def is_replay_children(self) -> bool:
        op = self.operation
        if not op:
            return False
        return op.context_details.replay_children if op.context_details else False

    def raise_callable_error(self, msg: str | None = None) -> None:
        if self.error is None:
            err_msg = (
                msg
                or "Unknown error. No ErrorObject exists on the Checkpoint Operation."
            )
            raise CallableRuntimeError(
                message=err_msg,
                error_type=None,
                data=None,
                stack_trace=None,
            )

        raise self.error.to_callable_runtime_error()

    def get_next_attempt_timestamp(self) -> datetime.datetime | None:
        if self.operation and self.operation.step_details:
            return self.operation.step_details.next_attempt_timestamp
        return None


# shared so don't need to create an instance for each not found check
CHECKPOINT_NOT_FOUND = CheckpointedResult.create_not_found()


class ReplayStatus(Enum):
    """Status indicating whether execution is replaying or executing new operations."""

    REPLAY = "replay"
    NEW = "new"


class ExecutionState:
    """Get, set and maintain execution state. This is mutable. Create and check checkpoints."""

    def __init__(
        self,
        durable_execution_arn: str,
        initial_checkpoint_token: str,
        operations: MutableMapping[str, Operation],
        service_client: DurableServiceClient,
        batcher_config: CheckpointBatcherConfig | None = None,
        replay_status: ReplayStatus = ReplayStatus.NEW,
    ):
        self.durable_execution_arn: str = durable_execution_arn
        self._current_checkpoint_token: str = initial_checkpoint_token
        self.operations: MutableMapping[str, Operation] = operations
        self._service_client: DurableServiceClient = service_client
        self._ordered_checkpoint_lock: OrderedLock = OrderedLock()
        self._operations_lock: Lock = Lock()

        # Checkpoint batching configuration
        self._batcher_config: CheckpointBatcherConfig = (
            batcher_config or CheckpointBatcherConfig()
        )

        # Checkpoint batching components
        self._checkpoint_queue: queue.Queue[QueuedOperation] = queue.Queue()
        self._overflow_queue: queue.Queue[QueuedOperation] = queue.Queue()
        self._checkpointing_stopped: threading.Event = threading.Event()
        self._checkpointing_failed: CompletionEvent = CompletionEvent()

        # Concurrency management for parallel operations: parent_id -> {child_operation_ids}
        self._parent_to_children: dict[str, set[str]] = {}

        # Operations whose parent has completed
        self._parent_done: set[str] = set()

        # Protects parent_to_children and parent_done
        self._parent_done_lock: Lock = Lock()
        self._replay_status: ReplayStatus = replay_status
        self._replay_status_lock: Lock = Lock()
        self._visited_operations: set[str] = set()

    def fetch_paginated_operations(
        self,
        initial_operations: list[Operation],
        checkpoint_token: str,
        next_marker: str | None,
    ) -> None:
        """Add initial operations and fetch all paginated operations from the Durable Functions API. This method is thread_safe.

        The checkpoint_token is passed explicitly as a parameter rather than using the instance variable to ensure thread safety.

        Args:
            initial_operations: initial operations to be added to ExecutionState
            checkpoint_token: checkpoint token used to call Durable Functions API.
            next_marker: a marker indicates that there are paginated operations.
        """
        all_operations: list[Operation] = (
            initial_operations.copy() if initial_operations else []
        )
        while next_marker:
            output: StateOutput = self._service_client.get_execution_state(
                durable_execution_arn=self.durable_execution_arn,
                checkpoint_token=checkpoint_token,
                next_marker=next_marker,
            )
            all_operations.extend(output.operations)
            next_marker = output.next_marker
        with self._operations_lock:
            self.operations.update({op.operation_id: op for op in all_operations})

    def track_replay(self, operation_id: str) -> None:
        """Check if operation exists with completed status; if not, transition to NEW status.

        This method is called before each operation (step, wait, invoke, etc.) to determine
        if we've reached the replay boundary. Once we encounter an operation that doesn't
        exist or isn't completed, we transition from REPLAY to NEW status, which enables
        logging for all subsequent code.

        Args:
            operation_id: The operation ID to check
        """
        with self._replay_status_lock:
            if self._replay_status == ReplayStatus.REPLAY:
                self._visited_operations.add(operation_id)
                completed_ops = {
                    op_id
                    for op_id, op in self.operations.items()
                    if op.operation_type != OperationType.EXECUTION
                    and op.status
                    in {
                        OperationStatus.SUCCEEDED,
                        OperationStatus.FAILED,
                        OperationStatus.CANCELLED,
                        OperationStatus.STOPPED,
                        OperationStatus.TIMED_OUT,
                    }
                }
                if completed_ops.issubset(self._visited_operations):
                    logger.debug(
                        "Transitioning from REPLAY to NEW status at operation %s",
                        operation_id,
                    )
                    self._replay_status = ReplayStatus.NEW

    def is_replaying(self) -> bool:
        """Check if execution is currently in replay mode.

        Returns:
            True if in REPLAY status, False if in NEW status
        """
        with self._replay_status_lock:
            return self._replay_status is ReplayStatus.REPLAY

    def get_checkpoint_result(self, checkpoint_id: str) -> CheckpointedResult:
        """Get checkpoint result.

        Note this does not invoke the Durable Functions API. It only checks
        against the checkpoints currently saved in ExecutionState. The current
        saved checkpoints are from InitialExecutionState as retrieved
        at the start of the current execution/replay (see execution.durable_execution),
        and from each create_checkpoint response.

        Args:
            checkpoint_id: str - id for checkpoint to retrieve.

        Returns:
            CheckpointedResult with is_succeeded True if the checkpoint exists and its
                status is SUCCEEDED. If the checkpoint exists but its status is not
                SUCCEEDED, or if the checkpoint doesn't exist, then return
                CheckpointedResult with is_succeeded=False,result=None.
        """
        # checking status are deliberately under a lighter non-serialized lock
        with self._operations_lock:
            if checkpoint := self.operations.get(checkpoint_id):
                return CheckpointedResult.create_from_operation(checkpoint)

        return CHECKPOINT_NOT_FOUND

    def create_checkpoint(
        self,
        operation_update: OperationUpdate | None = None,
        is_sync: bool = True,  # noqa: FBT001, FBT002
    ) -> None:
        """Create a checkpoint with optional synchronous behavior.

        This method enqueues a checkpoint operation for processing by the background
        batching thread. By default, the operation is synchronous (blocking) to ensure
        the checkpoint is persisted before continuing. For performance-critical paths
        where immediate confirmation is not required, set is_sync=False.

        Synchronous checkpoints (is_sync=True, default):
        - Block the caller until the checkpoint is processed by the background thread
        - Ensure the checkpoint is persisted before continuing
        - Safe default for correctness
        - Use cases: Most operations requiring confirmation before proceeding

        Asynchronous checkpoints (is_sync=False, opt-in):
        - Return immediately without waiting for the checkpoint to complete
        - Performance optimization for specific use cases
        - Use cases: observability checkpoints, fire-and-forget operations

        When to use synchronous checkpoints (is_sync=True, default):
        1. Step START with AtMostOncePerRetry semantics - prevents duplicate execution
        2. Operation completion (SUCCEED/FAIL) - ensures state persisted before returning
        3. Retry operations - ensures retry state recorded before continuing
        4. Callback START - must wait for API to generate callback ID
        5. Invoke START - ensures chained invoke recorded before proceeding
        6. Child context results - ensures results persisted before returning
        7. Large results - ensures results saved before returning to caller
        8. Wait for condition completion - ensures state recorded before proceeding
        9. Most operations - safe default

        When to use asynchronous checkpoints (is_sync=False, opt-in):
        1. Step START with AtLeastOncePerRetry semantics - performance optimization
        2. Child context START - fire-and-forget for performance
        3. Wait for condition START - observability only, no blocking needed
        4. Any checkpoint where immediate confirmation is not required AND performance matters

        Args:
            operation_update: The checkpoint to create. If None, creates an empty
                            checkpoint to get a fresh checkpoint token and updated
                            operations list.
            is_sync: If True (default), blocks until the checkpoint is processed.
                    If False, returns immediately without blocking for performance.

        Raises:
            Any exception from the background checkpoint processing will propagate
            through the ThreadPoolExecutor to the main thread, terminating the Lambda.

        Examples:
            # Synchronous checkpoint (default, safe)
            execution_state.create_checkpoint(operation_update)

            # Explicit synchronous checkpoint
            execution_state.create_checkpoint(operation_update, is_sync=True)

            # Asynchronous checkpoint (opt-in for performance)
            execution_state.create_checkpoint(operation_update, is_sync=False)

            # Empty checkpoint (sync by default)
            execution_state.create_checkpoint()

            # Empty checkpoint (async for performance)
            execution_state.create_checkpoint(is_sync=False)
        """
        # if this is CONTEXT complete, mark incomplete descendants as orphans so the children can't complete after the parent
        if operation_update is not None:
            # Use single lock to coordinate completion and checkpoint validation
            with self._parent_done_lock:
                # Build parent-to-children map as operations are created
                if operation_update.parent_id:
                    if operation_update.parent_id not in self._parent_to_children:
                        self._parent_to_children[operation_update.parent_id] = set()
                    self._parent_to_children[operation_update.parent_id].add(
                        operation_update.operation_id
                    )

                # Handle CONTEXT completion - mark descendants while holding lock
                if (
                    operation_update.operation_type == OperationType.CONTEXT
                    and operation_update.action
                    in {OperationAction.SUCCEED, OperationAction.FAIL}
                ):
                    self._mark_orphans(operation_update.operation_id)

                # Check if this operation's parent is done
                if operation_update.operation_id in self._parent_done:
                    logger.debug(
                        "Rejecting checkpoint for operation %s - parent is done",
                        operation_update.operation_id,
                    )
                    error_msg = (
                        "Parent context completed, child operation cannot checkpoint"
                    )
                    raise OrphanedChildException(
                        error_msg,
                        operation_id=operation_update.operation_id,
                    )

        # Check if background checkpointing has failed
        if self._checkpointing_failed.is_set():
            # This will raise the stored BackgroundThreadError
            self._checkpointing_failed.wait()

        # Conditionally create completion event based on is_sync parameter
        completion_event: CompletionEvent | None = (
            CompletionEvent() if is_sync else None
        )

        # Create wrapper object for queue
        queued_op = QueuedOperation(operation_update, completion_event)

        # Enqueue the wrapper object (operation_update can be None for empty checkpoints)
        self._checkpoint_queue.put(queued_op)

        # Conditionally wait for completion based on is_sync parameter
        if is_sync:
            logger.debug("Enqueued checkpoint operation for synchronous processing")
            if completion_event is None:  # pragma: no cover
                # this shouldn't ever be possible
                msg: str = "completion_event must be set for synchronous execution"
                raise DurableExecutionsError(msg)

            # Wait for completion - will raise BackgroundThreadError if background thread fails
            completion_event.wait()
        else:
            logger.debug("Enqueued checkpoint operation for asynchronous processing")

    def create_checkpoint_sync(
        self,
        operation_update: OperationUpdate | None = None,
    ) -> None:
        """Create a synchronous checkpoint that raises original errors instead of BackgroundThreadError.

        This method is identical to create_checkpoint(is_sync=True) except that if the background
        checkpoint processing fails, it raises the original exception directly instead of wrapping
        it in a BackgroundThreadError.

        This is useful in execution contexts where you want the original checkpoint error to
        propagate (e.g., CheckpointError, RuntimeError) rather than the wrapped BackgroundThreadError.
        The method always blocks until the checkpoint is processed.

        Args:
            operation_update: The checkpoint to create. If None, creates an empty checkpoint.

        Raises:
            The original exception from the background checkpoint processing if it fails,
            unwrapped from BackgroundThreadError (e.g., CheckpointError, RuntimeError).

        Example:
            # Instead of getting BackgroundThreadError wrapping a CheckpointError:
            execution_state.create_checkpoint_sync(operation_update)
            # Raises CheckpointError directly
        """
        try:
            self.create_checkpoint(operation_update, is_sync=True)
        except BackgroundThreadError as bg_error:
            # Background checkpoint system failed - unwrap the original error
            logger.exception("Checkpoint processing failed - unwrapping original error")
            self.stop_checkpointing()
            # Raise the original exception unwrapped
            raise bg_error.source_exception from bg_error

    def _mark_orphans(self, context_id: str) -> None:
        """Mark all descendants (direct and transitive) as orphaned.

        This method uses BFS (Breadth-First Search) to recursively collect all
        descendants of the given context operation and marks them as orphaned.
        Once marked, these operations will be rejected if they attempt to checkpoint.

        Must be called while holding _parent_done_lock.

        Args:
            context_id: The operation ID of the CONTEXT that has completed
        """
        # Collect all descendants recursively using BFS
        all_descendants = set()
        # Start with root
        to_process: set[str] = {context_id}

        while to_process:
            current_id = to_process.pop()

            # Skip if already processed (avoid cycles, though shouldn't happen)
            if current_id in all_descendants:
                continue

            all_descendants.add(current_id)

            # Add all direct children to processing queue
            direct_children = self._parent_to_children.get(current_id, set())
            to_process.update(direct_children)

        # Remove the root itself (we only want descendants)
        all_descendants.discard(context_id)

        # Mark all descendants as orphaned
        self._parent_done.update(all_descendants)
        logger.debug(
            "Marked %d descendants as parent-done for context %s",
            len(all_descendants),
            context_id,
        )

    def checkpoint_batches_forever(self) -> None:
        """Single background thread that batches operations and processes results.

        Runs until shutdown is signaled. This method processes checkpoint operations
        in batches, makes API calls to persist them, and updates the execution state
        with the results.

        The method maintains the checkpoint token locally and updates it after each
        successful batch processing. It continues running until stop_checkpointing()
        is called.

        Note: When shutdown is signaled, only non-essential async checkpoints may remain
        in the queue. All critical synchronous checkpoints (SUCCEED, FAIL, etc.) will
        have already completed because the main thread blocks on them. Therefore, we
        don't need to drain the queue - the Lambda timeout will handle cleanup.

        Raises:
            Any exception from the service client checkpoint call will propagate naturally,
            terminating the background thread and signaling an error to the main thread.
        """
        # Keep checkpoint token as local variable in the loop
        current_checkpoint_token: str = self._current_checkpoint_token

        while not self._checkpointing_stopped.is_set():
            # Collect operations into a batch
            batch: list[QueuedOperation] = self._collect_checkpoint_batch()

            if batch:
                # Extract OperationUpdates from QueuedOperations for API call
                updates: list[OperationUpdate] = [
                    q.operation_update for q in batch if q.operation_update is not None
                ]

                logger.debug(
                    "Processing checkpoint batch with %d operations (%d non-empty)",
                    len(batch),
                    len(updates),
                )

                try:
                    # Make API call with batched operations
                    output: CheckpointOutput = self._service_client.checkpoint(
                        durable_execution_arn=self.durable_execution_arn,
                        checkpoint_token=current_checkpoint_token,
                        updates=updates,
                        client_token=None,
                    )

                    logger.debug("Checkpoint batch processed successfully")

                    # Update local token for next iteration
                    current_checkpoint_token = output.checkpoint_token

                    # Fetch new operations from the API before unblocking sync waiters
                    self.fetch_paginated_operations(
                        output.new_execution_state.operations,
                        output.checkpoint_token,
                        output.new_execution_state.next_marker,
                    )

                    # Signal completion for any synchronous operations
                    for queued_op in batch:
                        if queued_op.completion_event is not None:
                            queued_op.completion_event.set()
                except Exception as e:
                    # Checkpoint failed - wake all blocked threads so they can raise error
                    # Drain both queues and signal all completion events
                    logger.exception("Checkpoint batch processing failed")
                    bg_error: BackgroundThreadError = BackgroundThreadError(
                        "Checkpoint creation failed", e
                    )

                    # FIFO: although at this point order not really import any anymore
                    # Signal completion events for the failed batch
                    for queued_op in batch:
                        if queued_op.completion_event is not None:
                            queued_op.completion_event.set(bg_error)

                    # overflow 1st: although at this point order not really import any anymore
                    while not self._overflow_queue.empty():
                        try:
                            item = self._overflow_queue.get_nowait()
                            if item.completion_event:
                                item.completion_event.set(bg_error)
                        except queue.Empty:
                            break

                    # finally Wake all blocked threads in main queue
                    while not self._checkpoint_queue.empty():
                        try:
                            item = self._checkpoint_queue.get_nowait()
                            if item.completion_event:
                                item.completion_event.set(bg_error)
                        except queue.Empty:
                            break

                    # Set the failure event so future checkpoint attempts fail immediately
                    self._checkpointing_failed.set(bg_error)

                    # Exit the loop - error has been signaled to main thread via completion events
                    break

        logger.debug("Background checkpoint processing stopped")

    def stop_checkpointing(self) -> None:
        """Signal background thread to stop checkpointing.

        This method sets the checkpointing stopped event, which signals the background
        thread to exit. Any remaining async checkpoints in the queue are non-essential
        (observability only) and will be abandoned. All critical synchronous checkpoints
        will have already completed before this is called.
        """
        logger.debug("Signaling background thread to stop checkpointing")
        self._checkpointing_stopped.set()

    def _collect_checkpoint_batch(self) -> list[QueuedOperation]:
        """Collect multiple checkpoint operations into a batch for API efficiency.

        Processes overflow queue first to maintain FIFO order, then collects from main queue.
        Respects configured size, time, and operation count limits. Blocks for the first
        operation if queues are empty, then collects additional operations within the time
        window.

        Returns:
            List of QueuedOperation objects ready for batch processing. Returns empty list
            if no operations are available.
        """
        batch: list[QueuedOperation] = []
        total_size = 0

        # First, drain overflow queue (FIFO order preserved)
        try:
            while len(batch) < self._batcher_config.max_batch_operations:
                overflow_op = self._overflow_queue.get_nowait()
                op_size = self._calculate_operation_size(overflow_op)

                if total_size + op_size > self._batcher_config.max_batch_size_bytes:
                    # Put back and stop
                    self._overflow_queue.put(overflow_op)
                    break

                batch.append(overflow_op)
                total_size += op_size
        except queue.Empty:
            pass

        # If batch is empty, get first operation from main queue
        if not batch:
            # Block for first operation, checking stop signal periodically
            while not self._checkpointing_stopped.is_set():
                try:
                    first_op = self._checkpoint_queue.get(
                        timeout=0.1
                    )  # Check stop signal every 100ms
                    self._checkpoint_queue.task_done()
                    batch.append(first_op)
                    total_size += self._calculate_operation_size(first_op)
                    break
                except queue.Empty:
                    continue

            # If stopped and no operation retrieved, return empty batch
            if not batch:
                return batch

        # Start batching window using configured time
        batch_deadline = time.time() + self._batcher_config.max_batch_time_seconds

        # Collect additional operations within the time window
        while (
            time.time() < batch_deadline
            and len(batch) < self._batcher_config.max_batch_operations
            and not self._checkpointing_stopped.is_set()
        ):
            remaining_time = min(
                batch_deadline - time.time(),
                0.1,  # Check stop signal every 100ms
            )

            if remaining_time <= 0:
                break

            try:
                additional_op = self._checkpoint_queue.get(timeout=remaining_time)
                self._checkpoint_queue.task_done()
                op_size = self._calculate_operation_size(additional_op)

                # Check if adding this operation would exceed size limit
                if total_size + op_size > self._batcher_config.max_batch_size_bytes:
                    # Put in overflow queue for next batch
                    self._overflow_queue.put(additional_op)
                    logger.debug(
                        "Batch size limit reached, moving operation to overflow queue"
                    )
                    break

                batch.append(additional_op)
                total_size += op_size

            except queue.Empty:
                break

        logger.debug(
            "Collected batch of %d operations, total size: %d bytes",
            len(batch),
            total_size,
        )
        return batch

    @staticmethod
    def _calculate_operation_size(queued_op: QueuedOperation) -> int:
        """Calculate the serialized size of a queued operation for batching limits.

        Uses JSON serialization to estimate the size of the operation update. Empty
        checkpoints (None operation_update) have zero size.

        Args:
            queued_op: The queued operation to calculate size for

        Returns:
            Size in bytes of the serialized operation, or 0 for empty checkpoints
        """
        # Empty checkpoints have no size
        if queued_op.operation_update is None:
            return 0

        # Use JSON serialization to estimate size
        serialized = json.dumps(queued_op.operation_update.to_dict()).encode("utf-8")
        return len(serialized)

    def close(self):
        self.stop_checkpointing()