com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures Maven / Gradle / Ivy
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.datastax.oss.driver.internal.core.util.concurrent;
import com.datastax.oss.driver.api.core.DriverException;
import com.datastax.oss.driver.api.core.DriverExecutionException;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import com.datastax.oss.driver.shaded.guava.common.base.Throwables;
import java.util.List;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;
public class CompletableFutures {
public static CompletableFuture failedFuture(Throwable cause) {
CompletableFuture future = new CompletableFuture<>();
future.completeExceptionally(cause);
return future;
}
/** Completes {@code target} with the outcome of {@code source}. */
public static void completeFrom(CompletionStage source, CompletableFuture target) {
source.whenComplete(
(t, error) -> {
if (error != null) {
target.completeExceptionally(error);
} else {
target.complete(t);
}
});
}
/** @return a completion stage that completes when all inputs are done (success or failure). */
public static CompletionStage allDone(List> inputs) {
CompletableFuture result = new CompletableFuture<>();
if (inputs.isEmpty()) {
result.complete(null);
} else {
final int todo = inputs.size();
final AtomicInteger done = new AtomicInteger();
for (CompletionStage input : inputs) {
input.whenComplete(
(v, error) -> {
if (done.incrementAndGet() == todo) {
result.complete(null);
}
});
}
}
return result;
}
/** Do something when all inputs are done (success or failure). */
public static void whenAllDone(
List> inputs, Runnable callback, Executor executor) {
allDone(inputs).thenRunAsync(callback, executor).exceptionally(UncaughtExceptions::log);
}
/**
* @return a completion stage that completes when all inputs are successful, or fails if any of
* them failed.
*/
public static CompletionStage allSuccessful(List> inputs) {
CompletableFuture result = new CompletableFuture<>();
if (inputs.isEmpty()) {
result.complete(null);
} else {
final int todo = inputs.size();
final AtomicInteger done = new AtomicInteger();
final CopyOnWriteArrayList errors = new CopyOnWriteArrayList<>();
for (CompletionStage input : inputs) {
input.whenComplete(
(v, error) -> {
if (error != null) {
errors.add(error);
}
if (done.incrementAndGet() == todo) {
if (errors.isEmpty()) {
result.complete(null);
} else {
Throwable finalError = errors.get(0);
for (int i = 1; i < errors.size(); i++) {
finalError.addSuppressed(errors.get(i));
}
result.completeExceptionally(finalError);
}
}
});
}
}
return result;
}
/** Get the result now, when we know for sure that the future is complete. */
public static T getCompleted(CompletionStage stage) {
CompletableFuture future = stage.toCompletableFuture();
Preconditions.checkArgument(future.isDone() && !future.isCompletedExceptionally());
try {
return future.get();
} catch (InterruptedException | ExecutionException e) {
// Neither can happen given the precondition
throw new AssertionError("Unexpected error", e);
}
}
/** Get the error now, when we know for sure that the future is failed. */
public static Throwable getFailed(CompletionStage stage) {
CompletableFuture future = stage.toCompletableFuture();
Preconditions.checkArgument(future.isCompletedExceptionally());
try {
future.get();
throw new AssertionError("future should be failed");
} catch (InterruptedException e) {
throw new AssertionError("Unexpected error", e);
} catch (ExecutionException e) {
return e.getCause();
}
}
public static T getUninterruptibly(CompletionStage stage) {
boolean interrupted = false;
try {
while (true) {
try {
return stage.toCompletableFuture().get();
} catch (InterruptedException e) {
interrupted = true;
} catch (ExecutionException e) {
Throwable cause = e.getCause();
if (cause instanceof DriverException) {
throw ((DriverException) cause).copy();
}
Throwables.throwIfUnchecked(cause);
throw new DriverExecutionException(cause);
}
}
} finally {
if (interrupted) {
Thread.currentThread().interrupt();
}
}
}
/**
* Executes a function on the calling thread and returns result in a {@link CompletableFuture}.
*
* Similar to {@link CompletableFuture#completedFuture} except takes a {@link Supplier} and if
* the supplier throws an unchecked exception, the returning future fails with that exception.
*
* @param supplier Function to execute
* @param Type of result
* @return result of function wrapped in future
*/
public static CompletableFuture wrap(Supplier supplier) {
try {
return CompletableFuture.completedFuture(supplier.get());
} catch (Throwable t) {
return failedFuture(t);
}
}
public static void whenCancelled(CompletionStage stage, Runnable action) {
stage.exceptionally(
(error) -> {
if (error instanceof CancellationException) {
action.run();
}
return null;
});
}
public static void propagateCancellation(CompletionStage source, CompletionStage target) {
whenCancelled(source, () -> target.toCompletableFuture().cancel(true));
}
}