FutureUtils { public static <T> CompletableFuture<T> retry( final Supplier<CompletableFuture<T>> operation, final int retries, final Executor executor) { final CompletableFuture<T> resultFuture = new CompletableFuture<>(); retryOperation(resultFuture, operation, retries, executor); return resultFuture; }  static CompletableFuture<T> retry(
			final Supplier<CompletableFuture<T>> operation,
			final int retries,
			final Executor executor); static CompletableFuture<T> retryWithDelay(
			final Supplier<CompletableFuture<T>> operation,
			final int retries,
			final Time retryDelay,
			final ScheduledExecutor scheduledExecutor); static ConjunctFuture<Collection<T>> combineAll(Collection<? extends CompletableFuture<? extends T>> futures); static ConjunctFuture<Void> waitForAll(Collection<? extends CompletableFuture<?>> futures); static CompletableFuture<T> completedExceptionally(Throwable cause); static FiniteDuration toFiniteDuration(Time time); static Time toTime(FiniteDuration finiteDuration); static CompletableFuture<T> toJava(Future<T> scalaFuture); static CompletableFuture<T> orTimeout(CompletableFuture<T> future, long timeout, TimeUnit timeUnit);  }
@Test public void testRetrySuccess() throws Exception { final int retries = 10; final AtomicInteger atomicInteger = new AtomicInteger(0); CompletableFuture<Boolean> retryFuture = FutureUtils.retry( () -> CompletableFuture.supplyAsync( () -> { if (atomicInteger.incrementAndGet() == retries) { return true; } else { throw new CompletionException(new FlinkException("Test exception")); } }, TestingUtils.defaultExecutor()), retries, TestingUtils.defaultExecutor()); assertTrue(retryFuture.get()); assertTrue(retries == atomicInteger.get()); }
@Test public void testRetryCancellation() throws Exception { final int retries = 10; final AtomicInteger atomicInteger = new AtomicInteger(0); final OneShotLatch notificationLatch = new OneShotLatch(); final OneShotLatch waitLatch = new OneShotLatch(); final AtomicReference<Throwable> atomicThrowable = new AtomicReference<>(null); CompletableFuture<?> retryFuture = FutureUtils.retry( () -> CompletableFuture.supplyAsync( () -> { if (atomicInteger.incrementAndGet() == 2) { notificationLatch.trigger(); try { waitLatch.await(); } catch (InterruptedException e) { atomicThrowable.compareAndSet(null, e); } } throw new CompletionException(new FlinkException("Test exception")); }, TestingUtils.defaultExecutor()), retries, TestingUtils.defaultExecutor()); notificationLatch.await(); assertFalse(retryFuture.isDone()); retryFuture.cancel(false); waitLatch.trigger(); assertTrue(retryFuture.isCancelled()); assertEquals(2, atomicInteger.get()); if (atomicThrowable.get() != null) { throw new FlinkException("Exception occurred in the retry operation.", atomicThrowable.get()); } }