/** Data transfer base class for {@link RocksDBKeyedStateBackend}. */
class RocksDBStateDataTransfer implements Closeable {

    protected final ExecutorService executorService;

    /** ExecutorService to run data transfer operations that can be retried on exceptions. */
    protected final ScheduledExecutorService retryExecutorService;

    private final int retryTimes;

    private final Duration retryInterval;

    RocksDBStateDataTransfer(int threadNum, int retryTimes, Duration retryInterval) {
        if (threadNum > 1) {
            executorService =
                    Executors.newFixedThreadPool(
                            threadNum, new ExecutorThreadFactory("Flink-RocksDBStateDataTransfer"));
        } else {
            executorService = newDirectExecutorService();
        }

        this.retryTimes = retryTimes;
        this.retryInterval = retryInterval;

        retryExecutorService =
                Executors.newSingleThreadScheduledExecutor(
                        new ExecutorThreadFactory("Flink-RocksDBStateDataTransfer-Retry"));
    }

    @Override
    public void close() {
        executorService.shutdownNow();
        retryExecutorService.shutdownNow();
    }

    protected <T> CompletableFuture<T> retry(CompletableFuture<T> dataTransferOperation) {
        if (this.retryTimes > 0) {
            return FutureUtils.retryWithDelay(
                    () -> dataTransferOperation,
                    new FixedRetryStrategy(retryTimes, retryInterval),
                    throwable -> true,
                    new ScheduledExecutorServiceAdapter(retryExecutorService));
        } else {
            // No retry will be performed.
            return dataTransferOperation;
        }
    }
}