feat(sdk-core): add SdkWarmUp.prime() for CRaC auto-priming

Author: joviegas

core/sdk-core/src/main/java/software/amazon/awssdk/core/crac/SdkWarmUp.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.crac;
+
+import java.util.ServiceLoader;
+import java.util.concurrent.atomic.AtomicBoolean;
+import software.amazon.awssdk.annotations.SdkPublicApi;
+import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.core.internal.crac.ClasspathWarmUpInvoker;
+
+/**
+ * Entry point for warming up SDK service request paths before a Coordinated Restore at Checkpoint (CRaC)
+ * checkpoint.
+ *
+ * <p>{@link #prime()} discovers every {@link SdkWarmUpProvider} registered on the classpath through {@link
+ * ServiceLoader} (via the {@code META-INF/services/software.amazon.awssdk.core.crac.SdkWarmUpProvider}
+ * resource) and invokes {@link SdkWarmUpProvider#warmUp()} on each.
+ *
+ * <p>Behavior contract:
+ * <ul>
+ *     <li><b>Idempotent:</b> {@code prime()} runs at most once per JVM. Subsequent calls return without
+ *     re-running discovery or invoking any provider.</li>
+ *     <li><b>Per-provider resilience:</b> a single provider that throws from {@code warmUp()}, or that fails
+ *     to load, does not prevent the remaining providers from running.</li>
+ *     <li><b>Safe when empty:</b> if no providers are registered, {@code prime()} is a no-op.</li>
+ * </ul>
+ *
+ * <p>Call this once during application initialization, before a CRaC checkpoint is taken.
+ */
+@ThreadSafe
+@SdkPublicApi
+public final class SdkWarmUp {
+
+    private static final AtomicBoolean PRIMED = new AtomicBoolean(false);
+
+    private SdkWarmUp() {
+    }
+
+    /**
+     * Discovers every {@link SdkWarmUpProvider} on the classpath and invokes {@link SdkWarmUpProvider#warmUp()}
+     * on each, honoring the idempotency, per-provider resilience, and empty-classpath behavior described on
+     * this class. Safe to call concurrently.
+     */
+    public static void prime() {
+        if (!PRIMED.compareAndSet(false, true)) {
+            return;
+        }
+
+        ClasspathWarmUpInvoker.create().invokeAll();
+    }
+}

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/crac/ClasspathWarmUpInvoker.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.internal.crac;
+
+import java.util.Iterator;
+import java.util.ServiceConfigurationError;
+import java.util.ServiceLoader;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
+import software.amazon.awssdk.core.crac.SdkWarmUpProvider;
+import software.amazon.awssdk.utils.Logger;
+
+/**
+ * {@link WarmUpInvoker} implementation that uses {@link ServiceLoader} to find {@link SdkWarmUpProvider}
+ * implementations on the classpath and invokes {@code warmUp()} on every one of them.
+ */
+@SdkInternalApi
+public final class ClasspathWarmUpInvoker implements WarmUpInvoker {
+
+    private static final Logger log = Logger.loggerFor(ClasspathWarmUpInvoker.class);
+
+    private final WarmUpServiceLoader serviceLoader;
+
+    @SdkTestInternalApi
+    ClasspathWarmUpInvoker(WarmUpServiceLoader serviceLoader) {
+        this.serviceLoader = serviceLoader;
+    }
+
+    @Override
+    public void invokeAll() {
+        Iterator<SdkWarmUpProvider> iterator = serviceLoader.loadProviders();
+        boolean invokedAny = false;
+
+        while (iterator.hasNext()) {
+            SdkWarmUpProvider provider;
+            try {
+                provider = iterator.next();
+            } catch (ServiceConfigurationError e) {
+                // next() has already advanced past the bad provider, so it is safe to continue to the next one.
+                log.warn(() -> "Skipping an SdkWarmUpProvider that could not be loaded.", e);
+                continue;
+            }
+
+            invokedAny = true;
+            try {
+                provider.warmUp();
+            } catch (RuntimeException e) {
+                log.warn(() -> "An SdkWarmUpProvider failed during warmUp() and was skipped.", e);
+            }
+        }
+
+        if (!invokedAny) {
+            log.debug(() -> "No SdkWarmUpProvider implementations were discovered on the classpath.");
+        }
+    }
+
+    /**
+     * @return ClasspathWarmUpInvoker that discovers {@link SdkWarmUpProvider}s from the classpath.
+     */
+    public static WarmUpInvoker create() {
+        return new ClasspathWarmUpInvoker(WarmUpServiceLoader.INSTANCE);
+    }
+}

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/crac/WarmUpInvoker.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.internal.crac;
+
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.crac.SdkWarmUpProvider;
+
+/**
+ * Discovers {@link SdkWarmUpProvider}s and invokes their warm-up behind the public {@code SdkWarmUp.prime()}.
+ * Mirrors the {@code SdkHttpServiceProvider} loader abstraction, except warm-up invokes every discovered
+ * provider rather than selecting one.
+ */
+@SdkInternalApi
+public interface WarmUpInvoker {
+
+    /**
+     * Invokes {@link SdkWarmUpProvider#warmUp()} on every discovered provider, containing per-provider failures
+     * so one failing provider does not stop the others.
+     */
+    void invokeAll();
+}

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/crac/WarmUpServiceLoader.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.internal.crac;
+
+import java.util.Iterator;
+import java.util.ServiceLoader;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.crac.SdkWarmUpProvider;
+import software.amazon.awssdk.core.internal.util.ClassLoaderHelper;
+
+/**
+ * Thin layer over {@link ServiceLoader} for {@link SdkWarmUpProvider}.
+ */
+@SdkInternalApi
+class WarmUpServiceLoader {
+
+    public static final WarmUpServiceLoader INSTANCE = new WarmUpServiceLoader();
+
+    Iterator<SdkWarmUpProvider> loadProviders() {
+        return ServiceLoader.load(SdkWarmUpProvider.class,
+                                  ClassLoaderHelper.classLoader(WarmUpServiceLoader.class)).iterator();
+    }
+}

core/sdk-core/src/test/java/software/amazon/awssdk/core/crac/CountingWarmUpProvider.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.crac;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Test-only {@link SdkWarmUpProvider} registered via test-scoped {@code META-INF/services} so the real static
+ * {@code SdkWarmUp.prime()} discovers and invokes it through {@link java.util.ServiceLoader}. It counts how many
+ * times {@code warmUp()} is invoked across the JVM. ServiceLoader requires a public no-arg constructor.
+ */
+public final class CountingWarmUpProvider implements SdkWarmUpProvider {
+
+    public static final AtomicInteger INVOCATIONS = new AtomicInteger();
+
+    @Override
+    public void warmUp() {
+        INVOCATIONS.incrementAndGet();
+    }
+}

core/sdk-core/src/test/java/software/amazon/awssdk/core/crac/SdkWarmUpTest.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.crac;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Tests for the static {@link SdkWarmUp#prime()} entry point, exercised end to end through {@link
+ * java.util.ServiceLoader} with a test-scoped {@code META-INF/services} registration of {@link
+ * CountingWarmUpProvider}.
+ *
+ * <p>{@code prime()} runs at most once per JVM and there is no reset hook, so assertions here are
+ * order-independent: regardless of how many threads call {@code prime()} (and regardless of any other test in
+ * this JVM having already called it), the counting provider is invoked exactly once in total.
+ */
+class SdkWarmUpTest {
+
+    @Test
+    void prime_concurrentCalls_invokeRegisteredProviderExactlyOnce() throws InterruptedException {
+        int threadCount = 16;
+        CountDownLatch start = new CountDownLatch(1);
+        CountDownLatch done = new CountDownLatch(threadCount);
+        List<Thread> threads = new ArrayList<>();
+
+        for (int i = 0; i < threadCount; i++) {
+            Thread thread = new Thread(() -> {
+                try {
+                    start.await();
+                    SdkWarmUp.prime();
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                } finally {
+                    done.countDown();
+                }
+            });
+            threads.add(thread);
+            thread.start();
+        }
+
+        start.countDown();
+        done.await();
+        for (Thread thread : threads) {
+            thread.join();
+        }
+
+        assertThat(CountingWarmUpProvider.INVOCATIONS.get()).isEqualTo(1);
+    }
+}