Add metric to track data downloaded from S3 (#886)

* Add metric to track data downloaded from S3

Author: aprudhomme

src/main/java/com/yelp/nrtsearch/server/config/NrtsearchConfig.java

@@ -115,6 +115,7 @@ public class NrtsearchConfig {
   private final DirectoryFactory.MMapGrouping mmapGrouping;
   private final boolean requireIdField;
   private final IsolatedReplicaConfig isolatedReplicaConfig;
+  private final boolean s3Metrics;
 
   @Inject
   public NrtsearchConfig(InputStream yamlStream) {
@@ -196,6 +197,7 @@ public NrtsearchConfig(InputStream yamlStream) {
     useSeparateCommitExecutor = configReader.getBoolean("useSeparateCommitExecutor", false);
     requireIdField = configReader.getBoolean("requireIdField", false);
     isolatedReplicaConfig = IsolatedReplicaConfig.fromConfig(configReader);
+    s3Metrics = configReader.getBoolean("s3Metrics", false);
 
     List<String> indicesWithOverrides = configReader.getKeysOrEmpty("indexLiveSettingsOverrides");
     Map<String, IndexLiveSettings> liveSettingsMap = new HashMap<>();
@@ -401,6 +403,10 @@ public IsolatedReplicaConfig getIsolatedReplicaConfig() {
     return isolatedReplicaConfig;
   }
 
+  public boolean getS3Metrics() {
+    return s3Metrics;
+  }
+
   /**
    * Substitute all sub strings of the form ${FOO} with the environment variable value env[FOO].
    * Variable names may only contain letters, numbers, and underscores. If a variable is not present

src/main/java/com/yelp/nrtsearch/server/grpc/NrtsearchServer.java

@@ -93,6 +93,7 @@
 import com.yelp.nrtsearch.server.monitoring.NrtsearchMonitoringServerInterceptor;
 import com.yelp.nrtsearch.server.monitoring.ProcStatCollector;
 import com.yelp.nrtsearch.server.monitoring.QueryCacheCollector;
+import com.yelp.nrtsearch.server.monitoring.S3DownloadStreamWrapper;
 import com.yelp.nrtsearch.server.monitoring.SearchResponseCollector;
 import com.yelp.nrtsearch.server.monitoring.ThreadPoolCollector;
 import com.yelp.nrtsearch.server.monitoring.ThreadPoolCollector.RejectionCounterWrapper;
@@ -272,6 +273,7 @@ private void registerMetrics(GlobalState globalState) {
     prometheusRegistry.register(new SearchResponseCollector(globalState));
     // register Indexing metrics such as individual addDocument, updateDocValue latencies and qps
     IndexingMetrics.register(prometheusRegistry);
+    S3DownloadStreamWrapper.register(prometheusRegistry);
   }
 
   /** Main launches the server from the command line. */

src/main/java/com/yelp/nrtsearch/server/monitoring/S3DownloadStreamWrapper.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2025 Yelp Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License 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 com.yelp.nrtsearch.server.monitoring;
+
+import io.prometheus.metrics.core.datapoints.CounterDataPoint;
+import io.prometheus.metrics.core.metrics.Counter;
+import io.prometheus.metrics.model.registry.PrometheusRegistry;
+import java.io.InputStream;
+import org.apache.commons.io.input.ProxyInputStream;
+
+/**
+ * InputStream wrapper that counts the number of bytes read from an S3 download stream and updates a
+ * Prometheus counter with the total bytes downloaded per index.
+ */
+public class S3DownloadStreamWrapper extends ProxyInputStream {
+  public static final Counter nrtS3DownloadBytes =
+      Counter.builder()
+          .name("nrt_s3_download_bytes_total")
+          .help("Total number of bytes downloaded from S3.")
+          .labelNames("index")
+          .build();
+
+  public static void register(PrometheusRegistry registry) {
+    registry.register(nrtS3DownloadBytes);
+  }
+
+  private final CounterDataPoint counterDataPoint;
+
+  public S3DownloadStreamWrapper(InputStream proxy, String indexName) {
+    super(proxy);
+    counterDataPoint = nrtS3DownloadBytes.labelValues(indexName);
+  }
+
+  @Override
+  protected void afterRead(int n) {
+    if (n != -1) {
+      counterDataPoint.inc(n);
+    }
+  }
+}

src/main/java/com/yelp/nrtsearch/server/remote/s3/S3Backend.java

@@ -34,9 +34,11 @@
 import com.amazonaws.services.s3.transfer.Upload;
 import com.google.common.annotations.VisibleForTesting;
 import com.yelp.nrtsearch.server.config.NrtsearchConfig;
+import com.yelp.nrtsearch.server.monitoring.S3DownloadStreamWrapper;
 import com.yelp.nrtsearch.server.nrt.state.NrtFileMetaData;
 import com.yelp.nrtsearch.server.nrt.state.NrtPointState;
 import com.yelp.nrtsearch.server.remote.RemoteBackend;
+import com.yelp.nrtsearch.server.state.BackendGlobalState;
 import com.yelp.nrtsearch.server.state.StateUtils;
 import com.yelp.nrtsearch.server.utils.TimeStringUtils;
 import com.yelp.nrtsearch.server.utils.ZipUtils;
@@ -88,6 +90,7 @@ public class S3Backend implements RemoteBackend {
   private final AmazonS3 s3;
   private final String serviceBucket;
   private final TransferManager transferManager;
+  private final boolean s3Metrics;
 
   /**
    * Pair of file names, one for the local file and one for the backend file.
@@ -104,17 +107,23 @@ record FileNamePair(String fileName, String backendFileName) {}
    * @param s3 s3 client
    */
   public S3Backend(NrtsearchConfig configuration, AmazonS3 s3) {
-    this(configuration.getBucketName(), configuration.getSavePluginBeforeUnzip(), s3);
+    this(
+        configuration.getBucketName(),
+        configuration.getSavePluginBeforeUnzip(),
+        configuration.getS3Metrics(),
+        s3);
   }
 
   /**
    * Constructor.
    *
    * @param serviceBucket bucket name
    * @param savePluginBeforeUnzip save plugin before unzipping
+   * @param s3Metrics enable s3 download metrics
    * @param s3 s3 client
    */
-  public S3Backend(String serviceBucket, boolean savePluginBeforeUnzip, AmazonS3 s3) {
+  public S3Backend(
+      String serviceBucket, boolean savePluginBeforeUnzip, boolean s3Metrics, AmazonS3 s3) {
     this.s3 = s3;
     this.executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(NUM_S3_THREADS);
     this.saveBeforeUnzip = savePluginBeforeUnzip;
@@ -125,6 +134,7 @@ public S3Backend(String serviceBucket, boolean savePluginBeforeUnzip, AmazonS3 s
             .withExecutorFactory(() -> executor)
             .withShutDownThreadPools(false)
             .build();
+    this.s3Metrics = s3Metrics;
   }
 
   public AmazonS3 getS3() {
@@ -346,7 +356,7 @@ public void uploadGlobalState(String service, byte[] data) throws IOException {
   @Override
   public InputStream downloadGlobalState(String service) throws IOException {
     String prefix = getGlobalStateResourcePrefix(service);
-    return downloadResource(prefix);
+    return downloadResource(prefix, null);
   }
 
   @Override
@@ -360,7 +370,7 @@ public void uploadIndexState(String service, String indexIdentifier, byte[] data
   @Override
   public InputStream downloadIndexState(String service, String indexIdentifier) throws IOException {
     String prefix = getIndexResourcePrefix(service, indexIdentifier, IndexResourceType.INDEX_STATE);
-    return downloadResource(prefix);
+    return downloadResource(prefix, indexIdentifier);
   }
 
   @Override
@@ -377,7 +387,7 @@ public InputStream downloadWarmingQueries(String service, String indexIdentifier
       throws IOException {
     String prefix =
         getIndexResourcePrefix(service, indexIdentifier, IndexResourceType.WARMING_QUERIES);
-    return downloadResource(prefix);
+    return downloadResource(prefix, indexIdentifier);
   }
 
   @Override
@@ -472,7 +482,8 @@ public InputStream downloadIndexFile(
     String backendFileName = getIndexBackendFileName(fileName, fileMetaData);
     String backendPrefix = getIndexDataPrefix(service, indexIdentifier);
     String backendKey = backendPrefix + backendFileName;
-    return downloadFromS3Path(serviceBucket, backendKey, false);
+    return wrapDownloadStream(
+        downloadFromS3Path(serviceBucket, backendKey, false), s3Metrics, indexIdentifier);
   }
 
   @VisibleForTesting
@@ -486,6 +497,17 @@ static List<FileNamePair> getFileNamePairs(Map<String, NrtFileMetaData> files) {
     return fileList;
   }
 
+  @VisibleForTesting
+  static InputStream wrapDownloadStream(
+      InputStream inputStream, boolean s3Metrics, String indexIdentifier) {
+    if (s3Metrics) {
+      String indexName = BackendGlobalState.getBaseIndexName(indexIdentifier);
+      return new S3DownloadStreamWrapper(inputStream, indexName);
+    } else {
+      return inputStream;
+    }
+  }
+
   @Override
   public void uploadPointState(
       String service, String indexIdentifier, NrtPointState nrtPointState, byte[] data)
@@ -521,7 +543,10 @@ public InputStreamWithTimestamp downloadPointState(
       return null;
     }
     return new InputStreamWithTimestamp(
-        downloadFromS3Path(serviceBucket, backendKeyWithTimestamp.backendKey(), true),
+        wrapDownloadStream(
+            downloadFromS3Path(serviceBucket, backendKeyWithTimestamp.backendKey(), true),
+            s3Metrics,
+            indexIdentifier),
         backendKeyWithTimestamp.timestamp());
   }
 
@@ -654,10 +679,14 @@ private void uploadResource(
     setCurrentResource(prefix, fileName);
   }
 
-  private InputStream downloadResource(String prefix) throws IOException {
+  private InputStream downloadResource(String prefix, String indexIdentifier) throws IOException {
     String fileName = getCurrentResourceName(prefix);
     String backendKey = prefix + fileName;
-    return downloadFromS3Path(serviceBucket, backendKey, true);
+    InputStream inputStream = downloadFromS3Path(serviceBucket, backendKey, true);
+    if (indexIdentifier != null) {
+      inputStream = wrapDownloadStream(inputStream, s3Metrics, indexIdentifier);
+    }
+    return inputStream;
   }
 
   @VisibleForTesting

src/main/java/com/yelp/nrtsearch/tools/nrt_utils/backup/CleanupSnapshotsCommand.java

@@ -123,7 +123,7 @@ public Integer call() throws Exception {
       s3Client =
           StateCommandUtils.createS3Client(bucketName, region, credsFile, credsProfile, maxRetry);
     }
-    S3Backend s3Backend = new S3Backend(bucketName, false, s3Client);
+    S3Backend s3Backend = new S3Backend(bucketName, false, false, s3Client);
 
     long deleteAfterMs = BackupCommandUtils.getTimeIntervalMs(deleteAfter);
     long minTimestampMs = System.currentTimeMillis() - deleteAfterMs;

src/main/java/com/yelp/nrtsearch/tools/nrt_utils/backup/RestoreCommand.java

@@ -147,7 +147,7 @@ public Integer call() throws Exception {
       s3Client =
           StateCommandUtils.createS3Client(bucketName, region, credsFile, credsProfile, maxRetry);
     }
-    S3Backend s3Backend = new S3Backend(bucketName, false, s3Client);
+    S3Backend s3Backend = new S3Backend(bucketName, false, false, s3Client);
 
     String resolvedSnapshotRoot =
         BackupCommandUtils.getSnapshotRoot(snapshotRoot, snapshotServiceName);

src/main/java/com/yelp/nrtsearch/tools/nrt_utils/backup/SnapshotCommand.java

@@ -121,7 +121,7 @@ public Integer call() throws Exception {
       s3Client =
           StateCommandUtils.createS3Client(bucketName, region, credsFile, credsProfile, maxRetry);
     }
-    S3Backend s3Backend = new S3Backend(bucketName, false, s3Client);
+    S3Backend s3Backend = new S3Backend(bucketName, false, false, s3Client);
 
     String resolvedIndexResource =
         StateCommandUtils.getResourceName(s3Backend, serviceName, indexName, exactResourceName);

src/main/java/com/yelp/nrtsearch/tools/nrt_utils/cleanup/CleanupDataCommand.java

@@ -128,7 +128,7 @@ public Integer call() throws Exception {
       s3Client =
           StateCommandUtils.createS3Client(bucketName, region, credsFile, credsProfile, maxRetry);
     }
-    S3Backend s3Backend = new S3Backend(bucketName, false, s3Client);
+    S3Backend s3Backend = new S3Backend(bucketName, false, false, s3Client);
 
     String resolvedIndexResource =
         StateCommandUtils.getResourceName(s3Backend, serviceName, indexName, exactResourceName);

src/main/java/com/yelp/nrtsearch/tools/nrt_utils/state/GetRemoteStateCommand.java

@@ -96,7 +96,7 @@ public Integer call() throws Exception {
       s3Client =
           StateCommandUtils.createS3Client(bucketName, region, credsFile, credsProfile, maxRetry);
     }
-    S3Backend s3Backend = new S3Backend(bucketName, false, s3Client);
+    S3Backend s3Backend = new S3Backend(bucketName, false, false, s3Client);
     String resolvedResourceName =
         StateCommandUtils.getResourceName(s3Backend, serviceName, resourceName, exactResourceName);
 

src/main/java/com/yelp/nrtsearch/tools/nrt_utils/state/GetResourceVersionCommand.java

@@ -99,7 +99,7 @@ public Integer call() throws Exception {
       s3Client =
           StateCommandUtils.createS3Client(bucketName, region, credsFile, credsProfile, maxRetry);
     }
-    S3Backend s3Backend = new S3Backend(bucketName, false, s3Client);
+    S3Backend s3Backend = new S3Backend(bucketName, false, false, s3Client);
     String resolvedResourceName =
         StateCommandUtils.getResourceName(s3Backend, serviceName, resourceName, exactResourceName);