public void transferAllStateDataToDirectory(
            Collection<StateHandleDownloadSpec> downloadRequests,
            CloseableRegistry closeableRegistry)
            throws Exception {

        // We use this closer for fine-grained shutdown of all parallel downloading.
        CloseableRegistry internalCloser = new CloseableRegistry();
        // Make sure we also react to external close signals.
        closeableRegistry.registerCloseable(internalCloser);
        try {
            // We have to wait for all futures to be completed, to make sure in
            // case of failure that we will clean up all the files
            FutureUtils.completeAll(
                            createDownloadRunnables(downloadRequests, internalCloser).stream()
                                    .map(
                                            runnable ->
                                                    CompletableFuture.runAsync(
                                                            runnable,
                                                            transfer.getExecutorService()))
                                    .collect(Collectors.toList()))
                    .get();
        } catch (Exception e) {
            downloadRequests.stream()
                    .map(StateHandleDownloadSpec::getDownloadDestination)
                    .map(Path::toFile)
                    .forEach(FileUtils::deleteDirectoryQuietly);
            // Error reporting
            Throwable throwable = ExceptionUtils.stripExecutionException(e);
            throwable = ExceptionUtils.stripException(throwable, RuntimeException.class);
            if (throwable instanceof IOException) {
                throw (IOException) throwable;
            } else {
                throw new FlinkRuntimeException("Failed to download data for state handles.", e);
            }
        } finally {
            // Unregister and close the internal closer.
            if (closeableRegistry.unregisterCloseable(internalCloser)) {
                IOUtils.closeQuietly(internalCloser);
            }
        }
    }
@Test
    public void testMultiThreadRestoreCorrectly() throws Exception {
        int numRemoteHandles = 3;
        int numSubHandles = 6;
        byte[][][] contents = createContents(numRemoteHandles, numSubHandles);
        List<StateHandleDownloadSpec> downloadRequests = new ArrayList<>(numRemoteHandles);
        for (int i = 0; i < numRemoteHandles; ++i) {
            downloadRequests.add(
                    createDownloadRequestForContent(
                            temporaryFolder.newFolder().toPath(), contents[i], i));
        }

        try (RocksDBStateDownloader rocksDBStateDownloader = new RocksDBStateDownloader(4)) {
            rocksDBStateDownloader.transferAllStateDataToDirectory(
                    downloadRequests, new CloseableRegistry());
        }

        for (int i = 0; i < numRemoteHandles; ++i) {
            StateHandleDownloadSpec downloadRequest = downloadRequests.get(i);
            Path dstPath = downloadRequest.getDownloadDestination();
            Assert.assertTrue(dstPath.toFile().exists());
            for (int j = 0; j < numSubHandles; ++j) {
                assertStateContentEqual(
                        contents[i][j], dstPath.resolve(String.format("sharedState-%d-%d", i, j)));
            }
        }
    }