Skip to content

Making the Subscribers use a common base class #154

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Making the Subscribers use a common base class
  • Loading branch information
bbakerman committed May 21, 2024
commit 74567fe99051ea6db7329bbe49683424cebcd3ed
167 changes: 88 additions & 79 deletions src/main/java/org/dataloader/DataLoaderHelper.java
Original file line number Diff line number Diff line change
Expand Up @@ -618,24 +618,23 @@ private static <T> DispatchResult<T> emptyDispatchResult() {
return (DispatchResult<T>) EMPTY_DISPATCH_RESULT;
}

private class DataLoaderSubscriber implements Subscriber<V> {
private abstract class DataLoaderSubscriberBase<T> implements Subscriber<T> {

private final CompletableFuture<List<V>> valuesFuture;
private final List<K> keys;
private final List<Object> callContexts;
private final List<CompletableFuture<V>> queuedFutures;
final CompletableFuture<List<V>> valuesFuture;
final List<K> keys;
final List<Object> callContexts;
final List<CompletableFuture<V>> queuedFutures;

private final List<K> clearCacheKeys = new ArrayList<>();
private final List<V> completedValues = new ArrayList<>();
private int idx = 0;
private boolean onErrorCalled = false;
private boolean onCompleteCalled = false;
List<K> clearCacheKeys = new ArrayList<>();
List<V> completedValues = new ArrayList<>();
boolean onErrorCalled = false;
boolean onCompleteCalled = false;

private DataLoaderSubscriber(
CompletableFuture<List<V>> valuesFuture,
List<K> keys,
List<Object> callContexts,
List<CompletableFuture<V>> queuedFutures
DataLoaderSubscriberBase(
CompletableFuture<List<V>> valuesFuture,
List<K> keys,
List<Object> callContexts,
List<CompletableFuture<V>> queuedFutures
) {
this.valuesFuture = valuesFuture;
this.keys = keys;
Expand All @@ -648,40 +647,87 @@ public void onSubscribe(Subscription subscription) {
subscription.request(keys.size());
}

// onNext may be called by multiple threads - for the time being, we pass 'synchronized' to guarantee
// correctness (at the cost of speed).
@Override
public synchronized void onNext(V value) {
public void onNext(T v) {
assertState(!onErrorCalled, () -> "onError has already been called; onNext may not be invoked.");
assertState(!onCompleteCalled, () -> "onComplete has already been called; onNext may not be invoked.");
}

K key = keys.get(idx);
Object callContext = callContexts.get(idx);
CompletableFuture<V> future = queuedFutures.get(idx);
@Override
public void onComplete() {
assertState(!onErrorCalled, () -> "onError has already been called; onComplete may not be invoked.");
onCompleteCalled = true;
}

@Override
public void onError(Throwable throwable) {
assertState(!onCompleteCalled, () -> "onComplete has already been called; onError may not be invoked.");
onErrorCalled = true;

stats.incrementBatchLoadExceptionCount(new IncrementBatchLoadExceptionCountStatisticsContext<>(keys, callContexts));
}

/*
* A value has arrived - how do we complete the future that's associated with it in a common way
*/
void onNextValue(K key, V value, Object callContext, CompletableFuture<V> future) {
if (value instanceof Try) {
// we allow the batch loader to return a Try so we can better represent a computation
// that might have worked or not.
//noinspection unchecked
Try<V> tryValue = (Try<V>) value;
if (tryValue.isSuccess()) {
future.complete(tryValue.get());
} else {
stats.incrementLoadErrorCount(new IncrementLoadErrorCountStatisticsContext<>(key, callContext));
future.completeExceptionally(tryValue.getThrowable());
clearCacheKeys.add(keys.get(idx));
clearCacheKeys.add(key);
}
} else {
future.complete(value);
}
}

Throwable unwrapThrowable(Throwable ex) {
if (ex instanceof CompletionException) {
ex = ex.getCause();
}
return ex;
}
}

private class DataLoaderSubscriber extends DataLoaderSubscriberBase<V> {

private int idx = 0;

private DataLoaderSubscriber(
CompletableFuture<List<V>> valuesFuture,
List<K> keys,
List<Object> callContexts,
List<CompletableFuture<V>> queuedFutures
) {
super(valuesFuture, keys, callContexts, queuedFutures);
}

// onNext may be called by multiple threads - for the time being, we pass 'synchronized' to guarantee
// correctness (at the cost of speed).
@Override
public synchronized void onNext(V value) {
super.onNext(value);

K key = keys.get(idx);
Object callContext = callContexts.get(idx);
CompletableFuture<V> future = queuedFutures.get(idx);
onNextValue(key, value, callContext, future);

completedValues.add(value);
idx++;
}


@Override
public void onComplete() {
assertState(!onErrorCalled, () -> "onError has already been called; onComplete may not be invoked.");
onCompleteCalled = true;

super.onComplete();
assertResultSize(keys, completedValues);

possiblyClearCacheEntriesOnExceptions(clearCacheKeys);
Expand All @@ -690,13 +736,8 @@ public void onComplete() {

@Override
public void onError(Throwable ex) {
assertState(!onCompleteCalled, () -> "onComplete has already been called; onError may not be invoked.");
onErrorCalled = true;

stats.incrementBatchLoadExceptionCount(new IncrementBatchLoadExceptionCountStatisticsContext<>(keys, callContexts));
if (ex instanceof CompletionException) {
ex = ex.getCause();
}
super.onError(ex);
ex = unwrapThrowable(ex);
// Set the remaining keys to the exception.
for (int i = idx; i < queuedFutures.size(); i++) {
K key = keys.get(i);
Expand All @@ -706,32 +747,23 @@ public void onError(Throwable ex) {
dataLoader.clear(key);
}
}

}

private class DataLoaderMapEntrySubscriber implements Subscriber<Map.Entry<K, V>> {
private final CompletableFuture<List<V>> valuesFuture;
private final List<K> keys;
private final List<Object> callContexts;
private final List<CompletableFuture<V>> queuedFutures;
private class DataLoaderMapEntrySubscriber extends DataLoaderSubscriberBase<Map.Entry<K, V>> {

private final Map<K, Object> callContextByKey;
private final Map<K, CompletableFuture<V>> queuedFutureByKey;

private final List<K> clearCacheKeys = new ArrayList<>();
private final Map<K, V> completedValuesByKey = new HashMap<>();
private boolean onErrorCalled = false;
private boolean onCompleteCalled = false;


private DataLoaderMapEntrySubscriber(
CompletableFuture<List<V>> valuesFuture,
List<K> keys,
List<Object> callContexts,
List<CompletableFuture<V>> queuedFutures
CompletableFuture<List<V>> valuesFuture,
List<K> keys,
List<Object> callContexts,
List<CompletableFuture<V>> queuedFutures
) {
this.valuesFuture = valuesFuture;
this.keys = keys;
this.callContexts = callContexts;
this.queuedFutures = queuedFutures;

super(valuesFuture,keys,callContexts,queuedFutures);
this.callContextByKey = new HashMap<>();
this.queuedFutureByKey = new HashMap<>();
for (int idx = 0; idx < queuedFutures.size(); idx++) {
Expand All @@ -743,42 +775,24 @@ private DataLoaderMapEntrySubscriber(
}
}

@Override
public void onSubscribe(Subscription subscription) {
subscription.request(keys.size());
}

@Override
public void onNext(Map.Entry<K, V> entry) {
assertState(!onErrorCalled, () -> "onError has already been called; onNext may not be invoked.");
assertState(!onCompleteCalled, () -> "onComplete has already been called; onNext may not be invoked.");
super.onNext(entry);
K key = entry.getKey();
V value = entry.getValue();

Object callContext = callContextByKey.get(key);
CompletableFuture<V> future = queuedFutureByKey.get(key);
if (value instanceof Try) {
// we allow the batch loader to return a Try so we can better represent a computation
// that might have worked or not.
Try<V> tryValue = (Try<V>) value;
if (tryValue.isSuccess()) {
future.complete(tryValue.get());
} else {
stats.incrementLoadErrorCount(new IncrementLoadErrorCountStatisticsContext<>(key, callContext));
future.completeExceptionally(tryValue.getThrowable());
clearCacheKeys.add(key);
}
} else {
future.complete(value);
}

onNextValue(key, value, callContext, future);

completedValuesByKey.put(key, value);
}

@Override
public void onComplete() {
assertState(!onErrorCalled, () -> "onError has already been called; onComplete may not be invoked.");
onCompleteCalled = true;
super.onComplete();

possiblyClearCacheEntriesOnExceptions(clearCacheKeys);
List<V> values = new ArrayList<>(keys.size());
Expand All @@ -791,13 +805,8 @@ public void onComplete() {

@Override
public void onError(Throwable ex) {
assertState(!onCompleteCalled, () -> "onComplete has already been called; onError may not be invoked.");
onErrorCalled = true;

stats.incrementBatchLoadExceptionCount(new IncrementBatchLoadExceptionCountStatisticsContext<>(keys, callContexts));
if (ex instanceof CompletionException) {
ex = ex.getCause();
}
super.onError(ex);
ex = unwrapThrowable(ex);
// Complete the futures for the remaining keys with the exception.
for (int idx = 0; idx < queuedFutures.size(); idx++) {
K key = keys.get(idx);
Expand Down