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

Example 1 with SlicedInputStream

use of org.opensearch.index.snapshots.blobstore.SlicedInputStream in project OpenSearch by opensearch-project.

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(), null);
        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
                    final int workers = Math.min(threadPool.info(ThreadPool.Names.SNAPSHOT).getMax(), 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) {
                        try {
                            executeOneFileRestore(files, allFilesListener);
                        } catch (Exception e) {
                            allFilesListener.onFailure(e);
                        }
                    }
                }
            }

            private void executeOneFileRestore(BlockingQueue<BlobStoreIndexShardSnapshot.FileInfo> files, ActionListener<Void> allFilesListener) throws InterruptedException {
                final BlobStoreIndexShardSnapshot.FileInfo fileToRecover = files.poll(0L, TimeUnit.MILLISECONDS);
                if (fileToRecover == null) {
                    allFilesListener.onResponse(null);
                } else {
                    executor.execute(ActionRunnable.wrap(allFilesListener, filesListener -> {
                        store.incRef();
                        try {
                            restoreFile(fileToRecover, store);
                        } finally {
                            store.decRef();
                        }
                        executeOneFileRestore(files, filesListener);
                    }));
                }
            }

            private void restoreFile(BlobStoreIndexShardSnapshot.FileInfo fileInfo, Store store) throws IOException {
                ensureNotClosing(store);
                logger.trace(() -> new ParameterizedMessage("[{}] restoring [{}] to [{}]", metadata.name(), fileInfo, store));
                boolean success = false;
                try (IndexOutput indexOutput = store.createVerifyingOutput(fileInfo.physicalName(), fileInfo.metadata(), IOContext.DEFAULT)) {
                    if (fileInfo.name().startsWith(VIRTUAL_DATA_BLOB_PREFIX)) {
                        final BytesRef hash = fileInfo.metadata().hash();
                        indexOutput.writeBytes(hash.bytes, hash.offset, hash.length);
                        recoveryState.getIndex().addRecoveredBytesToFile(fileInfo.physicalName(), hash.length);
                    } else {
                        try (InputStream stream = maybeRateLimitRestores(new SlicedInputStream(fileInfo.numberOfParts()) {

                            @Override
                            protected InputStream openSlice(int slice) throws IOException {
                                ensureNotClosing(store);
                                return container.readBlob(fileInfo.partName(slice));
                            }
                        })) {
                            final byte[] buffer = new byte[Math.toIntExact(Math.min(bufferSize, fileInfo.length()))];
                            int length;
                            while ((length = stream.read(buffer)) > 0) {
                                ensureNotClosing(store);
                                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());
                    }
                }
            }

            void ensureNotClosing(final Store store) throws AlreadyClosedException {
                assert store.refCount() > 0;
                if (store.isClosing()) {
                    throw new AlreadyClosedException("store is closing");
                }
            }
        }.restore(snapshotFiles, store, l);
    }));
}