Merge branch 'main' into fix/remove-invoke-config-timeout

Author: AdityaAudi

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py

@@ -5,11 +5,11 @@
 import hashlib
 import os
 import re
-from datetime import datetime, UTC
+from datetime import UTC, datetime
+
+from opentelemetry.sdk.trace import IdGenerator, RandomIdGenerator
 
-from opentelemetry.sdk.trace import RandomIdGenerator
 
-HASH_LENGTH = 16
 HASHED_ID_PATTERN = re.compile(r"^[0-9a-f]{16}$")
 
 
@@ -67,19 +67,25 @@ def operation_id_to_span_id(operation_id: str) -> int:
 
 class DeterministicIdGenerator(RandomIdGenerator):
     """An ID generator that produces deterministic span IDs when a pending
-    operation ID is set, and random IDs otherwise.
+    operation ID is set, and falls back to the provided generator otherwise.
 
     Trace IDs are deterministic when an execution ARN is set, ensuring all
-    invocations of the same durable execution share a single trace.
+    invocations of the same durable execution share a single trace. When no
+    deterministic ID is available, generation is delegated to the fallback
+    generator (the tracer provider's original ID generator by default).
 
     Trace IDs embed a real timestamp so they satisfy the X-Ray format
     requirement (first 8 hex chars = Unix epoch seconds).
+
+    Args:
+        fallback_id_generator: Generator used when no deterministic ID is
+            available. Defaults to a new ``RandomIdGenerator``.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, fallback_id_generator: IdGenerator | None = None) -> None:
         self._next_span_id: int | None = None
         self._execution_trace_id: int | None = None
-        self._random_id_generator = RandomIdGenerator()
+        self._fallback_id_generator = fallback_id_generator or RandomIdGenerator()
 
     def set_next_span_id(self, span_id: int | None) -> None:
         """Set the operation ID to use for the next span's ID.
@@ -101,9 +107,11 @@ def set_trace_id(
 
     def generate_trace_id(self) -> int:
         """Generate a 128-bit trace ID."""
-        return self._execution_trace_id or self._random_id_generator.generate_trace_id()
+        return (
+            self._execution_trace_id or self._fallback_id_generator.generate_trace_id()
+        )
 
     def generate_span_id(self) -> int:
         """Generate a 64-bit span ID."""
         span_id, self._next_span_id = self._next_span_id, None
-        return span_id or self._random_id_generator.generate_span_id()
+        return span_id or self._fallback_id_generator.generate_span_id()

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin.py

@@ -7,30 +7,31 @@
 import threading
 from typing import TYPE_CHECKING, Any
 
-from opentelemetry import trace, context
-from opentelemetry.context import Context
-from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider
-from opentelemetry.sdk.trace.sampling import TraceIdRatioBased
-from opentelemetry.trace import (
-    Tracer,
-    StatusCode,
-    SpanContext,
-    Span,
-    Link,
-    TraceFlags,
-)
-
 from aws_durable_execution_sdk_python.lambda_service import OperationType
 from aws_durable_execution_sdk_python.plugin import (
     DurableInstrumentationPlugin,
     InvocationEndInfo,
     InvocationStartInfo,
     OperationEndInfo,
     OperationStartInfo,
-    UserFunctionStartInfo,
     UserFunctionEndInfo,
     UserFunctionOutcome,
+    UserFunctionStartInfo,
 )
+from opentelemetry import context, trace
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider
+from opentelemetry.sdk.trace.sampling import TraceIdRatioBased
+from opentelemetry.trace import (
+    Link,
+    Span,
+    SpanContext,
+    StatusCode,
+    TraceFlags,
+    Tracer,
+)
+
 from aws_durable_execution_sdk_python_otel.context_extractors import (
     ContextExtractor,
     xray_context_extractor,
@@ -40,6 +41,7 @@
     operation_id_to_span_id,
 )
 
+
 if TYPE_CHECKING:
     pass
 
@@ -95,9 +97,13 @@ def __init__(
             context_extractor or xray_context_extractor
         )
 
-        self._id_generator: DeterministicIdGenerator = DeterministicIdGenerator()
-
         self._provider = trace_provider
+        # A ProxyTracerProvider (the API default from trace.get_tracer_provider()
+        # before an SDK provider is configured) has no id_generator; fall back to
+        # None so DeterministicIdGenerator uses its own default generator.
+        self._id_generator: DeterministicIdGenerator = DeterministicIdGenerator(
+            fallback_id_generator=getattr(self._provider, "id_generator", None)
+        )
         self._provider.id_generator = self._id_generator
         self._provider.sampler = TraceIdRatioBased(sampling_rate)
         self._tracer: Tracer = self._provider.get_tracer(instrument_name)

packages/aws-durable-execution-sdk-python-otel/tests/test_deterministic_id_generator.py

@@ -4,6 +4,8 @@
 
 from datetime import UTC, datetime
 
+from opentelemetry.sdk.trace import IdGenerator, RandomIdGenerator
+
 from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
     HASHED_ID_PATTERN,
     DeterministicIdGenerator,
@@ -14,6 +16,20 @@
 )
 
 
+class _StubIdGenerator(IdGenerator):
+    """An IdGenerator that returns fixed, identifiable IDs."""
+
+    def __init__(self, trace_id: int, span_id: int) -> None:
+        self._trace_id = trace_id
+        self._span_id = span_id
+
+    def generate_trace_id(self) -> int:
+        return self._trace_id
+
+    def generate_span_id(self) -> int:
+        return self._span_id
+
+
 def test_parse_xray_root_trace_id_returns_root_from_header():
     """Verify X-Ray Root trace ID parsing ignores other header fields."""
     trace_header = (
@@ -94,7 +110,7 @@ def test_deterministic_id_generator_falls_back_to_random_trace_id(monkeypatch):
     expected_trace_id = int("1" * 32, 16)
     generator = DeterministicIdGenerator()
     monkeypatch.setattr(
-        generator._random_id_generator,
+        generator._fallback_id_generator,
         "generate_trace_id",
         lambda: expected_trace_id,
     )
@@ -108,7 +124,7 @@ def test_deterministic_id_generator_uses_next_span_id_once(monkeypatch):
     random_span_id = int("3" * 16, 16)
     generator = DeterministicIdGenerator()
     monkeypatch.setattr(
-        generator._random_id_generator,
+        generator._fallback_id_generator,
         "generate_span_id",
         lambda: random_span_id,
     )
@@ -124,11 +140,66 @@ def test_deterministic_id_generator_accepts_cleared_next_span_id(monkeypatch):
     expected_span_id = int("4" * 16, 16)
     generator = DeterministicIdGenerator()
     monkeypatch.setattr(
-        generator._random_id_generator,
+        generator._fallback_id_generator,
         "generate_span_id",
         lambda: expected_span_id,
     )
 
     generator.set_next_span_id(None)
 
     assert generator.generate_span_id() == expected_span_id
+
+
+def test_deterministic_id_generator_defaults_to_random_fallback():
+    """Verify the fallback defaults to a RandomIdGenerator when none is given."""
+    generator = DeterministicIdGenerator()
+
+    assert isinstance(generator._fallback_id_generator, RandomIdGenerator)
+
+
+def test_deterministic_id_generator_uses_provided_fallback_for_trace_id(monkeypatch):
+    """Verify the supplied fallback generator produces trace IDs when no
+    execution trace ID is set."""
+    monkeypatch.delenv("_X_AMZN_TRACE_ID", raising=False)
+    fallback = _StubIdGenerator(trace_id=int("a" * 32, 16), span_id=int("b" * 16, 16))
+    generator = DeterministicIdGenerator(fallback_id_generator=fallback)
+
+    assert generator.generate_trace_id() == int("a" * 32, 16)
+
+
+def test_deterministic_id_generator_uses_provided_fallback_for_span_id():
+    """Verify the supplied fallback generator produces span IDs when no
+    deterministic span ID is pending."""
+    fallback = _StubIdGenerator(trace_id=int("a" * 32, 16), span_id=int("b" * 16, 16))
+    generator = DeterministicIdGenerator(fallback_id_generator=fallback)
+
+    assert generator.generate_span_id() == int("b" * 16, 16)
+
+
+def test_deterministic_id_generator_prefers_execution_trace_id_over_fallback(
+    monkeypatch,
+):
+    """Verify a configured execution trace ID takes precedence over the fallback."""
+    monkeypatch.delenv("_X_AMZN_TRACE_ID", raising=False)
+    fallback = _StubIdGenerator(trace_id=int("a" * 32, 16), span_id=int("b" * 16, 16))
+    generator = DeterministicIdGenerator(fallback_id_generator=fallback)
+
+    generator.set_trace_id(
+        "arn:aws:lambda:us-west-2:123456789012:function:workflow:$LATEST",
+        datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC),
+    )
+
+    assert generator.generate_trace_id() == int("65937d253aa8c3f7ffe36c50d65b1a6d", 16)
+
+
+def test_deterministic_id_generator_prefers_next_span_id_over_fallback():
+    """Verify a pending deterministic span ID takes precedence over the fallback."""
+    deterministic_span_id = int("c" * 16, 16)
+    fallback = _StubIdGenerator(trace_id=int("a" * 32, 16), span_id=int("b" * 16, 16))
+    generator = DeterministicIdGenerator(fallback_id_generator=fallback)
+
+    generator.set_next_span_id(deterministic_span_id)
+
+    assert generator.generate_span_id() == deterministic_span_id
+    # Subsequent calls fall back to the provided generator.
+    assert generator.generate_span_id() == int("b" * 16, 16)

packages/aws-durable-execution-sdk-python-otel/tests/test_plugin.py

@@ -5,11 +5,6 @@
 from concurrent.futures import ThreadPoolExecutor
 from datetime import UTC, datetime
 
-from opentelemetry.context import Context
-from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.sdk.trace.export import SimpleSpanProcessor
-from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
-
 from aws_durable_execution_sdk_python.lambda_service import (
     InvocationStatus,
     OperationStatus,
@@ -24,6 +19,11 @@
     UserFunctionOutcome,
     UserFunctionStartInfo,
 )
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
 from aws_durable_execution_sdk_python_otel.deterministic_id_generator import (
     operation_id_to_span_id,
 )
@@ -131,7 +131,9 @@ def test_operation_end_without_start_emits_continuation_span_with_link():
     plugin.on_invocation_start(_invocation_start_info())
     operation_id = "wait-existing"
     random_span_id = int("1234567890abcdef", 16)
-    plugin._id_generator._random_id_generator.generate_span_id = lambda: random_span_id
+    plugin._id_generator._fallback_id_generator.generate_span_id = lambda: (
+        random_span_id
+    )
 
     plugin.on_operation_end(
         OperationEndInfo(

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/retries.py

@@ -125,6 +125,34 @@ def retry_strategy(error: Exception, attempts_made: int) -> RetryDecision:
     return retry_strategy
 
 
+def create_linear_retry_strategy(
+    max_attempts: int = 6,
+    initial_delay: Duration | None = None,
+    increment: Duration | None = None,
+) -> Callable[[Exception, int], RetryDecision]:
+    """Linearly increasing delay between retries: initial + increment * (attempts_made - 1).
+
+    Mirrors the JS SDK's ``createLinearRetryStrategy``. With the defaults this
+    yields delays of 1s, 2s, 3s, 4s, 5s. No jitter is applied and there is no
+    upper cap on the delay; callers who need either can build their own
+    strategy via ``create_retry_strategy``.
+    """
+    initial: Duration = (
+        initial_delay if initial_delay is not None else Duration.from_seconds(1)
+    )
+    step: Duration = increment if increment is not None else Duration.from_seconds(1)
+
+    def linear_retry_strategy(_error: Exception, attempts_made: int) -> RetryDecision:
+        if attempts_made >= max_attempts:
+            return RetryDecision.no_retry()
+        delay_seconds: int = initial.to_seconds() + step.to_seconds() * (
+            attempts_made - 1
+        )
+        return RetryDecision.retry(Duration(seconds=delay_seconds))
+
+    return linear_retry_strategy
+
+
 class RetryPresets:
     """Default retry presets."""
 
@@ -180,6 +208,31 @@ def critical(cls) -> Callable[[Exception, int], RetryDecision]:
             )
         )
 
+    @classmethod
+    def linear(cls) -> Callable[[Exception, int], RetryDecision]:
+        """Linearly increasing delay between retries: 1s, 2s, 3s, 4s, 5s."""
+        return create_linear_retry_strategy(
+            max_attempts=6,
+            initial_delay=Duration.from_seconds(1),
+            increment=Duration.from_seconds(1),
+        )
+
+    @classmethod
+    def fixed(
+        cls, interval: Duration | None = None
+    ) -> Callable[[Exception, int], RetryDecision]:
+        """Constant delay between retries (5s by default, no jitter)."""
+        delay: Duration = interval if interval is not None else Duration.from_seconds(5)
+        return create_retry_strategy(
+            RetryStrategyConfig(
+                max_attempts=5,
+                initial_delay=delay,
+                max_delay=delay,
+                backoff_rate=1,
+                jitter_strategy=JitterStrategy.NONE,
+            )
+        )
+
 
 @dataclass(frozen=True)
 class WithRetryConfig(Generic[T]):