Add overloaded menthod for SdkWarmUp.prime overloaded to accept speci…

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

@@ -15,11 +15,21 @@
 
 package software.amazon.awssdk.core.crac;
 
+import java.util.LinkedHashSet;
 import java.util.ServiceLoader;
+import java.util.Set;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.core.ClientType;
+import software.amazon.awssdk.core.SdkClient;
 import software.amazon.awssdk.core.internal.crac.ClasspathWarmUpInvoker;
+import software.amazon.awssdk.core.internal.crac.PrimedClientRegistry;
+import software.amazon.awssdk.core.internal.crac.TargetedWarmUpInvoker;
+import software.amazon.awssdk.core.internal.crac.TargetedWarmUpResult;
+import software.amazon.awssdk.core.internal.http.loader.AsyncHttpClientWarmer;
 import software.amazon.awssdk.core.internal.http.loader.ClasspathHttpWarmupInvoker;
+import software.amazon.awssdk.core.internal.http.loader.SyncHttpClientWarmer;
+import software.amazon.awssdk.utils.Logger;
 
 /**
  * Entry point for warming up SDK service request paths before a Coordinated Restore at Checkpoint (CRaC)
@@ -45,17 +55,24 @@
 @SdkPublicApi
 public final class SdkWarmUp {
 
+    private static final Logger log = Logger.loggerFor(SdkWarmUp.class);
+
     private static final Object PRIME_LOCK = new Object();
 
     private static volatile boolean primed = false;
 
+    private static final PrimedClientRegistry PRIMED_CLIENTS = new PrimedClientRegistry();
+
     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.
+     *
+     * <p>This method and {@link #prime(Class[])} track primed state independently: this method warms
+     * every provider and does not skip clients that were warmed by a targeted {@link #prime(Class[])} call.
      */
     public static void prime() {
         if (primed) {
@@ -71,4 +88,46 @@ public static void prime() {
             primed = true;
         }
     }
+
+    /**
+     * Primes only the given service clients, warming the sync path for a sync client and the async path for an async
+     * client. A client already primed by this method is skipped; a class matched by no provider, or whose warm-up
+     * fails, is logged at warn and retried on the next call.Best-effort and safe to call concurrently.
+     *
+     * <p>This method and {@link #prime()} track primed state independently: this method does not skip
+     * clients that {@link #prime()} already warmed.
+     *
+     * @param clients the service client classes to prime, for example {@code ServiceClient.class} or
+     *                {@code ServiceAsyncClient.class}.
+     */
+    @SafeVarargs
+    public static void prime(Class<? extends SdkClient>... clients) {
+        if (clients == null || clients.length == 0) {
+            log.debug(() -> "SdkWarmUp.prime(Class...) called with no clients; nothing to do.");
+            return;
+        }
+
+        Set<String> requested = new LinkedHashSet<>();
+        for (Class<? extends SdkClient> client : clients) {
+            if (client != null) {
+                requested.add(client.getName());
+            }
+        }
+        Set<String> toPrime = PRIMED_CLIENTS.selectUnprimed(requested);
+        if (toPrime.isEmpty()) {
+            return;
+        }
+
+        // Racing calls may double-warm the same client; warming is idempotent, so that is harmless.
+        TargetedWarmUpResult result = TargetedWarmUpInvoker.create().invoke(toPrime);
+        if (result.matchedClientTypes().contains(ClientType.SYNC)) {
+            SyncHttpClientWarmer.create().warmAll();
+        }
+        if (result.matchedClientTypes().contains(ClientType.ASYNC)) {
+            AsyncHttpClientWarmer.create().warmAll();
+        }
+
+        // Only successfully warmed names are recorded; unmatched or failed ones are retried on a later call.
+        PRIMED_CLIENTS.markPrimed(result.warmedClientNames());
+    }
 }

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

@@ -0,0 +1,55 @@
+/*
+ * 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.Collection;
+import java.util.LinkedHashSet;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.ThreadSafe;
+
+/**
+ * Tracks which service clients {@code SdkWarmUp.prime(Class...)} has already warmed, so each is warmed at most once per
+ * JVM. State is per-instance, so it is unit-testable with fresh instances; production uses a single long-lived instance.
+ */
+@ThreadSafe
+@SdkInternalApi
+public final class PrimedClientRegistry {
+
+    private final Set<String> primed = ConcurrentHashMap.newKeySet();
+
+    /**
+     * Returns the not-yet-primed subset of {@code clientClassNames}, in encounter order, ignoring nulls. Warm the
+     * returned names, then pass them to {@link #markPrimed(Collection)}.
+     */
+    public Set<String> selectUnprimed(Collection<String> clientClassNames) {
+        Set<String> unprimed = new LinkedHashSet<>();
+        for (String name : clientClassNames) {
+            if (name != null && !primed.contains(name)) {
+                unprimed.add(name);
+            }
+        }
+        return unprimed;
+    }
+
+    /**
+     * Records the given names as primed. Call only after warming completes, so a failed run is retried by a later call.
+     */
+    public void markPrimed(Collection<String> clientClassNames) {
+        primed.addAll(clientClassNames);
+    }
+}

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

@@ -0,0 +1,106 @@
+/*
+ * 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.ArrayList;
+import java.util.Collection;
+import java.util.EnumSet;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Set;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
+import software.amazon.awssdk.core.ClientType;
+import software.amazon.awssdk.core.crac.SdkWarmUpProvider;
+import software.amazon.awssdk.utils.Logger;
+
+/**
+ * Warms only the service clients named by {@code SdkWarmUp.prime(Class...)}. Matches each requested client class name
+ * against the discovered {@link SdkWarmUpProvider}s' sync and async class names, then warms the matched client type only.
+ */
+@SdkInternalApi
+public final class TargetedWarmUpInvoker {
+
+    private static final Logger log = Logger.loggerFor(TargetedWarmUpInvoker.class);
+
+    private final WarmUpServiceLoader serviceLoader;
+
+    @SdkTestInternalApi
+    TargetedWarmUpInvoker(WarmUpServiceLoader serviceLoader) {
+        this.serviceLoader = serviceLoader;
+    }
+
+    public static TargetedWarmUpInvoker create() {
+        return new TargetedWarmUpInvoker(WarmUpServiceLoader.INSTANCE);
+    }
+
+    /**
+     * Warms the provider client type matching each requested client class name. An unmatched name is logged
+     * at warn and skipped, and a provider that fails does not stop the others.
+     *
+     * @param requestedClassNames the fully qualified sync or async client class names to warm.
+     * @return the matched client types and the names that warmed successfully.
+     */
+    public TargetedWarmUpResult invoke(Collection<String> requestedClassNames) {
+        Set<ClientType> matchedClientTypes = EnumSet.noneOf(ClientType.class);
+        Set<String> warmedClientNames = new LinkedHashSet<>();
+        if (requestedClassNames.isEmpty()) {
+            return new TargetedWarmUpResult(matchedClientTypes, warmedClientNames);
+        }
+
+        List<SdkWarmUpProvider> providers = loadProviders();
+        for (String requested : requestedClassNames) {
+            boolean warmFailed = false;
+            Set<ClientType> matched = EnumSet.noneOf(ClientType.class);
+            for (SdkWarmUpProvider provider : providers) {
+                ClientType clientType = clientTypeFor(requested, provider);
+                if (clientType == null) {
+                    continue;
+                }
+                matched.add(clientType);
+                try {
+                    provider.warmUpClient(clientType);
+                } catch (RuntimeException | LinkageError e) {
+                    warmFailed = true;
+                    log.warn(() -> "Warm-up failed for " + provider.getClass().getName() + " and was skipped.", e);
+                }
+            }
+            if (matched.isEmpty()) {
+                log.warn(() -> "No warm-up provider matched client class " + requested + "; skipping.");
+            } else if (!warmFailed) {
+                warmedClientNames.add(requested);
+            }
+            matchedClientTypes.addAll(matched);
+        }
+        return new TargetedWarmUpResult(matchedClientTypes, warmedClientNames);
+    }
+
+    private ClientType clientTypeFor(String requested, SdkWarmUpProvider provider) {
+        if (requested.equals(provider.syncClientClassName())) {
+            return ClientType.SYNC;
+        }
+        if (requested.equals(provider.asyncClientClassName())) {
+            return ClientType.ASYNC;
+        }
+        return null;
+    }
+
+    private List<SdkWarmUpProvider> loadProviders() {
+        List<SdkWarmUpProvider> providers = new ArrayList<>();
+        WarmUpDiscovery.forEachDiscovered(serviceLoader.loadProviders(), providers::add);
+        return providers;
+    }
+}

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

@@ -0,0 +1,78 @@
+/*
+ * 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.Collections;
+import java.util.Set;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.ClientType;
+import software.amazon.awssdk.utils.ToString;
+
+/**
+ * Outcome of a {@link TargetedWarmUpInvoker} run: the client types that matched and the client names that warmed
+ * successfully.
+ */
+@SdkInternalApi
+public final class TargetedWarmUpResult {
+
+    private final Set<ClientType> matchedClientTypes;
+    private final Set<String> warmedClientNames;
+
+    TargetedWarmUpResult(Set<ClientType> matchedClientTypes, Set<String> warmedClientNames) {
+        this.matchedClientTypes = Collections.unmodifiableSet(matchedClientTypes);
+        this.warmedClientNames = Collections.unmodifiableSet(warmedClientNames);
+    }
+
+    public Set<ClientType> matchedClientTypes() {
+        return matchedClientTypes;
+    }
+
+    public Set<String> warmedClientNames() {
+        return warmedClientNames;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
+        TargetedWarmUpResult that = (TargetedWarmUpResult) o;
+
+        if (!matchedClientTypes.equals(that.matchedClientTypes)) {
+            return false;
+        }
+        return warmedClientNames.equals(that.warmedClientNames);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = matchedClientTypes.hashCode();
+        result = 31 * result + warmedClientNames.hashCode();
+        return result;
+    }
+
+    @Override
+    public String toString() {
+        return ToString.builder("TargetedWarmUpResult")
+                       .add("matchedClientTypes", matchedClientTypes)
+                       .add("warmedClientNames", warmedClientNames)
+                       .build();
+    }
+}

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

@@ -0,0 +1,24 @@
+/*
+ * 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 software.amazon.awssdk.core.SdkClient;
+
+/**
+ * Test-only client whose fully qualified name matches {@link RegisteredWarmUpProvider#asyncClientClassName()}.
+ */
+public interface RegisteredAsyncClient extends SdkClient {
+}

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

@@ -0,0 +1,25 @@
+/*
+ * 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 software.amazon.awssdk.core.SdkClient;
+
+/**
+ * Test-only client whose fully qualified name matches {@link RegisteredWarmUpProvider#syncClientClassName()}, so
+ * {@code SdkWarmUp.prime(RegisteredSyncClient.class)} matches that provider's sync client type.
+ */
+public interface RegisteredSyncClient extends SdkClient {
+}