chore(ci): Improve reliability of retries in TracingE2ET

powertools-e2e-tests/pom.xml

@@ -150,6 +150,16 @@
             <version>2.4</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>com.fasterxml.jackson.core</groupId>
+            <artifactId>jackson-databind</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>com.fasterxml.jackson.datatype</groupId>
+            <artifactId>jackson-datatype-jsr310</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.aspectj</groupId>
             <artifactId>aspectjrt</artifactId>

powertools-e2e-tests/src/test/java/software/amazon/lambda/powertools/MetricsE2ET.java

@@ -34,7 +34,9 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
+import software.amazon.lambda.powertools.testutils.DataNotReadyException;
 import software.amazon.lambda.powertools.testutils.Infrastructure;
+import software.amazon.lambda.powertools.testutils.RetryUtils;
 import software.amazon.lambda.powertools.testutils.lambda.InvocationResult;
 import software.amazon.lambda.powertools.testutils.metrics.MetricsFetcher;
 
@@ -89,16 +91,20 @@ void test_recordMetrics() {
                         { "FunctionName", functionName },
                         { "Service", SERVICE } }).collect(Collectors.toMap(data -> data[0], data -> data[1])));
         assertThat(coldStart.get(0)).isEqualTo(1);
-        List<Double> orderMetrics = metricsFetcher.fetchMetrics(invocationResult.getStart(), invocationResult.getEnd(),
-                60, NAMESPACE,
-                "orders", Collections.singletonMap("Environment", "test"));
+        List<Double> orderMetrics = RetryUtils.withRetry(() -> {
+            List<Double> metrics = metricsFetcher.fetchMetrics(invocationResult.getStart(), invocationResult.getEnd(),
+                    60, NAMESPACE, "orders", Collections.singletonMap("Environment", "test"));
+            if (metrics.get(0) != 2.0) {
+                throw new DataNotReadyException("Expected 2.0 orders but got " + metrics.get(0));
+            }
+            return metrics;
+        }, "orderMetricsRetry", DataNotReadyException.class).get();
         assertThat(orderMetrics.get(0)).isEqualTo(2);
         List<Double> productMetrics = metricsFetcher.fetchMetrics(invocationResult.getStart(),
                 invocationResult.getEnd(), 60, NAMESPACE,
                 "products", Collections.singletonMap("Environment", "test"));
 
         // When searching across a 1 minute time period with a period of 60 we find both metrics and the sum is 12
-
         assertThat(productMetrics.get(0)).isEqualTo(12);
 
         orderMetrics = metricsFetcher.fetchMetrics(invocationResult.getStart(), invocationResult.getEnd(), 60,

powertools-e2e-tests/src/test/java/software/amazon/lambda/powertools/testutils/RetryUtils.java

@@ -47,7 +47,21 @@ private RetryUtils() {
      */
     @SafeVarargs
     public static Retry createRetry(String name, Class<? extends Throwable>... retryOnThrowables) {
-        RetryConfig config = RetryConfig.from(DEFAULT_RETRY_CONFIG)
+        return createRetry(name, DEFAULT_RETRY_CONFIG, retryOnThrowables);
+    }
+
+    /**
+     * Creates a retry instance with custom configuration for the specified throwable types.
+     * 
+     * @param name the name for the retry instance
+     * @param customConfig the custom retry configuration
+     * @param retryOnThrowables the throwable classes to retry on
+     * @return configured Retry instance
+     */
+    @SafeVarargs
+    public static Retry createRetry(String name, RetryConfig customConfig,
+            Class<? extends Throwable>... retryOnThrowables) {
+        RetryConfig config = RetryConfig.from(customConfig)
                 .retryExceptions(retryOnThrowables)
                 .build();
 
@@ -72,4 +86,20 @@ public static <T> Supplier<T> withRetry(Supplier<T> supplier, String name,
         Retry retry = createRetry(name, retryOnThrowables);
         return Retry.decorateSupplier(retry, supplier);
     }
+
+    /**
+     * Decorates a supplier with custom retry logic for the specified throwable types.
+     * 
+     * @param supplier the supplier to decorate
+     * @param name the name for the retry instance
+     * @param customConfig the custom retry configuration
+     * @param retryOnThrowables the throwable classes to retry on
+     * @return decorated supplier with retry logic
+     */
+    @SafeVarargs
+    public static <T> Supplier<T> withRetry(Supplier<T> supplier, String name, RetryConfig customConfig,
+            Class<? extends Throwable>... retryOnThrowables) {
+        Retry retry = createRetry(name, customConfig, retryOnThrowables);
+        return Retry.decorateSupplier(retry, supplier);
+    }
 }

powertools-e2e-tests/src/test/java/software/amazon/lambda/powertools/testutils/tracing/TraceFetcher.java

@@ -14,6 +14,7 @@
 
 package software.amazon.lambda.powertools.testutils.tracing;
 
+import java.time.Duration;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
 import java.util.Arrays;
@@ -27,8 +28,11 @@
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.fasterxml.jackson.databind.MapperFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.json.JsonMapper;
 
+import io.github.resilience4j.retry.RetryConfig;
 import software.amazon.awssdk.http.SdkHttpClient;
 import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient;
 import software.amazon.awssdk.regions.Region;
@@ -46,9 +50,10 @@
  * Class in charge of retrieving the actual traces of a Lambda execution on X-Ray
  */
 public class TraceFetcher {
-
-    private static final ObjectMapper MAPPER = new ObjectMapper()
-            .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
+    private static final ObjectMapper MAPPER = JsonMapper.builder()
+            .disable(MapperFeature.REQUIRE_HANDLERS_FOR_JAVA8_TIMES)
+            .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)
+            .build();
     private static final Logger LOG = LoggerFactory.getLogger(TraceFetcher.class);
     private static final SdkHttpClient httpClient = UrlConnectionHttpClient.builder().build();
     private static final Region region = Region.of(System.getProperty("AWS_DEFAULT_REGION", "eu-west-1"));
@@ -90,7 +95,12 @@ public Trace fetchTrace() {
             return getTrace(traceIds);
         };
 
-        return RetryUtils.withRetry(supplier, "trace-fetcher", TraceNotFoundException.class).get();
+        RetryConfig customConfig = RetryConfig.custom()
+                .maxAttempts(120) // 120 attempts over 10 minutes
+                .waitDuration(Duration.ofSeconds(5)) // 5 seconds between attempts
+                .build();
+
+        return RetryUtils.withRetry(supplier, "trace-fetcher", customConfig, TraceNotFoundException.class).get();
     }
 
     /**
@@ -104,26 +114,41 @@ private Trace getTrace(List<String> traceIds) {
                 .traceIds(traceIds)
                 .build());
         if (!tracesResponse.hasTraces()) {
-            throw new TraceNotFoundException("No trace found");
+            throw new TraceNotFoundException(String.format("No trace found for traceIds %s", traceIds));
         }
+
         Trace traceRes = new Trace();
         tracesResponse.traces().forEach(trace -> {
             if (trace.hasSegments()) {
                 trace.segments().forEach(segment -> {
                     try {
                         SegmentDocument document = MAPPER.readValue(segment.document(), SegmentDocument.class);
                         if ("AWS::Lambda::Function".equals(document.getOrigin()) && document.hasSubsegments()) {
-                            getNestedSubSegments(document.getSubsegments(), traceRes,
-                                    Collections.emptyList());
+                            LOG.debug("Populating subsegments for document {}", MAPPER.writeValueAsString(document));
+                            getNestedSubSegments(document.getSubsegments(), traceRes, Collections.emptyList());
+                            // If only the default (excluded) subsegments were populated we need to keep retrying for
+                            // our custom subsegments. They might appear later.
+                            if (traceRes.getSubsegments().isEmpty()) {
+                                throw new TraceNotFoundException(
+                                        "Found AWS::Lambda::Function SegmentDocument with no non-excluded subsegments.");
+                            }
+                        } else if ("AWS::Lambda::Function".equals(document.getOrigin())) {
+                            LOG.debug(
+                                    "Found AWS::Lambda::Function SegmentDocument with no subsegments. Retrying {}",
+                                    MAPPER.writeValueAsString(document));
+                            throw new TraceNotFoundException(
+                                    "Found AWS::Lambda::Function SegmentDocument with no subsegments.");
                         }
-
                     } catch (JsonProcessingException e) {
                         LOG.error("Failed to parse segment document: " + e.getMessage());
                         throw new RuntimeException(e);
                     }
                 });
+            } else {
+                throw new TraceNotFoundException(String.format("No segments found in trace %s", trace.id()));
             }
         });
+
         return traceRes;
     }
 
@@ -149,21 +174,30 @@ private void getNestedSubSegments(List<SubSegment> subsegments, Trace traceRes,
      * @return a list of trace ids
      */
     private List<String> getTraceIds() {
+        LOG.debug("Searching for traces from {} to {} with filter: {}", start, end, filterExpression);
         GetTraceSummariesResponse traceSummaries = xray.getTraceSummaries(GetTraceSummariesRequest.builder()
                 .startTime(start)
                 .endTime(end)
-                .timeRangeType(TimeRangeType.EVENT)
+                .timeRangeType(TimeRangeType.TRACE_ID)
                 .sampling(false)
                 .filterExpression(filterExpression)
                 .build());
+
+        LOG.debug("Found {} trace summaries",
+                traceSummaries.hasTraceSummaries() ? traceSummaries.traceSummaries().size() : 0);
+
         if (!traceSummaries.hasTraceSummaries()) {
-            throw new TraceNotFoundException("No trace id found");
+            throw new TraceNotFoundException(String.format("No trace id found for filter '%s' between %s and %s",
+                    filterExpression, start, end));
         }
         List<String> traceIds = traceSummaries.traceSummaries().stream().map(TraceSummary::id)
                 .collect(Collectors.toList());
         if (traceIds.isEmpty()) {
-            throw new TraceNotFoundException("No trace id found");
+            throw new TraceNotFoundException(
+                    String.format("Empty trace summary found for filter '%s' between %s and %s",
+                            filterExpression, start, end));
         }
+        LOG.debug("Found trace IDs: {}", traceIds);
         return traceIds;
     }
 
@@ -183,9 +217,13 @@ public TraceFetcher build() {
             if (end == null) {
                 end = start.plus(1, ChronoUnit.MINUTES);
             }
-            LOG.debug("Looking for traces from {} to {} with filter {} and excluded segments {}", start, end,
-                    filterExpression, excludedSegments);
-            return new TraceFetcher(start, end, filterExpression, excludedSegments);
+            // Expand search window by 1 minute on each side to account for timing imprecisions
+            Instant expandedStart = start.minus(1, ChronoUnit.MINUTES);
+            Instant expandedEnd = end.plus(1, ChronoUnit.MINUTES);
+            LOG.debug(
+                    "Looking for traces from {} to {} (expanded from {} to {}) with filter {} and excluded segments {}",
+                    expandedStart, expandedEnd, start, end, filterExpression, excludedSegments);
+            return new TraceFetcher(expandedStart, expandedEnd, filterExpression, excludedSegments);
         }
 
         public Builder start(Instant start) {