Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance. Project price only 1 $
You can buy this project and download/modify it how often you want.
package org.opensearch.migrations.replay.util;
import java.time.Duration;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.IntStream;
import java.util.stream.Stream;
import lombok.Getter;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
/**
* This class wraps CompletableFutures into traceable and identifiable pieces so that when
* dealing with thousands of completable futures at a time, one can clearly understand how
* work is proceeding and why there may be issues.
*
* This is adding a great amount of clarity, though using it may still be a challenge. Much
* more work is expected to improve the UX for developers.
* @param The type of object that will be returned to represent diagnostic information
* @param The type of value of the underlying (internal) CompletableFuture's result
*/
@Slf4j
public class TrackedFuture {
public final CompletableFuture future;
protected AtomicReference> innerComposedPendingCompletableFutureReference;
@Getter
public final Supplier diagnosticSupplier;
protected final AtomicReference> parentDiagnosticFutureRef;
private TrackedFuture() {
throw new IllegalCallerException();
}
/**
* This factory class is here so that subclasses can write their own versions that return objects
* of their own subclass.
*/
public static class Factory {
private Factory() {}
public static TrackedFuture failedFuture(
@NonNull Throwable e,
@NonNull Supplier diagnosticSupplier
) {
return new TrackedFuture<>(CompletableFuture.failedFuture(e), diagnosticSupplier, null);
}
public static TrackedFuture completedFuture(U v, @NonNull Supplier diagnosticSupplier) {
return new TrackedFuture<>(CompletableFuture.completedFuture(v), diagnosticSupplier, null);
}
}
private TrackedFuture(
@NonNull CompletableFuture future,
@NonNull Supplier diagnosticSupplier,
TrackedFuture parentFuture
) {
this.future = future;
this.diagnosticSupplier = diagnosticSupplier;
this.parentDiagnosticFutureRef = new AtomicReference<>();
setParentDiagnosticFuture(parentFuture);
}
public TrackedFuture(@NonNull CompletableFuture future, @NonNull Supplier diagnosticSupplier) {
this(future, diagnosticSupplier, null);
}
public static Throwable unwindPossibleCompletionException(Throwable t) {
while (t instanceof CompletionException) {
t = t.getCause();
}
return t;
}
public TrackedFuture getParentDiagnosticFuture() {
var p = parentDiagnosticFutureRef.get();
if (future.isDone() && p != null) {
p.setParentDiagnosticFuture(null);
}
return p;
}
protected void setParentDiagnosticFuture(TrackedFuture parent) {
if (parent == null) {
parentDiagnosticFutureRef.set(null);
return;
}
var wasSet = parentDiagnosticFutureRef.compareAndSet(null, parent);
if (!wasSet) {
throw new IllegalStateException(
"dependencyDiagnosticFutureRef was already set to " + parentDiagnosticFutureRef.get()
);
}
// the parent is a pretty good breadcrumb for the current stack... but the grandparent of the most recently
// finished ancestor begins to have diminished value immediately, so cut the ancestry tree at this point
future.whenComplete(
(v, t) -> Optional.ofNullable(getParentDiagnosticFuture()).ifPresent(p -> p.setParentDiagnosticFuture(null))
);
}
/**
* @throws IllegalStateException if the dependentFuture has already been passed to this method
* before or if it has already been marked as completed or was initialized with a parent.
*/
public TrackedFuture propagateCompletionToDependentFuture(
TrackedFuture dependentFuture,
BiConsumer, CompletableFuture> consume,
@NonNull Supplier diagnosticSupplier
) {
dependentFuture.setParentDiagnosticFuture(this);
return this.whenComplete((v, t) -> consume.accept(this.future, dependentFuture.future), diagnosticSupplier);
}
public TrackedFuture getInnerComposedPendingCompletableFuture() {
return Optional.ofNullable(innerComposedPendingCompletableFutureReference)
.map(AtomicReference::get)
.orElse(null);
}
public TrackedFuture map(
@NonNull Function, CompletableFuture> fn,
@NonNull Supplier diagnosticSupplier
) {
var newCf = fn.apply(future);
return new TrackedFuture<>(newCf, diagnosticSupplier, this);
}
public TrackedFuture thenAccept(Consumer fn, @NonNull Supplier diagnosticSupplier) {
return this.map(cf -> cf.thenAccept(fn), diagnosticSupplier);
}
public TrackedFuture thenApply(Function fn, @NonNull Supplier diagnosticSupplier) {
return this.map(cf -> cf.thenApply(fn), diagnosticSupplier);
}
public TrackedFuture exceptionally(Function fn, @NonNull Supplier diagnosticSupplier) {
return this.map(cf -> cf.exceptionally(fn), diagnosticSupplier);
}
public TrackedFuture whenComplete(
BiConsumer fn,
@NonNull Supplier diagnosticSupplier
) {
return map(cf -> cf.whenComplete(fn::accept), diagnosticSupplier);
}
public TrackedFuture thenCompose(
@NonNull Function> fn,
@NonNull Supplier diagnosticSupplier
) {
var innerComposedCompletableFutureReference = new AtomicReference>();
var newCf = this.future.thenCompose(v -> {
var innerFuture = fn.apply(v);
innerComposedCompletableFutureReference.set(innerFuture);
return innerFuture.future;
});
var wrappedDiagnosticFuture = new TrackedFuture<>(newCf, diagnosticSupplier, this);
wrappedDiagnosticFuture.innerComposedPendingCompletableFutureReference =
innerComposedCompletableFutureReference;
wrappedDiagnosticFuture.future.whenComplete((v2, t2) -> innerComposedCompletableFutureReference.set(null));
return wrappedDiagnosticFuture;
}
/**
* Run handle(), which can take care of either value completions or Exceptions by sending
* those through the fn argument. This returns the DiagnosticTrackableCompletableFuture
* that the application of fn will yield.
*
* NB/TODO - I can't yet figure out how to enforce type-checking on the incoming fn parameter.
* For example, I can pass a fn as a lambda which returns a String or an integer and the
* compiler doesn't give an error.
* @param fn
* @param diagnosticSupplier
* @return
* @param
*/
public TrackedFuture getDeferredFutureThroughHandle(
@NonNull BiFunction> fn,
@NonNull Supplier diagnosticSupplier
) {
var innerComposedCompletableFutureReference = new AtomicReference>();
CompletableFuture> handledFuture = this.future.handle((v, t) -> {
var innerFuture = fn.apply(v, t);
innerComposedCompletableFutureReference.set(innerFuture);
return innerFuture;
});
var newCf = handledFuture.thenCompose(wcf -> wcf.future);
var wrappedDiagnosticFuture = new TrackedFuture<>(newCf, diagnosticSupplier, this);
wrappedDiagnosticFuture.innerComposedPendingCompletableFutureReference =
innerComposedCompletableFutureReference;
// TODO: Add a count to how many futures have been completed and are falling away?
wrappedDiagnosticFuture.future.whenComplete((v2, t2) -> innerComposedCompletableFutureReference.set(null));
return wrappedDiagnosticFuture;
}
public TrackedFuture handle(
@NonNull BiFunction fn,
@NonNull Supplier diagnosticSupplier
) {
CompletableFuture newCf = this.future.handle(fn);
return new TrackedFuture<>(newCf, diagnosticSupplier, this);
}
public T get() throws ExecutionException, InterruptedException {
return future.get();
}
public T get(@NonNull Duration timeout) throws ExecutionException, InterruptedException, TimeoutException {
var millis = timeout.toMillis() + (timeout.minusNanos(timeout.toNanosPart()).equals(timeout) ? 0 : 1); // round
// up
return future.get(millis, TimeUnit.MILLISECONDS);
}
public boolean isDone() {
return future.isDone();
}
public Stream> walkParentsAsStream() {
AtomicReference> chainHeadReference = new AtomicReference<>(this);
return IntStream.generate(() -> chainHeadReference.get() != null ? 1 : 0).takeWhile(x -> x == 1).mapToObj(i -> {
var trackedFuture = chainHeadReference.get();
chainHeadReference.set(trackedFuture.getParentDiagnosticFuture());
return trackedFuture;
});
}
@Override
public String toString() {
return formatAsString(x -> null);
}
public String formatAsString(@NonNull Function, String> resultFormatter) {
return TrackedFutureStringFormatter.format(this, resultFormatter);
}
public String formatAsJson(@NonNull Function, String> resultFormatter) {
return TrackedFutureJsonFormatter.format(this, resultFormatter);
}
}
