org.opensearch.migrations.replay.util.TrackedFuture Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of trafficReplayer Show documentation
Show all versions of trafficReplayer Show documentation
Everything opensearch migrations
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);
}
}