Cache clients by purpose, add purpose to HttpStatsCollector

Author: nicktindall

modules/repository-gcs/src/main/java/org/elasticsearch/repositories/gcs/GoogleCloudStorageBlobStore.java

@@ -130,7 +130,7 @@ class GoogleCloudStorageBlobStore implements BlobStore {
     }
 
     private Storage client(OperationPurpose purpose) throws IOException {
-        return storageService.client(clientName, repositoryName, stats);
+        return storageService.client(clientName, repositoryName, purpose, stats);
     }
 
     @Override
@@ -140,7 +140,7 @@ public BlobContainer blobContainer(BlobPath path) {
 
     @Override
     public void close() {
-        storageService.closeRepositoryClient(repositoryName);
+        storageService.closeRepositoryClients(repositoryName);
     }
 
     /**

modules/repository-gcs/src/main/java/org/elasticsearch/repositories/gcs/GoogleCloudStorageHttpStatsCollector.java

@@ -14,6 +14,8 @@
 import com.google.api.client.http.HttpResponse;
 import com.google.api.client.http.HttpResponseInterceptor;
 
+import org.elasticsearch.common.blobstore.OperationPurpose;
+
 import java.util.List;
 import java.util.Locale;
 import java.util.function.Consumer;
@@ -43,10 +45,12 @@ final class GoogleCloudStorageHttpStatsCollector implements HttpResponseIntercep
     );
 
     private final GoogleCloudStorageOperationsStats gcsOperationStats;
+    private final OperationPurpose operationPurpose;
     private final List<HttpRequestTracker> trackers;
 
-    GoogleCloudStorageHttpStatsCollector(final GoogleCloudStorageOperationsStats gcsOperationStats) {
+    GoogleCloudStorageHttpStatsCollector(final GoogleCloudStorageOperationsStats gcsOperationStats, OperationPurpose operationPurpose) {
         this.gcsOperationStats = gcsOperationStats;
+        this.operationPurpose = operationPurpose;
         this.trackers = trackerFactories.stream()
             .map(trackerFactory -> trackerFactory.apply(gcsOperationStats.getTrackedBucket()))
             .toList();

modules/repository-gcs/src/main/java/org/elasticsearch/repositories/gcs/GoogleCloudStorageService.java

@@ -25,6 +25,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.blobstore.OperationPurpose;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.SuppressForbidden;
@@ -40,6 +41,7 @@
 import java.net.URL;
 import java.security.KeyStore;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Collections.emptyMap;
@@ -51,12 +53,15 @@ public class GoogleCloudStorageService {
 
     private volatile Map<String, GoogleCloudStorageClientSettings> clientSettings = emptyMap();
 
+    private record ClientKey(OperationPurpose purpose, String repositoryName) {}
+
     /**
      * Dictionary of client instances. Client instances are built lazily from the
-     * latest settings. Each repository has its own client instance identified by
-     * the repository name.
+     * latest settings. Clients are cached by a composite OperationPurpose/repositoryName
+     * key.
+     * @see ClientKey
      */
-    private volatile Map<String, Storage> clientCache = emptyMap();
+    private volatile Map<ClientKey, Storage> clientCache = emptyMap();
 
     /**
      * Refreshes the client settings and clears the client cache. Subsequent calls to
@@ -79,20 +84,26 @@ public synchronized void refreshAndClearCache(Map<String, GoogleCloudStorageClie
      *
      * @param clientName name of the client settings used to create the client
      * @param repositoryName name of the repository that would use the client
+     * @param operationPurpose the purpose for which the client will be used
      * @param stats the stats collector used to gather information about the underlying SKD API calls.
      * @return a cached client storage instance that can be used to manage objects
      *         (blobs)
      */
-    public Storage client(final String clientName, final String repositoryName, final GoogleCloudStorageOperationsStats stats)
-        throws IOException {
+    public Storage client(
+        final String clientName,
+        final String repositoryName,
+        final OperationPurpose operationPurpose,
+        final GoogleCloudStorageOperationsStats stats
+    ) throws IOException {
+        ClientKey clientKey = new ClientKey(operationPurpose, repositoryName);
         {
-            final Storage storage = clientCache.get(repositoryName);
+            final Storage storage = clientCache.get(clientKey);
             if (storage != null) {
                 return storage;
             }
         }
         synchronized (this) {
-            final Storage existing = clientCache.get(repositoryName);
+            final Storage existing = clientCache.get(clientKey);
 
             if (existing != null) {
                 return existing;
@@ -110,26 +121,33 @@ public Storage client(final String clientName, final String repositoryName, fina
             }
 
             logger.debug(() -> format("creating GCS client with client_name [%s], endpoint [%s]", clientName, settings.getHost()));
-            final Storage storage = createClient(settings, stats);
-            clientCache = Maps.copyMapWithAddedEntry(clientCache, repositoryName, storage);
+            final Storage storage = createClient(settings, stats, operationPurpose);
+            clientCache = Maps.copyMapWithAddedEntry(clientCache, clientKey, storage);
             return storage;
         }
     }
 
-    synchronized void closeRepositoryClient(String repositoryName) {
-        clientCache = Maps.copyMapWithRemovedEntry(clientCache, repositoryName);
+    synchronized void closeRepositoryClients(String repositoryName) {
+        clientCache = clientCache.entrySet()
+            .stream()
+            .filter(entry -> entry.getKey().repositoryName().equals(repositoryName) == false)
+            .collect(Collectors.toUnmodifiableMap(Map.Entry::getKey, Map.Entry::getValue));
     }
 
     /**
      * Creates a client that can be used to manage Google Cloud Storage objects. The client is thread-safe.
      *
      * @param gcsClientSettings client settings to use, including secure settings
      * @param stats the stats collector to use by the underlying SDK
+     * @param operationPurpose the purpose this client will be used for
      * @return a new client storage instance that can be used to manage objects
      *         (blobs)
      */
-    private Storage createClient(GoogleCloudStorageClientSettings gcsClientSettings, GoogleCloudStorageOperationsStats stats)
-        throws IOException {
+    private Storage createClient(
+        GoogleCloudStorageClientSettings gcsClientSettings,
+        GoogleCloudStorageOperationsStats stats,
+        OperationPurpose operationPurpose
+    ) throws IOException {
         final HttpTransport httpTransport = SocketAccess.doPrivilegedIOException(() -> {
             final NetHttpTransport.Builder builder = new NetHttpTransport.Builder();
             // requires java.lang.RuntimePermission "setFactory"
@@ -149,7 +167,7 @@ private Storage createClient(GoogleCloudStorageClientSettings gcsClientSettings,
             return builder.build();
         });
 
-        final GoogleCloudStorageHttpStatsCollector httpStatsCollector = new GoogleCloudStorageHttpStatsCollector(stats);
+        final GoogleCloudStorageHttpStatsCollector httpStatsCollector = new GoogleCloudStorageHttpStatsCollector(stats, operationPurpose);
 
         final HttpTransportOptions httpTransportOptions = new HttpTransportOptions(
             HttpTransportOptions.newBuilder()