update

Author: mhl-b

modules/repository-gcs/src/internalClusterTest/java/org/elasticsearch/repositories/gcs/GoogleCloudStorageBlobStoreRepositoryTests.java

@@ -277,7 +277,7 @@ public Map<String, Repository.Factory> getRepositories(
                     clusterService,
                     bigArrays,
                     recoverySettings,
-                    repositoriesMetrics
+                    new RepositoryStatsCollector()
                 ) {
                     @Override
                     protected GoogleCloudStorageBlobStore createBlobStore() {
@@ -289,7 +289,7 @@ protected GoogleCloudStorageBlobStore createBlobStore() {
                             bigArrays,
                             randomIntBetween(1, 8) * 1024,
                             BackoffPolicy.noBackoff(),
-                            repositoriesMetrics
+                            this.statsCollector()
                         ) {
                             @Override
                             long getLargeBlobThresholdInBytes() {

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

@@ -43,7 +43,6 @@
 import org.elasticsearch.core.Streams;
 import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.core.TimeValue;
-import org.elasticsearch.repositories.RepositoriesMetrics;
 import org.elasticsearch.rest.RestStatus;
 
 import java.io.ByteArrayInputStream;
@@ -123,14 +122,14 @@ class GoogleCloudStorageBlobStore implements BlobStore {
         BigArrays bigArrays,
         int bufferSize,
         BackoffPolicy casBackoffPolicy,
-        RepositoriesMetrics repositoriesMetrics
+        RepositoryStatsCollector statsCollector
     ) {
         this.bucketName = bucketName;
         this.clientName = clientName;
         this.repositoryName = repositoryName;
         this.storageService = storageService;
         this.bigArrays = bigArrays;
-        this.statsCollector = new RepositoryStatsCollector(repositoryName, repositoriesMetrics);
+        this.statsCollector = statsCollector;
         this.bufferSize = bufferSize;
         this.casBackoffPolicy = casBackoffPolicy;
     }

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

@@ -62,7 +62,7 @@ public Map<String, Repository.Factory> getRepositories(
                 clusterService,
                 bigArrays,
                 recoverySettings,
-                repositoriesMetrics
+                new RepositoryStatsCollector(() -> clusterService.threadPool().absoluteTimeInMillis(), metadata, repositoriesMetrics)
             )
         );
     }

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

@@ -22,7 +22,6 @@
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.indices.recovery.RecoverySettings;
-import org.elasticsearch.repositories.RepositoriesMetrics;
 import org.elasticsearch.repositories.RepositoryException;
 import org.elasticsearch.repositories.blobstore.MeteredBlobStoreRepository;
 import org.elasticsearch.xcontent.NamedXContentRegistry;
@@ -86,7 +85,7 @@ class GoogleCloudStorageRepository extends MeteredBlobStoreRepository {
     private final TimeValue retryThrottledCasDelayIncrement;
     private final int retryThrottledCasMaxNumberOfRetries;
     private final TimeValue retryThrottledCasMaxDelay;
-    private final RepositoriesMetrics repositoriesMetrics;
+    private final RepositoryStatsCollector statsCollector;
 
     GoogleCloudStorageRepository(
         final RepositoryMetadata metadata,
@@ -95,7 +94,7 @@ class GoogleCloudStorageRepository extends MeteredBlobStoreRepository {
         final ClusterService clusterService,
         final BigArrays bigArrays,
         final RecoverySettings recoverySettings,
-        final RepositoriesMetrics repositoriesMetrics
+        final RepositoryStatsCollector statsCollector
     ) {
         super(
             metadata,
@@ -113,7 +112,7 @@ class GoogleCloudStorageRepository extends MeteredBlobStoreRepository {
         this.retryThrottledCasDelayIncrement = RETRY_THROTTLED_CAS_DELAY_INCREMENT.get(metadata.settings());
         this.retryThrottledCasMaxNumberOfRetries = RETRY_THROTTLED_CAS_MAX_NUMBER_OF_RETRIES.get(metadata.settings());
         this.retryThrottledCasMaxDelay = RETRY_THROTTLED_CAS_MAXIMUM_DELAY.get(metadata.settings());
-        this.repositoriesMetrics = repositoriesMetrics;
+        this.statsCollector = statsCollector;
         logger.debug("using bucket [{}], base_path [{}], chunk_size [{}], compress [{}]", bucket, basePath(), chunkSize, isCompress());
     }
 
@@ -144,7 +143,7 @@ protected GoogleCloudStorageBlobStore createBlobStore() {
             bigArrays,
             bufferSize,
             BackoffPolicy.linearBackoff(retryThrottledCasDelayIncrement, retryThrottledCasMaxNumberOfRetries, retryThrottledCasMaxDelay),
-            repositoriesMetrics
+            statsCollector
         );
     }
 
@@ -153,6 +152,10 @@ protected ByteSizeValue chunkSize() {
         return chunkSize;
     }
 
+    RepositoryStatsCollector statsCollector() {
+        return statsCollector;
+    }
+
     /**
      * Get a given setting from the repository settings, throwing a {@link RepositoryException} if the setting does not exist or is empty.
      */

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

@@ -171,7 +171,7 @@ public HttpRequestInitializer getHttpRequestInitializer(ServiceOptions<?, ?> ser
 
                 return (httpRequest) -> {
                     if (requestInitializer != null) requestInitializer.initialize(httpRequest);
-                    httpRequest.setResponseInterceptor(new MeteringResponseInterceptor());
+                    httpRequest.setResponseInterceptor(RepositoryStatsCollector.METERING_INTERCEPTOR);
                 };
             }
         };

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

@@ -23,9 +23,9 @@
 import com.google.cloud.storage.StorageOptions;
 import com.google.cloud.storage.spi.v1.HttpStorageRpc;
 
+import org.apache.lucene.util.IORunnable;
 import org.elasticsearch.SpecialPermission;
 import org.elasticsearch.common.blobstore.OperationPurpose;
-import org.elasticsearch.core.CheckedRunnable;
 import org.elasticsearch.core.SuppressForbidden;
 
 import java.io.IOException;
@@ -42,6 +42,8 @@
 
 /**
  * A wrapper for GCP {@link Storage} client. Provides metering and telemetry.
+ * It's mostly boilerplate code that wraps Storage calls with thread-local metrics.
+ * And special cases, such as Paginated List calls, WriteChannel, ReadChannel.
  */
 public class MeteredStorage {
     private final Storage storage;
@@ -108,13 +110,17 @@ public MeteredObjectsGetRequest meteredObjectsGet(OperationPurpose purpose, Stri
         return new MeteredObjectsGetRequest(statsCollector, purpose, storageRpc.objects().get(bucket, blob));
     }
 
-    public MeteredWriteChannel meteredWriter(OperationPurpose purpose, BlobInfo blobInfo, Storage.BlobWriteOption... writeOptions) {
-        var initStats = OperationStats.initAndGet(purpose, INSERT);
-        return new MeteredWriteChannel(statsCollector, initStats, storage.writer(blobInfo, writeOptions));
+    public MeteredWriteChannel meteredWriter(OperationPurpose purpose, BlobInfo blobInfo, Storage.BlobWriteOption... writeOptions)
+        throws IOException {
+        var initStats = new OperationStats(purpose, INSERT);
+        return statsCollector.continueAndCollect(
+            initStats,
+            () -> new MeteredWriteChannel(statsCollector, initStats, storage.writer(blobInfo, writeOptions))
+        );
     }
 
     public MeteredReadChannel meteredReader(OperationPurpose purpose, BlobId blobId, Storage.BlobSourceOption... options) {
-        var initStats = OperationStats.initAndGet(purpose, GET);
+        var initStats = new OperationStats(purpose, GET);
         return new MeteredReadChannel(statsCollector, initStats, storage.reader(blobId, options));
     }
 
@@ -186,7 +192,7 @@ public boolean isOpen() {
 
         @Override
         public void close() throws IOException {
-            statsCollector.finishAndCollect(stats, (CheckedRunnable<IOException>) writeChannel::close);
+            statsCollector.finishAndCollect(stats, (IORunnable) writeChannel::close);
         }
     }
 

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

@@ -13,7 +13,6 @@
 
 public class OperationStats {
 
-    private static final ThreadLocal<OperationStats> OPERATION_STATS = new ThreadLocal<>();
     final OperationPurpose purpose;
     final StorageOperation operation;
 
@@ -58,24 +57,4 @@ public class OperationStats {
         this.startTimeMs = 0;
     }
 
-    static OperationStats initAndGet(OperationPurpose purpose, StorageOperation operation) {
-        var stats = new OperationStats(purpose, operation);
-        OPERATION_STATS.set(stats);
-        return stats;
-    }
-
-    static void set(OperationStats stats) {
-        OPERATION_STATS.set(stats);
-    }
-
-    static OperationStats get() {
-        var stats = OPERATION_STATS.get();
-        assert stats != null : "must initialize operation stats";
-        return stats;
-    }
-
-    static void clear() {
-        OPERATION_STATS.remove();
-    }
-
 }