/*
 * Copyright (C) 2009 The Guava Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.common.cache;

import static com.google.common.cache.ReflectionFreeAssertThrows.assertThrows;
import static com.google.common.cache.TestingCacheLoaders.constantLoader;
import static com.google.common.cache.TestingCacheLoaders.identityLoader;
import static com.google.common.cache.TestingRemovalListeners.countingRemovalListener;
import static com.google.common.cache.TestingRemovalListeners.nullRemovalListener;
import static com.google.common.cache.TestingRemovalListeners.queuingRemovalListener;
import static com.google.common.cache.TestingWeighers.constantWeigher;
import static com.google.common.truth.Truth.assertThat;
import static java.util.concurrent.Executors.newFixedThreadPool;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.concurrent.TimeUnit.NANOSECONDS;
import static java.util.concurrent.TimeUnit.SECONDS;

import com.google.common.annotations.GwtCompatible;
import com.google.common.annotations.GwtIncompatible;
import com.google.common.base.Ticker;
import com.google.common.cache.TestingRemovalListeners.CountingRemovalListener;
import com.google.common.cache.TestingRemovalListeners.QueuingRemovalListener;
import com.google.common.collect.Sets;
import com.google.common.testing.NullPointerTester;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import java.time.Duration;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import junit.framework.TestCase;
import org.jspecify.annotations.NullUnmarked;

/** Unit tests for CacheBuilder. */
@GwtCompatible
// We are intentionally testing the TimeUnit overloads, too.
@SuppressWarnings("LongTimeUnit_ExpireAfterWrite_Seconds")
@NullUnmarked
public class CacheBuilderTest extends TestCase {

  public void testNewBuilder() {
    CacheLoader<Object, Integer> loader = constantLoader(1);

    LoadingCache<String, Integer> cache =
        CacheBuilder.newBuilder().removalListener(countingRemovalListener()).build(loader);

    assertEquals(Integer.valueOf(1), cache.getUnchecked("one"));
    assertEquals(1, cache.size());
  }

  public void testInitialCapacity_negative() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder();
    assertThrows(IllegalArgumentException.class, () -> builder.initialCapacity(-1));
  }

  public void testInitialCapacity_setTwice() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder().initialCapacity(16);
    assertThrows(IllegalStateException.class, () -> builder.initialCapacity(16));
  }

  @GwtIncompatible // CacheTesting
  public void testInitialCapacity_small() {
    LoadingCache<?, ?> cache = CacheBuilder.newBuilder().initialCapacity(5).build(identityLoader());
    LocalCache<?, ?> map = CacheTesting.toLocalCache(cache);

    assertThat(map.segments).hasLength(4);
    assertEquals(2, map.segments[0].table.length());
    assertEquals(2, map.segments[1].table.length());
    assertEquals(2, map.segments[2].table.length());
    assertEquals(2, map.segments[3].table.length());
  }

  @GwtIncompatible // CacheTesting
  public void testInitialCapacity_smallest() {
    LoadingCache<?, ?> cache = CacheBuilder.newBuilder().initialCapacity(0).build(identityLoader());
    LocalCache<?, ?> map = CacheTesting.toLocalCache(cache);

    assertThat(map.segments).hasLength(4);
    // 1 is as low as it goes, not 0. it feels dirty to know this/test this.
    assertEquals(1, map.segments[0].table.length());
    assertEquals(1, map.segments[1].table.length());
    assertEquals(1, map.segments[2].table.length());
    assertEquals(1, map.segments[3].table.length());
  }

  public void testInitialCapacity_large() {
    CacheBuilder.newBuilder().initialCapacity(Integer.MAX_VALUE);
    // that the builder didn't blow up is enough;
    // don't actually create this monster!
  }

  public void testConcurrencyLevel_zero() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder();
    assertThrows(IllegalArgumentException.class, () -> builder.concurrencyLevel(0));
  }

  public void testConcurrencyLevel_setTwice() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder().concurrencyLevel(16);
    assertThrows(IllegalStateException.class, () -> builder.concurrencyLevel(16));
  }

  @GwtIncompatible // CacheTesting
  public void testConcurrencyLevel_small() {
    LoadingCache<?, ?> cache =
        CacheBuilder.newBuilder().concurrencyLevel(1).build(identityLoader());
    LocalCache<?, ?> map = CacheTesting.toLocalCache(cache);
    assertThat(map.segments).hasLength(1);
  }

  public void testConcurrencyLevel_large() {
    CacheBuilder.newBuilder().concurrencyLevel(Integer.MAX_VALUE);
    // don't actually build this beast
  }

  public void testMaximumSize_negative() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder();
    assertThrows(IllegalArgumentException.class, () -> builder.maximumSize(-1));
  }

  public void testMaximumSize_setTwice() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder().maximumSize(16);
    assertThrows(IllegalStateException.class, () -> builder.maximumSize(16));
  }

  @GwtIncompatible // maximumWeight
  public void testMaximumSize_andWeight() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder().maximumSize(16);
    assertThrows(IllegalStateException.class, () -> builder.maximumWeight(16));
  }

  @GwtIncompatible // digs into internals of the non-GWT implementation
  public void testMaximumSize_largerThanInt() {
    CacheBuilder<Object, Object> builder =
        CacheBuilder.newBuilder().initialCapacity(512).maximumSize(Long.MAX_VALUE);
    LocalCache<?, ?> cache = ((LocalCache.LocalManualCache<?, ?>) builder.build()).localCache;
    assertThat(cache.segments.length * cache.segments[0].table.length()).isEqualTo(512);
  }

  @GwtIncompatible // maximumWeight
  public void testMaximumWeight_negative() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder();
    assertThrows(IllegalArgumentException.class, () -> builder.maximumWeight(-1));
  }

  @GwtIncompatible // maximumWeight
  public void testMaximumWeight_setTwice() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder().maximumWeight(16);
    assertThrows(IllegalStateException.class, () -> builder.maximumWeight(16));
    assertThrows(IllegalStateException.class, () -> builder.maximumSize(16));
  }

  @GwtIncompatible // maximumWeight
  public void testMaximumWeight_withoutWeigher() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder().maximumWeight(1);
    assertThrows(IllegalStateException.class, () -> builder.build(identityLoader()));
  }

  @GwtIncompatible // weigher
  public void testWeigher_withoutMaximumWeight() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder().weigher(constantWeigher(42));
    assertThrows(IllegalStateException.class, () -> builder.build(identityLoader()));
  }

  @GwtIncompatible // weigher
  public void testWeigher_withMaximumSize() {
    assertThrows(
        IllegalStateException.class,
        () -> CacheBuilder.newBuilder().weigher(constantWeigher(42)).maximumSize(1));
    assertThrows(
        IllegalStateException.class,
        () -> CacheBuilder.newBuilder().maximumSize(1).weigher(constantWeigher(42)));
  }

  @GwtIncompatible // weakKeys
  public void testKeyStrengthSetTwice() {
    CacheBuilder<Object, Object> builder1 = CacheBuilder.newBuilder().weakKeys();
    assertThrows(IllegalStateException.class, () -> builder1.weakKeys());
  }

  @GwtIncompatible // weakValues
  public void testValueStrengthSetTwice() {
    CacheBuilder<Object, Object> builder1 = CacheBuilder.newBuilder().weakValues();
    assertThrows(IllegalStateException.class, () -> builder1.weakValues());
    assertThrows(IllegalStateException.class, () -> builder1.softValues());

    CacheBuilder<Object, Object> builder2 = CacheBuilder.newBuilder().softValues();
    assertThrows(IllegalStateException.class, () -> builder2.softValues());
    assertThrows(IllegalStateException.class, () -> builder2.weakValues());
  }

  @GwtIncompatible // Duration
  @IgnoreJRERequirement // No more dangerous than wherever the caller got the Duration from
  public void testLargeDurationsAreOk() {
    Duration threeHundredYears = Duration.ofDays(365 * 300);
    CacheBuilder<Object, Object> unused =
        CacheBuilder.newBuilder()
            .expireAfterWrite(threeHundredYears)
            .expireAfterAccess(threeHundredYears)
            .refreshAfterWrite(threeHundredYears);
  }

  public void testTimeToLive_negative() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder();
    assertThrows(IllegalArgumentException.class, () -> builder.expireAfterWrite(-1, SECONDS));
  }

  @GwtIncompatible // Duration
  public void testTimeToLive_negative_duration() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder();
    assertThrows(
        IllegalArgumentException.class, () -> builder.expireAfterWrite(Duration.ofSeconds(-1)));
  }

  @SuppressWarnings("ReturnValueIgnored")
  public void testTimeToLive_small() {
    CacheBuilder.newBuilder().expireAfterWrite(1, NANOSECONDS).build(identityLoader());
    // well, it didn't blow up.
  }

  public void testTimeToLive_setTwice() {
    CacheBuilder<Object, Object> builder =
        CacheBuilder.newBuilder().expireAfterWrite(3600, SECONDS);
    assertThrows(IllegalStateException.class, () -> builder.expireAfterWrite(3600, SECONDS));
  }

  @GwtIncompatible // Duration
  public void testTimeToLive_setTwice_duration() {
    CacheBuilder<Object, Object> builder =
        CacheBuilder.newBuilder().expireAfterWrite(Duration.ofHours(1));
    assertThrows(IllegalStateException.class, () -> builder.expireAfterWrite(Duration.ofHours(1)));
  }

  public void testTimeToIdle_negative() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder();
    assertThrows(IllegalArgumentException.class, () -> builder.expireAfterAccess(-1, SECONDS));
  }

  @GwtIncompatible // Duration
  public void testTimeToIdle_negative_duration() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder();
    assertThrows(
        IllegalArgumentException.class, () -> builder.expireAfterAccess(Duration.ofSeconds(-1)));
  }

  @SuppressWarnings("ReturnValueIgnored")
  public void testTimeToIdle_small() {
    CacheBuilder.newBuilder().expireAfterAccess(1, NANOSECONDS).build(identityLoader());
    // well, it didn't blow up.
  }

  public void testTimeToIdle_setTwice() {
    CacheBuilder<Object, Object> builder =
        CacheBuilder.newBuilder().expireAfterAccess(3600, SECONDS);
    assertThrows(IllegalStateException.class, () -> builder.expireAfterAccess(3600, SECONDS));
  }

  @GwtIncompatible // Duration
  public void testTimeToIdle_setTwice_duration() {
    CacheBuilder<Object, Object> builder =
        CacheBuilder.newBuilder().expireAfterAccess(Duration.ofHours(1));
    assertThrows(IllegalStateException.class, () -> builder.expireAfterAccess(Duration.ofHours(1)));
  }

  public void testTimeToIdleAndToLive() {
    LoadingCache<?, ?> unused =
        CacheBuilder.newBuilder()
            .expireAfterWrite(1, NANOSECONDS)
            .expireAfterAccess(1, NANOSECONDS)
            .build(identityLoader());
    // well, it didn't blow up.
  }

  @GwtIncompatible // refreshAfterWrite
  public void testRefresh_zero() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder();
    assertThrows(IllegalArgumentException.class, () -> builder.refreshAfterWrite(0, SECONDS));
  }

  @GwtIncompatible // Duration
  public void testRefresh_zero_duration() {
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder();
    assertThrows(IllegalArgumentException.class, () -> builder.refreshAfterWrite(Duration.ZERO));
  }

  @GwtIncompatible // refreshAfterWrite
  public void testRefresh_setTwice() {
    CacheBuilder<Object, Object> builder =
        CacheBuilder.newBuilder().refreshAfterWrite(3600, SECONDS);
    assertThrows(IllegalStateException.class, () -> builder.refreshAfterWrite(3600, SECONDS));
  }

  @GwtIncompatible // Duration
  public void testRefresh_setTwice_duration() {
    CacheBuilder<Object, Object> builder =
        CacheBuilder.newBuilder().refreshAfterWrite(Duration.ofHours(1));
    assertThrows(IllegalStateException.class, () -> builder.refreshAfterWrite(Duration.ofHours(1)));
  }

  public void testTicker_setTwice() {
    Ticker testTicker = Ticker.systemTicker();
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder().ticker(testTicker);
    assertThrows(IllegalStateException.class, () -> builder.ticker(testTicker));
  }

  public void testRemovalListener_setTwice() {
    RemovalListener<Object, Object> testListener = nullRemovalListener();
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder().removalListener(testListener);
    assertThrows(IllegalStateException.class, () -> builder.removalListener(testListener));
  }

  public void testValuesIsNotASet() {
    assertFalse(CacheBuilder.newBuilder().build().asMap().values() instanceof Set);
  }

  @GwtIncompatible // CacheTesting
  public void testNullCache() {
    CountingRemovalListener<Object, Object> listener = countingRemovalListener();
    LoadingCache<Object, Object> nullCache =
        CacheBuilder.newBuilder().maximumSize(0).removalListener(listener).build(identityLoader());
    assertEquals(0, nullCache.size());
    Object key = new Object();
    assertSame(key, nullCache.getUnchecked(key));
    assertEquals(1, listener.getCount());
    assertEquals(0, nullCache.size());
    CacheTesting.checkEmpty(nullCache.asMap());
  }

  @GwtIncompatible // QueuingRemovalListener
  public void testRemovalNotification_clear() throws InterruptedException {
    // If a clear() happens while a computation is pending, we should not get a removal
    // notification.

    AtomicBoolean shouldWait = new AtomicBoolean(false);
    CountDownLatch computingLatch = new CountDownLatch(1);
    CacheLoader<String, String> computingFunction =
        new CacheLoader<String, String>() {
          @Override
          public String load(String key) throws InterruptedException {
            if (shouldWait.get()) {
              computingLatch.await();
            }
            return key;
          }
        };
    QueuingRemovalListener<String, String> listener = queuingRemovalListener();

    LoadingCache<String, String> cache =
        CacheBuilder.newBuilder()
            .concurrencyLevel(1)
            .removalListener(listener)
            .build(computingFunction);

    // seed the map, so its segment's count > 0
    cache.getUnchecked("a");
    shouldWait.set(true);

    CountDownLatch computationStarted = new CountDownLatch(1);
    CountDownLatch computationComplete = new CountDownLatch(1);
    new Thread(
            new Runnable() {
              @Override
              public void run() {
                computationStarted.countDown();
                cache.getUnchecked("b");
                computationComplete.countDown();
              }
            })
        .start();

    // wait for the computingEntry to be created
    computationStarted.await();
    cache.invalidateAll();
    // let the computation proceed
    computingLatch.countDown();
    // don't check cache.size() until we know the get("b") call is complete
    computationComplete.await();

    // At this point, the listener should be holding the seed value (a -> a), and the map should
    // contain the computed value (b -> b), since the clear() happened before the computation
    // completed.
    assertEquals(1, listener.size());
    RemovalNotification<String, String> notification = listener.remove();
    assertEquals("a", notification.getKey());
    assertEquals("a", notification.getValue());
    assertEquals(1, cache.size());
    assertEquals("b", cache.getUnchecked("b"));
  }

  // "Basher tests", where we throw a bunch of stuff at a LoadingCache and check basic invariants.

  /**
   * This is a less carefully-controlled version of {@link #testRemovalNotification_clear} - this is
   * a black-box test that tries to create lots of different thread-interleavings, and asserts that
   * each computation is affected by a call to {@code clear()} (and therefore gets passed to the
   * removal listener), or else is not affected by the {@code clear()} (and therefore exists in the
   * cache afterward).
   */
  @GwtIncompatible // QueuingRemovalListener

  @SuppressWarnings("ThreadPriorityCheck") // TODO: b/175898629 - Consider onSpinWait.
  public void testRemovalNotification_clear_basher() throws InterruptedException {
    // If a clear() happens close to the end of computation, one of two things should happen:
    // - computation ends first: the removal listener is called, and the cache does not contain the
    //   key/value pair
    // - clear() happens first: the removal listener is not called, and the cache contains the pair
    AtomicBoolean computationShouldWait = new AtomicBoolean();
    CountDownLatch computationLatch = new CountDownLatch(1);
    QueuingRemovalListener<String, String> listener = queuingRemovalListener();
    LoadingCache<String, String> cache =
        CacheBuilder.newBuilder()
            .removalListener(listener)
            .concurrencyLevel(20)
            .build(new DelayingIdentityLoader<String>(computationShouldWait, computationLatch));

    int nThreads = 100;
    int nTasks = 1000;
    int nSeededEntries = 100;
    Set<String> expectedKeys = Sets.newHashSetWithExpectedSize(nTasks + nSeededEntries);
    // seed the map, so its segments have a count>0; otherwise, clear() won't visit the in-progress
    // entries
    for (int i = 0; i < nSeededEntries; i++) {
      String s = "b" + i;
      cache.getUnchecked(s);
      expectedKeys.add(s);
    }
    computationShouldWait.set(true);

    AtomicInteger computedCount = new AtomicInteger();
    ExecutorService threadPool = newFixedThreadPool(nThreads);
    CountDownLatch tasksFinished = new CountDownLatch(nTasks);
    for (int i = 0; i < nTasks; i++) {
      String s = "a" + i;
      @SuppressWarnings("unused") // https://errorprone.info/bugpattern/FutureReturnValueIgnored
      Future<?> possiblyIgnoredError =
          threadPool.submit(
              new Runnable() {
                @Override
                public void run() {
                  cache.getUnchecked(s);
                  computedCount.incrementAndGet();
                  tasksFinished.countDown();
                }
              });
      expectedKeys.add(s);
    }

    computationLatch.countDown();
    // let some computations complete
    while (computedCount.get() < nThreads) {
      Thread.yield();
    }
    cache.invalidateAll();
    tasksFinished.await();

    // Check all of the removal notifications we received: they should have had correctly-associated
    // keys and values. (An earlier bug saw removal notifications for in-progress computations,
    // which had real keys with null values.)
    Map<String, String> removalNotifications = new HashMap<>();
    for (RemovalNotification<String, String> notification : listener) {
      removalNotifications.put(notification.getKey(), notification.getValue());
      assertEquals(
          "Unexpected key/value pair passed to removalListener",
          notification.getKey(),
          notification.getValue());
    }

    // All of the seed values should have been visible, so we should have gotten removal
    // notifications for all of them.
    for (int i = 0; i < nSeededEntries; i++) {
      assertEquals("b" + i, removalNotifications.get("b" + i));
    }

    // Each of the values added to the map should either still be there, or have seen a removal
    // notification.
    assertEquals(expectedKeys, Sets.union(cache.asMap().keySet(), removalNotifications.keySet()));
    assertTrue(Sets.intersection(cache.asMap().keySet(), removalNotifications.keySet()).isEmpty());
    threadPool.shutdown();
    threadPool.awaitTermination(300, SECONDS);
  }

  /**
   * Calls get() repeatedly from many different threads, and tests that all of the removed entries
   * (removed because of size limits or expiration) trigger appropriate removal notifications.
   */
  @GwtIncompatible // QueuingRemovalListener

  public void testRemovalNotification_get_basher() throws InterruptedException {
    int nTasks = 1000;
    int nThreads = 100;
    int getsPerTask = 1000;
    int nUniqueKeys = 10000;
    Random random = new Random(); // Randoms.insecureRandom();

    QueuingRemovalListener<String, String> removalListener = queuingRemovalListener();
    AtomicInteger computeCount = new AtomicInteger();
    AtomicInteger exceptionCount = new AtomicInteger();
    AtomicInteger computeNullCount = new AtomicInteger();
    @SuppressWarnings("CacheLoaderNull") // test of handling of erroneous implementation
    CacheLoader<String, String> countingIdentityLoader =
        new CacheLoader<String, String>() {
          @Override
          public String load(String key) throws InterruptedException {
            int behavior = random.nextInt(4);
            if (behavior == 0) { // throw an exception
              exceptionCount.incrementAndGet();
              throw new RuntimeException("fake exception for test");
            } else if (behavior == 1) { // return null
              computeNullCount.incrementAndGet();
              return null;
            } else if (behavior == 2) { // slight delay before returning
              Thread.sleep(5);
              computeCount.incrementAndGet();
              return key;
            } else {
              computeCount.incrementAndGet();
              return key;
            }
          }
        };
    LoadingCache<String, String> cache =
        CacheBuilder.newBuilder()
            .recordStats()
            .concurrencyLevel(2)
            .expireAfterWrite(100, MILLISECONDS)
            .removalListener(removalListener)
            .maximumSize(5000)
            .build(countingIdentityLoader);

    ExecutorService threadPool = newFixedThreadPool(nThreads);
    for (int i = 0; i < nTasks; i++) {
      @SuppressWarnings("unused") // https://errorprone.info/bugpattern/FutureReturnValueIgnored
      Future<?> possiblyIgnoredError =
          threadPool.submit(
              new Runnable() {
                @Override
                public void run() {
                  for (int j = 0; j < getsPerTask; j++) {
                    try {
                      cache.getUnchecked("key" + random.nextInt(nUniqueKeys));
                    } catch (RuntimeException e) {
                    }
                  }
                }
              });
    }

    threadPool.shutdown();
    threadPool.awaitTermination(300, SECONDS);

    // Since we're not doing any more cache operations, and the cache only expires/evicts when doing
    // other operations, the cache and the removal queue won't change from this point on.

    // Verify that each received removal notification was valid
    for (RemovalNotification<String, String> notification : removalListener) {
      assertEquals("Invalid removal notification", notification.getKey(), notification.getValue());
    }

    CacheStats stats = cache.stats();
    assertEquals(removalListener.size(), stats.evictionCount());
    assertEquals(computeCount.get(), stats.loadSuccessCount());
    assertEquals(exceptionCount.get() + computeNullCount.get(), stats.loadExceptionCount());
    // each computed value is still in the cache, or was passed to the removal listener
    assertEquals(computeCount.get(), cache.size() + removalListener.size());
  }

  @GwtIncompatible // NullPointerTester
  public void testNullParameters() throws Exception {
    NullPointerTester tester = new NullPointerTester();
    CacheBuilder<Object, Object> builder = CacheBuilder.newBuilder();
    tester.testAllPublicInstanceMethods(builder);
  }

  @GwtIncompatible // CacheTesting
  public void testSizingDefaults() {
    LoadingCache<?, ?> cache = CacheBuilder.newBuilder().build(identityLoader());
    LocalCache<?, ?> map = CacheTesting.toLocalCache(cache);
    assertThat(map.segments).hasLength(4); // concurrency level
    assertEquals(4, map.segments[0].table.length()); // capacity / conc level
  }

  @GwtIncompatible // CountDownLatch
  static final class DelayingIdentityLoader<T> extends CacheLoader<T, T> {
    private final AtomicBoolean shouldWait;
    private final CountDownLatch delayLatch;

    DelayingIdentityLoader(AtomicBoolean shouldWait, CountDownLatch delayLatch) {
      this.shouldWait = shouldWait;
      this.delayLatch = delayLatch;
    }

    @CanIgnoreReturnValue // Sure, why not?
    @Override
    public T load(T key) throws InterruptedException {
      if (shouldWait.get()) {
        delayLatch.await();
      }
      return key;
    }
  }
}
