Java map-like that lazily and atomically creates values for keys. Supports global shutdown also.

LazyInitMapValues.java

package breeze.collections;

import java.util.ArrayList;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
import java.util.function.Function;

final public class LazyInitMapValues<K, V> {
  final private ConcurrentHashMap<K, Future<V>> objects = new ConcurrentHashMap<>();
  final private Function<K, V> valueCreator;
  final private BiConsumer<K, V> valueDestroyer;
  final private ExecutorService creationExecutor;
  private volatile boolean isShutdown = false;


  /**
   * A map-like object which lazily and atomically creates values for keys.
   *
   * @param valueCreator     Function from map key K to stateful value V. For a given key this
   *                         function is guaranteed to run exactly once and other readers will block
   *                         until the value is produced. This function should be prepared to receive
   *                         thread interruptions and cleanup after itself if the map is shutting down
   *                         while the creation function runs.
   * @param valueDestroyer   Destroys a stateful value V given K and V. Only run on shutdown
   * @param creationExecutor Where valueCreator and valueDestroyer are run
   */
  public LazyInitMapValues(Function<K, V> valueCreator, BiConsumer<K, V> valueDestroyer,
                           ExecutorService creationExecutor) {
    this.valueCreator = valueCreator;
    this.valueDestroyer = valueDestroyer;
    this.creationExecutor = creationExecutor;
  }

  private void checkShutdown() {
    if (isShutdown) throw new IllegalStateException("LazyInitMapValues is shutdown");
  }


  /**
   * Retrieve the value for a key if that key was already computed, or compute it using
   * valueCreator and creationExecutor and block until the value is computed. The computing
   * function will be run exactly once per key, so it is safe to read from multiple threads
   * and perform non-reentrant or unrepeatable actions in the computing function.
   *
   * @param key Key whose value we want to create or retrieve
   * @return The value for this key whether cached or computed
   * @throws IllegalStateException when the map is shutdown, or shutdown occurred during value
   *                               computation
   * @throws InterruptedException  when the attempt to compute the value was interrupted for an
   *                               unknown reason unrelated to map shutdown
   * @throws CancellationException when the attempt to compute the value was interrupted for an
   *                               unknown reason unrelated to map shutdown
   * @throws ExecutionException when the computation function throws an exception while running
   */
  public V getOrCreate(K key) throws InterruptedException, ExecutionException {
    checkShutdown();
    final Future<V> box = objects.computeIfAbsent(key,
        (k) -> creationExecutor.submit(new InvokeValueCreator(k)));
    try {
      return box.get();
    } catch (InterruptedException | CancellationException e) {
      checkShutdown();
      throw e;
    }
  }


  /**
   * Reject all getOrCreate(key) attempts, interrupted-ly cancel all value creation, and destroy
   * all existing values using valueDestroyer running on the creationExecutor. Waits for at most
   * the specified timeout for all destruction to occur. This method <em>will not</em>
   * shutdown the executor.
   *
   * @param timeout
   * @param unit
   */
  public void shutdown(int timeout, TimeUnit unit) {
    isShutdown = true;
    final ArrayList<Callable<Boolean>> todestroy = new ArrayList<>();
    for (Map.Entry<K, Future<V>> entry : objects.entrySet()) {
      final K k = entry.getKey();
      final Future<V> box = entry.getValue();
      if (!box.cancel(true)) {
        todestroy.add(() -> {
          try {
            valueDestroyer.accept(k, box.get(timeout, unit));
          } catch (CancellationException | InterruptedException | ExecutionException e) {
            return false;
          }
          return true;
        });
      }
    }
    objects.clear();
    if (todestroy.size() > 0) {
      try {
        creationExecutor.invokeAll(todestroy, timeout, unit);
      } catch (RejectedExecutionException | InterruptedException ignored) {

      }
    }
  }

  private class InvokeValueCreator implements Callable<V> {
    final private K key;

    private InvokeValueCreator(K key) {
      this.key = key;
    }

    @Override
    public V call() {
      checkShutdown();
      return valueCreator.apply(key);
    }
  }
}