Skip to content

Commit

Permalink
Implement follower rate limiting for file restore
Browse files Browse the repository at this point in the history
This is related to elastic#35975. This commit implements rate limiting on the
follower side using the `RateLimitingInputStream`.
  • Loading branch information
Tim-Brooks committed Jan 15, 2019
1 parent 36a3b84 commit adae8e8
Showing 1 changed file with 19 additions and 13 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
package org.elasticsearch.xpack.ccr.repository;

import org.apache.lucene.index.IndexCommit;
import org.apache.lucene.store.RateLimiter;
import org.elasticsearch.Version;
import org.elasticsearch.action.admin.cluster.state.ClusterStateRequest;
import org.elasticsearch.action.admin.cluster.state.ClusterStateResponse;
Expand All @@ -25,6 +26,7 @@
import org.elasticsearch.common.component.AbstractLifecycleComponent;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.unit.ByteSizeUnit;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.index.Index;
import org.elasticsearch.index.IndexSettings;
Expand All @@ -35,6 +37,7 @@
import org.elasticsearch.index.snapshots.IndexShardRestoreFailedException;
import org.elasticsearch.index.snapshots.IndexShardSnapshotStatus;
import org.elasticsearch.index.snapshots.blobstore.BlobStoreIndexShardSnapshot;
import org.elasticsearch.index.snapshots.blobstore.RateLimitingInputStream;
import org.elasticsearch.index.snapshots.blobstore.SnapshotFiles;
import org.elasticsearch.index.store.Store;
import org.elasticsearch.index.store.StoreFileMetaData;
Expand Down Expand Up @@ -82,6 +85,7 @@ public class CcrRepository extends AbstractLifecycleComponent implements Reposit
private final String remoteClusterAlias;
private final Client client;
private final CcrLicenseChecker ccrLicenseChecker;
private final RateLimiter.SimpleRateLimiter rateLimiter;

public CcrRepository(RepositoryMetaData metadata, Client client, CcrLicenseChecker ccrLicenseChecker, Settings settings) {
super(settings);
Expand All @@ -90,6 +94,7 @@ public CcrRepository(RepositoryMetaData metadata, Client client, CcrLicenseCheck
this.remoteClusterAlias = Strings.split(metadata.name(), NAME_PREFIX)[1];
this.ccrLicenseChecker = ccrLicenseChecker;
this.client = client;
this.rateLimiter = new RateLimiter.SimpleRateLimiter(new ByteSizeValue(40, ByteSizeUnit.MB).getMbFrac());
}

@Override
Expand Down Expand Up @@ -258,7 +263,7 @@ public void restoreShard(IndexShard indexShard, SnapshotId snapshotId, Version v
// TODO: There should be some local timeout. And if the remote cluster returns an unknown session
// response, we should be able to retry by creating a new session.
String name = metadata.name();
try (RestoreSession restoreSession = RestoreSession.openSession(name, remoteClient, leaderShardId, indexShard, recoveryState)) {
try (RestoreSession restoreSession = openSession(name, remoteClient, leaderShardId, indexShard, recoveryState)) {
restoreSession.restoreFiles();
} catch (Exception e) {
throw new IndexShardRestoreFailedException(indexShard.shardId(), "failed to restore snapshot [" + snapshotId + "]", e);
Expand Down Expand Up @@ -286,7 +291,16 @@ private void maybeUpdateMappings(Client localClient, Client remoteClient, Index
}
}

private static class RestoreSession extends FileRestoreContext implements Closeable {
private RestoreSession openSession(String repositoryName, Client remoteClient, ShardId leaderShardId, IndexShard indexShard,
RecoveryState recoveryState) {
String sessionUUID = UUIDs.randomBase64UUID();
PutCcrRestoreSessionAction.PutCcrRestoreSessionResponse response = remoteClient.execute(PutCcrRestoreSessionAction.INSTANCE,
new PutCcrRestoreSessionRequest(sessionUUID, leaderShardId)).actionGet();
return new RestoreSession(repositoryName, remoteClient, sessionUUID, response.getNode(), indexShard, recoveryState,
response.getStoreFileMetaData());
}

private class RestoreSession extends FileRestoreContext implements Closeable {

private static final int BUFFER_SIZE = 1 << 16;

Expand All @@ -304,15 +318,6 @@ private static class RestoreSession extends FileRestoreContext implements Closea
this.sourceMetaData = sourceMetaData;
}

static RestoreSession openSession(String repositoryName, Client remoteClient, ShardId leaderShardId, IndexShard indexShard,
RecoveryState recoveryState) {
String sessionUUID = UUIDs.randomBase64UUID();
PutCcrRestoreSessionAction.PutCcrRestoreSessionResponse response = remoteClient.execute(PutCcrRestoreSessionAction.INSTANCE,
new PutCcrRestoreSessionRequest(sessionUUID, leaderShardId)).actionGet();
return new RestoreSession(repositoryName, remoteClient, sessionUUID, response.getNode(), indexShard, recoveryState,
response.getStoreFileMetaData());
}

void restoreFiles() throws IOException {
ArrayList<BlobStoreIndexShardSnapshot.FileInfo> fileInfos = new ArrayList<>();
for (StoreFileMetaData fileMetaData : sourceMetaData) {
Expand All @@ -325,7 +330,8 @@ void restoreFiles() throws IOException {

@Override
protected InputStream fileInputStream(BlobStoreIndexShardSnapshot.FileInfo fileInfo) {
return new RestoreFileInputStream(remoteClient, sessionUUID, node, fileInfo.metadata());
RestoreFileInputStream restoreInputStream = new RestoreFileInputStream(remoteClient, sessionUUID, node, fileInfo.metadata());
return new RateLimitingInputStream(restoreInputStream, rateLimiter, (n) -> {});
}

@Override
Expand All @@ -336,7 +342,7 @@ public void close() {
}
}

private static class RestoreFileInputStream extends InputStream {
private class RestoreFileInputStream extends InputStream {

private final Client remoteClient;
private final String sessionUUID;
Expand Down

0 comments on commit adae8e8

Please sign in to comment.