feat: add plugin interface

Author: zhongkechen

src/aws_durable_execution_sdk_python/execution.py

@@ -96,6 +96,15 @@ class OperationSubType(Enum):
     CHAINED_INVOKE = "ChainedInvoke"
 
 
+class InvocationStatus(Enum):
+    SUCCEEDED = "SUCCEEDED"
+    FAILED = "FAILED"
+    PENDING = "PENDING"
+
+    # Used internally only: the invocation failed and the backend will retry
+    RETRY = "RETRY"
+
+
 @dataclass(frozen=True)
 class ExecutionDetails:
     input_payload: str | None = None

src/aws_durable_execution_sdk_python/lambda_service.py

@@ -0,0 +1,299 @@
+import datetime
+import logging
+from abc import ABC
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from aws_durable_execution_sdk_python.lambda_service import (
+    OperationType,
+    OperationStatus,
+    OperationAction,
+    OperationSubType,
+    ErrorObject,
+    InvocationStatus,
+    Operation,
+)
+from aws_durable_execution_sdk_python.types import LambdaContext
+
+if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python.execution import (
+        DurableExecutionInvocationOutput,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OperationStartInfo:
+    operation_id: str
+    operation_type: OperationType
+    sub_type: OperationSubType | None = None
+    name: str | None = None
+    parent_id: str | None = None
+    start_timestamp: datetime.datetime | None = None
+
+
+@dataclass
+class OperationEndInfo(OperationStartInfo):
+    status: OperationStatus = OperationStatus.SUCCEEDED
+    end_timestamp: datetime.datetime | None = None
+    attempt: int | None = None
+    error: ErrorObject | None = None
+
+
+@dataclass
+class AttemptStartInfo(OperationStartInfo):
+    attempt: int = 1
+
+
+@dataclass
+class AttemptEndInfo(AttemptStartInfo):
+    succeeded: bool | None = None
+    end_timestamp: datetime.datetime | None = None
+    error: ErrorObject | None = None
+    next_attempt_delay_seconds: int | None = None
+
+
+@dataclass
+class InvocationStartInfo:
+    request_id: str | None
+    execution_arn: str | None
+    start_timestamp: datetime.datetime | None
+
+
+@dataclass
+class InvocationEndInfo(InvocationStartInfo):
+    status: InvocationStatus = InvocationStatus.SUCCEEDED
+    end_timestamp: datetime.datetime | None = None
+    error: ErrorObject | None = None
+
+
+@dataclass
+class ExecutionStartInfo(InvocationStartInfo):
+    pass
+
+
+@dataclass
+class ExecutionEndInfo(ExecutionStartInfo):
+    status: InvocationStatus = InvocationStatus.SUCCEEDED
+    end_timestamp: datetime.datetime | None = None
+    error: ErrorObject | None = None
+
+
+class DurableExecutionPlugin(ABC):
+    """Base class for plugins. Override only the methods you need."""
+
+    def on_execution_start(self, info: ExecutionStartInfo) -> None:
+        pass
+
+    def on_execution_end(self, info: ExecutionEndInfo) -> None:
+        pass
+
+    def on_invocation_start(self, info: InvocationStartInfo) -> None:
+        pass
+
+    def on_invocation_end(self, info: InvocationEndInfo) -> None:
+        pass
+
+    def on_operation_start(self, info: OperationStartInfo) -> None:
+        pass
+
+    def on_operation_end(self, info: OperationEndInfo) -> None:
+        pass
+
+    def on_operation_attempt_start(self, info: AttemptStartInfo) -> None:
+        pass
+
+    def on_operation_attempt_end(self, info: AttemptEndInfo) -> None:
+        pass
+
+    # Todo: further discussions required to finalize the following interface
+    # def enrich_log_context(self, info: OperationStartInfo | None) -> Dict[str, Any] | None: pass
+
+
+class PluginExecutor:
+    def __init__(self, plugins: list[DurableExecutionPlugin] | None):
+        self.plugins = plugins or []
+
+    def execute_plugins(self, info):
+        for plugin in self.plugins:
+            try:
+                match info:
+                    case ExecutionEndInfo():
+                        plugin.on_execution_end(info)
+                    case InvocationEndInfo():
+                        plugin.on_invocation_end(info)
+                    case ExecutionStartInfo():
+                        plugin.on_execution_start(info)
+                    case InvocationStartInfo():
+                        plugin.on_invocation_start(info)
+                    case AttemptEndInfo():
+                        plugin.on_operation_attempt_end(info)
+                    case OperationEndInfo():
+                        plugin.on_operation_end(info)
+                    case AttemptStartInfo():
+                        plugin.on_operation_attempt_start(info)
+                    case OperationStartInfo():
+                        plugin.on_operation_start(info)
+                    case _:
+                        raise ValueError(f"Unknown info type: {type(info)}")
+            except Exception:
+                # log and ignore the exception
+                logger.exception(
+                    "Plugin %s exception ignored", plugin.__class__.__name__
+                )
+
+    def on_invocation_start(
+        self,
+        durable_execution_arn: str,
+        context: LambdaContext | None,
+        execution_operation: Operation | None,
+        is_replaying: bool,
+    ) -> None:
+        aws_request_id = context.aws_request_id if context else None
+        start_timestamp = (
+            execution_operation.start_timestamp if execution_operation else None
+        )
+
+        if not is_replaying:
+            self.execute_plugins(
+                ExecutionStartInfo(
+                    request_id=aws_request_id,
+                    execution_arn=durable_execution_arn,
+                    start_timestamp=start_timestamp,
+                )
+            )
+
+        self.execute_plugins(
+            InvocationStartInfo(
+                request_id=aws_request_id,
+                execution_arn=durable_execution_arn,
+                start_timestamp=start_timestamp,
+            )
+        )
+
+    def on_invocation_end(
+        self,
+        durable_execution_arn: str | None,
+        context: LambdaContext,
+        execution_operation: Operation | None,
+        output: "DurableExecutionInvocationOutput",
+    ) -> None:
+        start_timestamp = (
+            execution_operation.start_timestamp if execution_operation else None
+        )
+        # the actual end timestamp may be unknown because it's not checkpointed yet
+        end_timestamp: datetime.datetime = (
+            execution_operation.end_timestamp if execution_operation else None
+        ) or datetime.datetime.now()
+        request_id = context.aws_request_id if context else None
+
+        self.execute_plugins(
+            InvocationEndInfo(
+                request_id=request_id,
+                execution_arn=durable_execution_arn,
+                start_timestamp=start_timestamp,
+                status=output.status,
+                end_timestamp=end_timestamp,
+                error=output.error,
+            )
+        )
+
+        if output.status in [InvocationStatus.SUCCEEDED, InvocationStatus.FAILED]:
+            self.execute_plugins(
+                ExecutionEndInfo(
+                    request_id=request_id,
+                    execution_arn=durable_execution_arn,
+                    start_timestamp=start_timestamp,
+                    status=output.status,
+                    end_timestamp=end_timestamp,
+                    error=output.error,
+                )
+            )
+
+    def on_operation_action(self, update):
+        """Execute any registered plugins for a given operation before it is updated.
+
+        Args:
+            update: the operation update that is pending checkpoint
+        """
+        if update.action is OperationAction.START:
+            self.execute_plugins(
+                OperationStartInfo(
+                    operation_id=update.operation_id,
+                    operation_type=update.operation_type,
+                    sub_type=update.sub_type,
+                    name=update.name,
+                    parent_id=update.parent_id,
+                    start_timestamp=datetime.datetime.now(),
+                )
+            )
+
+    def on_operation_update(self, operation):
+        """Execute any registered plugins for a given operation after it is updated.
+
+        Updates such as STARTED might be omitted because START and completion action (e.g. SUCCEED/FAIL) may be
+        checkpointed in batch and the backend returns only the terminal status (e.g. SUCCEEDED/PENDING/FAILED).
+
+        Args:
+            operation: the operation is just checkpointed
+        """
+        params = dict(
+            operation_id=operation.operation_id,
+            operation_type=operation.operation_type,
+            sub_type=operation.sub_type,
+            name=operation.name,
+            parent_id=operation.parent_id,
+            start_timestamp=operation.start_timestamp,
+        )
+        # todo: Python SDK doesn't submit a START update when retrying
+        if operation.step_details and (
+            self._is_terminal_status(operation.status)
+            # PENDING in addition to terminal status
+            or operation.status is OperationStatus.PENDING
+        ):
+            self.execute_plugins(AttemptStartInfo(**params))
+            self.execute_plugins(
+                AttemptEndInfo(
+                    **params,
+                    end_timestamp=operation.end_timestamp,
+                    attempt=operation.step_details.attempt,
+                    succeeded=operation.status is OperationStatus.SUCCEEDED,
+                    error=operation.step_details.error,
+                )
+            )
+
+        if self._is_terminal_status(operation.status):
+            attempt = operation.step_details.attempt if operation.step_details else None
+            self.execute_plugins(OperationStartInfo(**params))
+            self.execute_plugins(
+                OperationEndInfo(
+                    **params,
+                    end_timestamp=operation.end_timestamp,
+                    status=operation.status,
+                    error=self._extract_error(operation),
+                    attempt=attempt,
+                )
+            )
+
+    @staticmethod
+    def _extract_error(operation: Operation):
+        if operation.step_details and operation.step_details.error:
+            return operation.step_details.error
+        if operation.callback_details and operation.callback_details.error:
+            return operation.callback_details.error
+        if operation.chained_invoke_details and operation.chained_invoke_details.error:
+            return operation.chained_invoke_details.error
+        if operation.context_details and operation.context_details.error:
+            return operation.context_details.error
+        return None
+
+    @staticmethod
+    def _is_terminal_status(status):
+        return status in [
+            OperationStatus.SUCCEEDED,
+            OperationStatus.FAILED,
+            OperationStatus.TIMED_OUT,
+            OperationStatus.CANCELLED,
+            OperationStatus.STOPPED,
+        ]

src/aws_durable_execution_sdk_python/plugin.py

@@ -7,6 +7,8 @@
 import queue
 import threading
 import time
+from concurrent.futures import Executor
+from contextlib import contextmanager
 from dataclasses import dataclass
 from enum import Enum
 from threading import Lock
@@ -30,6 +32,9 @@
     OperationUpdate,
     StateOutput,
 )
+from aws_durable_execution_sdk_python.plugin import (
+    PluginExecutor,
+)
 from aws_durable_execution_sdk_python.threading import CompletionEvent, OrderedLock
 
 if TYPE_CHECKING:
@@ -229,13 +234,15 @@ def __init__(
         initial_checkpoint_token: str,
         operations: MutableMapping[str, Operation],
         service_client: DurableServiceClient,
+        plugin_executor: PluginExecutor | None = None,
         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._plugin_executor: PluginExecutor = plugin_executor or PluginExecutor(None)
         self._ordered_checkpoint_lock: OrderedLock = OrderedLock()
         self._operations_lock: Lock = Lock()
 
@@ -267,7 +274,7 @@ def fetch_paginated_operations(
         initial_operations: list[Operation],
         checkpoint_token: str,
         next_marker: str | None,
-    ) -> None:
+    ) -> list[Operation]:
         """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.
@@ -276,6 +283,8 @@ def fetch_paginated_operations(
             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.
+        Returns:
+            List of all operations fetched from the Durable Functions API
 
         Raises:
             GetExecutionStateError: If the API call fails. The error is logged
@@ -308,6 +317,7 @@ def fetch_paginated_operations(
                     self.operations.update(
                         {op.operation_id: op for op in all_operations}
                     )
+        return all_operations
 
     def get_input_payload(self) -> str | None:
         # It is possible that backend will not provide an execution operation
@@ -682,12 +692,18 @@ def checkpoint_batches_forever(self) -> None:
                     current_checkpoint_token = output.checkpoint_token
 
                     # Fetch new operations from the API before unblocking sync waiters
-                    self.fetch_paginated_operations(
+                    updated_operations = self.fetch_paginated_operations(
                         output.new_execution_state.operations,
                         output.checkpoint_token,
                         output.new_execution_state.next_marker,
                     )
 
+                    for update in updates:
+                        self._plugin_executor.on_operation_action(update)
+
+                    for operation in updated_operations:
+                        self._plugin_executor.on_operation_update(operation)
+
                     # Signal completion for any synchronous operations
                     for queued_op in batch:
                         if queued_op.completion_event is not None: