Implement leader rate limiting for file restore (#37677)

This is related to #35975. This commit implements rate limiting on the
leader side using the CombinedRateLimiter.

Author: Tim-Brooks

x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/repository/CcrRestoreSourceService.java

@@ -16,8 +16,10 @@
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.component.AbstractLifecycleComponent;
 import org.elasticsearch.common.lease.Releasable;
+import org.elasticsearch.common.metrics.CounterMetric;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.TimeValue;
+import org.elasticsearch.common.util.CombinedRateLimiter;
 import org.elasticsearch.common.util.concurrent.AbstractRefCounted;
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.common.util.concurrent.KeyedLock;
@@ -42,6 +44,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.function.Consumer;
+import java.util.function.LongConsumer;
 
 public class CcrRestoreSourceService extends AbstractLifecycleComponent implements IndexEventListener {
 
@@ -52,6 +55,7 @@ public class CcrRestoreSourceService extends AbstractLifecycleComponent implemen
     private final CopyOnWriteArrayList<Consumer<String>> closeSessionListeners = new CopyOnWriteArrayList<>();
     private final ThreadPool threadPool;
     private final CcrSettings ccrSettings;
+    private final CounterMetric throttleTime = new CounterMetric();
 
     public CcrRestoreSourceService(ThreadPool threadPool, CcrSettings ccrSettings) {
         this.threadPool = threadPool;
@@ -136,7 +140,7 @@ public synchronized SessionReader getSessionReader(String sessionUUID) {
             throw new IllegalArgumentException("session [" + sessionUUID + "] not found");
         }
         restore.idle = false;
-        return new SessionReader(restore);
+        return new SessionReader(restore, ccrSettings, throttleTime::inc);
     }
 
     private void internalCloseSession(String sessionUUID, boolean throwIfSessionMissing) {
@@ -182,6 +186,10 @@ private void maybeTimeout(String sessionUUID) {
         }
     }
 
+    public long getThrottleTime() {
+        return this.throttleTime.count();
+    }
+
     private static class RestoreSession extends AbstractRefCounted {
 
         private final String sessionUUID;
@@ -254,9 +262,13 @@ protected void closeInternal() {
     public static class SessionReader implements Closeable {
 
         private final RestoreSession restoreSession;
+        private final CcrSettings ccrSettings;
+        private final LongConsumer throttleListener;
 
-        private SessionReader(RestoreSession restoreSession) {
+        private SessionReader(RestoreSession restoreSession, CcrSettings ccrSettings, LongConsumer throttleListener) {
             this.restoreSession = restoreSession;
+            this.ccrSettings = ccrSettings;
+            this.throttleListener = throttleListener;
             restoreSession.incRef();
         }
 
@@ -270,6 +282,9 @@ private SessionReader(RestoreSession restoreSession) {
          * @throws IOException if the read fails
          */
         public long readFileBytes(String fileName, BytesReference reference) throws IOException {
+            CombinedRateLimiter rateLimiter = ccrSettings.getRateLimiter();
+            long throttleTime = rateLimiter.maybePause(reference.length());
+            throttleListener.accept(throttleTime);
             return restoreSession.readFileBytes(fileName, reference);
         }
 

x-pack/plugin/ccr/src/test/java/org/elasticsearch/xpack/ccr/CcrRepositoryIT.java

@@ -239,9 +239,15 @@ public void testDocsAreRecovered() throws Exception {
     }
 
     public void testRateLimitingIsEmployed() throws Exception {
+        boolean followerRateLimiting = randomBoolean();
+
         ClusterUpdateSettingsRequest settingsRequest = new ClusterUpdateSettingsRequest();
         settingsRequest.persistentSettings(Settings.builder().put(CcrSettings.RECOVERY_MAX_BYTES_PER_SECOND.getKey(), "10K"));
-        assertAcked(followerClient().admin().cluster().updateSettings(settingsRequest).actionGet());
+        if (followerRateLimiting) {
+            assertAcked(followerClient().admin().cluster().updateSettings(settingsRequest).actionGet());
+        } else {
+            assertAcked(leaderClient().admin().cluster().updateSettings(settingsRequest).actionGet());
+        }
 
         String leaderClusterRepoName = CcrRepository.NAME_PREFIX + "leader_cluster";
         String leaderIndex = "index1";
@@ -257,11 +263,15 @@ public void testRateLimitingIsEmployed() throws Exception {
         final ClusterService clusterService = getFollowerCluster().getCurrentMasterNodeInstance(ClusterService.class);
 
         List<CcrRepository> repositories = new ArrayList<>();
+        List<CcrRestoreSourceService> restoreSources = new ArrayList<>();
 
         for (RepositoriesService repositoriesService : getFollowerCluster().getDataOrMasterNodeInstances(RepositoriesService.class)) {
             Repository repository = repositoriesService.repository(leaderClusterRepoName);
             repositories.add((CcrRepository) repository);
         }
+        for (CcrRestoreSourceService restoreSource : getLeaderCluster().getDataOrMasterNodeInstances(CcrRestoreSourceService.class)) {
+            restoreSources.add(restoreSource);
+        }
 
         logger.info("--> indexing some data");
         for (int i = 0; i < 100; i++) {
@@ -284,12 +294,20 @@ public void testRateLimitingIsEmployed() throws Exception {
         restoreService.restoreSnapshot(restoreRequest, waitForRestore(clusterService, future));
         future.actionGet();
 
-        assertTrue(repositories.stream().anyMatch(cr -> cr.getRestoreThrottleTimeInNanos() > 0));
+        if (followerRateLimiting) {
+            assertTrue(repositories.stream().anyMatch(cr -> cr.getRestoreThrottleTimeInNanos() > 0));
+        } else {
+            assertTrue(restoreSources.stream().anyMatch(cr -> cr.getThrottleTime() > 0));
+        }
 
         settingsRequest = new ClusterUpdateSettingsRequest();
         ByteSizeValue defaultValue = CcrSettings.RECOVERY_MAX_BYTES_PER_SECOND.getDefault(Settings.EMPTY);
         settingsRequest.persistentSettings(Settings.builder().put(CcrSettings.RECOVERY_MAX_BYTES_PER_SECOND.getKey(), defaultValue));
-        assertAcked(followerClient().admin().cluster().updateSettings(settingsRequest).actionGet());
+        if (followerRateLimiting) {
+            assertAcked(followerClient().admin().cluster().updateSettings(settingsRequest).actionGet());
+        } else {
+            assertAcked(leaderClient().admin().cluster().updateSettings(settingsRequest).actionGet());
+        }
     }
 
     @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/37887")