Auto batching

[treetime]

TL;DR: For AsyncLoaderCache Is there some mechanism by which concurrent invocations of the loader can be limited, so that incoming incoming .get's that miss get batched and summoned together using loadAll?

---

Good day

There is a large existing system with a custom caches, which is being replaced with Caffeine. This cache is called concurrently by many callers, each needing singular values, and given the established nature of the system this is hard to change. The data fetching operation to the DB is expensive, and it would be convenient to say that only ever eg. two loading operations can happen at the same time, with singular requests that arrive while they are in flight (that need fetching) instead being fetched using the loadAll method.

I can see how one could achieve this by building around Caffeine, but is there some way to achieve it inside it?

Thanks for you help.

[ben-manes]

You can look at the contributed coalescing bulkloader. I believe this returns the singular futures, queues the requests, batches loads them, and completes the pending futures. This type of functionality is nice but a little out of scope for this library, similar to retry and fallbacks being served by combining with Failsafe. Unfortunately I don’t know of a generic coalescing library, but that would compose nicely with our cache loader.

[treetime]

Thank you.

[ben-manes]

fwiw, I suspect using reactor’s bufferTimeout and similar might be useful, but I’ve never used that library. The contrib code is simple, but that might be less code if your team is already familiar with it.

[ben-manes]

I thought this would be a fun excuse to try out a reactive library; see the sample below. It does seem a bit over engineered which makes it confusing to find the right configuration, but otherwise nicely fits this use-case. The core setup was a lean 30 LOC, as I spent about the same trying to make it fluent through a builder. Hope that helps.

CoalescingBulkLoader

import static com.google.common.truth.Truth.assertThat;
import static java.util.Objects.requireNonNull;
import static java.util.stream.Collectors.toSet;

import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.ThreadLocalRandom;
import java.util.function.Function;

import org.testng.annotations.Test;

import com.github.benmanes.caffeine.cache.AsyncCacheLoader;
import com.github.benmanes.caffeine.cache.AsyncLoadingCache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.google.common.base.Stopwatch;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.Uninterruptibles;

import reactor.core.publisher.Sinks;
import reactor.core.publisher.Sinks.Many;
import reactor.core.scheduler.Schedulers;

public final class Reactive {

  @Test
  public void reactor() {
    var stopwatch = Stopwatch.createUnstarted();
    AsyncLoadingCache<Integer, Integer> cache = Caffeine.newBuilder()
        .buildAsync(CoalescingBulkLoader.builder()
            .maxTime(Duration.ofMillis(50))
            .parallelism(5)
            .maxSize(5)
            .build(keys -> {
              System.out.printf("Thread #%d: Loading %d keys at time %s%n",
                  Thread.currentThread().getId(), keys.size(), stopwatch);
              return Maps.toMap(keys, key -> -key);
            }));

    var results = new HashMap<Integer, CompletableFuture<Integer>>();
    stopwatch.start();
    for (int i = 0; i < 82; i++) {
      Uninterruptibles.sleepUninterruptibly(Duration.ofMillis(
          ThreadLocalRandom.current().nextInt(0, 30)));
      results.put(i, cache.get(i));
    }
    results.forEach((key, future) -> {
      assertThat(future.join()).isEqualTo(-key);
    });
  }

  static final class CoalescingBulkLoader<K, V> implements AsyncCacheLoader<K, V> {
    private final Function<Set<K>, Map<K, V>> mappingFunction;
    private final Many<Request<K, V>> sink;

    private CoalescingBulkLoader(Builder builder, Function<Set<K>, Map<K, V>> mappingFunction) {
      this.mappingFunction = requireNonNull(mappingFunction);
      sink = Sinks.many().unicast().onBackpressureBuffer();
      sink.asFlux()
          .bufferTimeout(builder.maxSize, builder.maxTime)
          .parallel(builder.parallelism)
          .runOn(Schedulers.boundedElastic())
          .subscribe(this::handle);
    }

    @Override public CompletableFuture<V> asyncLoad(K key, Executor executor) {
      var result = new CompletableFuture<V>();
      sink.tryEmitNext(new Request<>(key, result)).orThrow();
      return result;
    }

    private void handle(List<Request<K, V>> requests) {
      try {
        var results = mappingFunction.apply(requests.stream().map(Request::key).collect(toSet()));
        requests.forEach(request -> request.result.complete(results.get(request.key())));
      } catch (Throwable t) {
        requests.forEach(request -> request.result.completeExceptionally(t));
      }
    }

    public static Builder builder() {
      return new Builder();
    }

    public static final class Builder {
      private Duration maxTime = Duration.ofNanos(Long.MAX_VALUE);
      private int maxSize = Integer.MAX_VALUE;
      private int parallelism = 1;

      public Builder maxTime(Duration maxTime) {
        this.maxTime = maxTime;
        return this;
      }

      public Builder maxSize(int maxSize) {
        this.maxSize = maxSize;
        return this;
      }

      public Builder parallelism(int parallelism) {
        this.parallelism = parallelism;
        return this;
      }

      public <K, V> CoalescingBulkLoader<K, V> build(Function<Set<K>, Map<K, V>> mappingFunction) {
        return new CoalescingBulkLoader<>(this, mappingFunction);
      }
    }

    private record Request<K, V>(K key, CompletableFuture<V> result) {}
  }
}

Output with random delays

Thread #16: Loading 3 keys at time 86.87 ms
Thread #17: Loading 5 keys at time 126.0 ms
Thread #18: Loading 4 keys at time 183.0 ms
Thread #19: Loading 3 keys at time 243.0 ms
Thread #20: Loading 5 keys at time 286.8 ms
Thread #16: Loading 4 keys at time 363.3 ms
Thread #17: Loading 4 keys at time 422.6 ms
Thread #18: Loading 5 keys at time 477.5 ms
Thread #19: Loading 5 keys at time 550.5 ms
Thread #20: Loading 3 keys at time 613.6 ms
Thread #16: Loading 2 keys at time 663.8 ms
Thread #17: Loading 5 keys at time 726.0 ms
Thread #18: Loading 4 keys at time 809.9 ms
Thread #19: Loading 3 keys at time 870.4 ms
Thread #20: Loading 3 keys at time 925.0 ms
Thread #16: Loading 4 keys at time 994.4 ms
Thread #17: Loading 4 keys at time 1.061 s
Thread #18: Loading 4 keys at time 1.126 s
Thread #19: Loading 3 keys at time 1.178 s
Thread #20: Loading 4 keys at time 1.234 s
Thread #16: Loading 4 keys at time 1.300 s
Thread #17: Loading 1 keys at time 1.366 s
PASSED: Reactive.reactor

Output with no delays

Thread #16: Loading 5 keys at time 8.163 ms
Thread #20: Loading 5 keys at time 28.37 ms
Thread #19: Loading 5 keys at time 28.89 ms
Thread #18: Loading 5 keys at time 29.28 ms
Thread #17: Loading 5 keys at time 29.73 ms
Thread #17: Loading 5 keys at time 47.16 ms
Thread #18: Loading 5 keys at time 48.30 ms
Thread #18: Loading 5 keys at time 49.50 ms
Thread #19: Loading 5 keys at time 51.91 ms
Thread #20: Loading 5 keys at time 52.60 ms
Thread #16: Loading 5 keys at time 53.24 ms
Thread #16: Loading 5 keys at time 54.02 ms
Thread #20: Loading 5 keys at time 55.31 ms
Thread #19: Loading 5 keys at time 56.12 ms
Thread #17: Loading 5 keys at time 58.37 ms
Thread #16: Loading 5 keys at time 59.50 ms
Thread #17: Loading 2 keys at time 62.43 ms
PASSED: Reactive.reactor