Examples with SlicedInputStream org.elasticsearch.index.snapshots.blobstore.SlicedInputStream used on opensource projects

Example 1 with SlicedInputStream

use of org.elasticsearch.index.snapshots.blobstore.SlicedInputStream in project crate by crate.

the class BlobStoreRepository method restoreShard.

@Override
public void restoreShard(Store store, SnapshotId snapshotId, IndexId indexId, ShardId snapshotShardId, RecoveryState recoveryState, ActionListener<Void> listener) {
    final ShardId shardId = store.shardId();
    final ActionListener<Void> restoreListener = ActionListener.delegateResponse(listener, (l, e) -> l.onFailure(new IndexShardRestoreFailedException(shardId, "failed to restore snapshot [" + snapshotId + "]", e)));
    final Executor executor = threadPool.executor(ThreadPool.Names.SNAPSHOT);
    final BlobContainer container = shardContainer(indexId, snapshotShardId);
    executor.execute(ActionRunnable.wrap(restoreListener, l -> {
        final BlobStoreIndexShardSnapshot snapshot = loadShardSnapshot(container, snapshotId);
        final SnapshotFiles snapshotFiles = new SnapshotFiles(snapshot.snapshot(), snapshot.indexFiles());
        new FileRestoreContext(metadata.name(), shardId, snapshotId, recoveryState) {

            @Override
            protected void restoreFiles(List<BlobStoreIndexShardSnapshot.FileInfo> filesToRecover, Store store, ActionListener<Void> listener) {
                if (filesToRecover.isEmpty()) {
                    listener.onResponse(null);
                } else {
                    // Start as many workers as fit into the snapshot pool at once at the most
                    int maxPoolSize = executor instanceof ThreadPoolExecutor ? ((ThreadPoolExecutor) executor).getMaximumPoolSize() : 1;
                    final int workers = Math.min(maxPoolSize, snapshotFiles.indexFiles().size());
                    final BlockingQueue<BlobStoreIndexShardSnapshot.FileInfo> files = new LinkedBlockingQueue<>(filesToRecover);
                    final ActionListener<Void> allFilesListener = fileQueueListener(files, workers, ActionListener.map(listener, v -> null));
                    // restore the files from the snapshot to the Lucene store
                    for (int i = 0; i < workers; ++i) {
                        executor.execute(ActionRunnable.run(allFilesListener, () -> {
                            store.incRef();
                            try {
                                BlobStoreIndexShardSnapshot.FileInfo fileToRecover;
                                while ((fileToRecover = files.poll(0L, TimeUnit.MILLISECONDS)) != null) {
                                    restoreFile(fileToRecover, store);
                                }
                            } finally {
                                store.decRef();
                            }
                        }));
                    }
                }
            }

            private void restoreFile(BlobStoreIndexShardSnapshot.FileInfo fileInfo, Store store) throws IOException {
                boolean success = false;
                try (InputStream stream = maybeRateLimit(new SlicedInputStream(fileInfo.numberOfParts()) {

                    @Override
                    protected InputStream openSlice(long slice) throws IOException {
                        return container.readBlob(fileInfo.partName(slice));
                    }
                }, restoreRateLimiter, restoreRateLimitingTimeInNanos)) {
                    try (IndexOutput indexOutput = store.createVerifyingOutput(fileInfo.physicalName(), fileInfo.metadata(), IOContext.DEFAULT)) {
                        final byte[] buffer = new byte[BUFFER_SIZE];
                        int length;
                        while ((length = stream.read(buffer)) > 0) {
                            indexOutput.writeBytes(buffer, 0, length);
                            recoveryState.getIndex().addRecoveredBytesToFile(fileInfo.physicalName(), length);
                        }
                        Store.verify(indexOutput);
                        indexOutput.close();
                        store.directory().sync(Collections.singleton(fileInfo.physicalName()));
                        success = true;
                    } catch (CorruptIndexException | IndexFormatTooOldException | IndexFormatTooNewException ex) {
                        try {
                            store.markStoreCorrupted(ex);
                        } catch (IOException e) {
                            LOGGER.warn("store cannot be marked as corrupted", e);
                        }
                        throw ex;
                    } finally {
                        if (success == false) {
                            store.deleteQuiet(fileInfo.physicalName());
                        }
                    }
                }
            }
        }.restore(snapshotFiles, store, l);
    }));
}