Introduce unloaders on PerThreadCache entries

RepoRefCache holds reference to a Repository object
which would need once the cache entry is removed.

Add the logic to cleanup the entry resource when the
cache is closed and therefore entries need to be
cleaned up before being removed from cache.

Release-Notes: skip
Change-Id: Id201d4a93792b27d2e1ab8214aa6377387eb567f

3 files changed, 131 insertions, 22 deletions

diff --git a/java/com/google/gerrit/server/cache/PerThreadCache.java b/java/com/google/gerrit/server/cache/PerThreadCache.java
index fe27bfb597..459650d90c 100644
--- a/java/com/google/gerrit/server/cache/PerThreadCache.java
+++ b/java/com/google/gerrit/server/cache/PerThreadCache.java
@@ -21,7 +21,11 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Maps;
 import com.google.gerrit.common.Nullable;
 import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Optional;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
+import java.util.stream.Stream;
 import javax.servlet.http.HttpServletRequest;
 
 /**
@@ -64,6 +68,9 @@ public class PerThreadCache implements AutoCloseable {
    */
   private final boolean readOnlyRequest;
 
+  private final Map<Key<?>, Consumer<Object>> unloaders =
+      Maps.newHashMapWithExpectedSize(PER_THREAD_CACHE_SIZE);
+
   /**
    * Unique key for key-value mappings stored in PerThreadCache. The key is based on the value's
    * class and a list of identifiers that in combination uniquely set the object apart form others
@@ -141,6 +148,38 @@ public class PerThreadCache implements AutoCloseable {
     return CACHE.get();
   }
 
+  /**
+   * Return a cached value associated with a key fetched with a loader and released with an unloader
+   * function.
+   *
+   * @param <T> The data type of the cached value
+   * @param key the key associated with the value
+   * @param loader the loader function for fetching the value from the key
+   * @param unloader the unloader function for releasing the value when unloaded from the cache
+   * @return Optional of the cached value or empty if the value could not be cached for any reason
+   *     (e.g. cache full)
+   */
+  public static <T> Optional<T> get(Key<T> key, Supplier<T> loader, Consumer<T> unloader) {
+    return Optional.ofNullable(get()).flatMap(c -> c.getWithLoader(key, loader, unloader));
+  }
+
+  /**
+   * Legacy way for retrieving a cached element through a loader.
+   *
+   * <p>This method is deprecated because it was error-prone due to the unclear ownership of the
+   * objects created through the loader. When the cache has space available, the entries are loaded
+   * and cached, hence owned and reused by the cache.
+   *
+   * <p>When the cache is full, this method just short-circuit to the invocation of the loader and
+   * the objects created aren't owned or stored by the cache, leaving the space for potential memory
+   * and resources leaks.
+   *
+   * <p>Because of the unclear semantics of the method (who owns the instances? are they reused?)
+   * this is now deprecated the the caller should use instead the {@link PerThreadCache#get(Key,
+   * Supplier, Consumer)} which has a clear ownership policy.
+   *
+   * @deprecated use {@link PerThreadCache#get(Key, Supplier, Consumer)}
+   */
   public static <T> T getOrCompute(Key<T> key, Supplier<T> loader) {
     PerThreadCache cache = get();
     return cache != null ? cache.get(key, loader) : loader.get();
@@ -157,19 +196,39 @@ public class PerThreadCache implements AutoCloseable {
   }
 
   /**
-   * Returns an instance of {@code T} that was either loaded from the cache or obtained from the
-   * provided {@link Supplier}.
+   * Legacy way of retrieving an instance of {@code T} that was either loaded from the cache or
+   * obtained from the provided {@link Supplier}.
+   *
+   * <p>This method is deprecated because it was error-prone due to the unclear ownership of the
+   * objects created through the loader. When the cache has space available, the entries are loaded
+   * and cached, hence owned and reused by the cache.
+   *
+   * <p>When the cache is full, this method just short-circuit to the invocation of the loader and
+   * the objects created aren't owned or stored by the cache, leaving the space for potential memory
+   * and resources leaks.
+   *
+   * <p>Because of the unclear semantics of the method (who owns the instances? are they reused?)
+   * this is now deprecated the the caller should use instead the {@link PerThreadCache#get(Key,
+   * Supplier, Consumer)} which has a clear ownership policy.
+   *
+   * @deprecated use {@link PerThreadCache#getWithLoader(Key, Supplier, Consumer)}
    */
   public <T> T get(Key<T> key, Supplier<T> loader) {
-    @SuppressWarnings("unchecked")
+    return getWithLoader(key, loader, null).orElse(loader.get());
+  }
+
+  @SuppressWarnings("unchecked")
+  public <T> Optional<T> getWithLoader(
+      Key<T> key, Supplier<T> loader, @Nullable Consumer<T> unloader) {
     T value = (T) cache.get(key);
-    if (value == null) {
+    if (value == null && cache.size() < PER_THREAD_CACHE_SIZE) {
       value = loader.get();
-      if (cache.size() < PER_THREAD_CACHE_SIZE) {
-        cache.put(key, value);
+      cache.put(key, value);
+      if (unloader != null) {
+        unloaders.put(key, (Consumer<Object>) unloader);
       }
     }
-    return value;
+    return Optional.ofNullable(value);
   }
 
   /** Returns true if the associated request is read-only */
@@ -179,6 +238,15 @@ public class PerThreadCache implements AutoCloseable {
 
   @Override
   public void close() {
+    Optional.of(CACHE.get())
+        .map(v -> v.unloaders.entrySet().stream())
+        .orElse(Stream.empty())
+        .forEach(this::unload);
+
     CACHE.remove();
   }
+
+  private <T> void unload(Entry<Key<?>, Consumer<Object>> unloaderEntry) {
+    unloaderEntry.getValue().accept(cache.get(unloaderEntry.getKey()));
+  }
 }
diff --git a/java/com/google/gerrit/server/git/RepoRefCache.java b/java/com/google/gerrit/server/git/RepoRefCache.java
index 1fe905e5e3..2f60f580e0 100644
--- a/java/com/google/gerrit/server/git/RepoRefCache.java
+++ b/java/com/google/gerrit/server/git/RepoRefCache.java
@@ -34,9 +34,12 @@ public class RepoRefCache implements RefCache {
   public static Optional<RefCache> getOptional(Repository repo) {
     PerThreadCache cache = PerThreadCache.get();
     if (cache != null && cache.hasReadonlyRequest()) {
-      return Optional.of(
-          cache.get(
-              PerThreadCache.Key.create(RepoRefCache.class, repo), () -> new RepoRefCache(repo)));
+      return cache
+          .getWithLoader(
+              PerThreadCache.Key.create(RepoRefCache.class, repo),
+              () -> new RepoRefCache(repo),
+              RepoRefCache::close)
+          .map(c -> (RefCache) c);
     }
 
     return Optional.empty();
diff --git a/javatests/com/google/gerrit/server/cache/PerThreadCacheTest.java b/javatests/com/google/gerrit/server/cache/PerThreadCacheTest.java
index 5a69401df8..2905e570ac 100644
--- a/javatests/com/google/gerrit/server/cache/PerThreadCacheTest.java
+++ b/javatests/com/google/gerrit/server/cache/PerThreadCacheTest.java
@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
 
 import com.google.gerrit.util.http.testutil.FakeHttpServletRequest;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletRequestWrapper;
@@ -50,14 +51,13 @@ public class PerThreadCacheTest {
       PerThreadCache cache = PerThreadCache.get();
       PerThreadCache.Key<String> key1 = PerThreadCache.Key.create(String.class);
 
-      String value1 = cache.get(key1, () -> "value1");
-      assertThat(value1).isEqualTo("value1");
+      assertThat(cache.getWithLoader(key1, () -> "value1", null)).hasValue("value1");
 
       Supplier<String> neverCalled =
           () -> {
             throw new IllegalStateException("this method must not be called");
           };
-      assertThat(cache.get(key1, neverCalled)).isEqualTo("value1");
+      assertThat(cache.getWithLoader(key1, neverCalled, null)).hasValue("value1");
     }
   }
 
@@ -66,20 +66,30 @@ public class PerThreadCacheTest {
     PerThreadCache.Key<String> key = PerThreadCache.Key.create(String.class);
     try (PerThreadCache ignored = PerThreadCache.create(null)) {
       PerThreadCache cache = PerThreadCache.get();
-      String value1 = cache.get(key, () -> "value1");
-      assertThat(value1).isEqualTo("value1");
+      assertThat(cache.getWithLoader(key, () -> "value1", null)).hasValue("value1");
     }
 
     // Create a second cache and assert that it is not connected to the first one.
     // This ensures that the cleanup is actually working.
     try (PerThreadCache ignored = PerThreadCache.create(null)) {
       PerThreadCache cache = PerThreadCache.get();
-      String value1 = cache.get(key, () -> "value2");
-      assertThat(value1).isEqualTo("value2");
+      assertThat(cache.getWithLoader(key, () -> "value2", null)).hasValue("value2");
     }
   }
 
   @Test
+  public void unloaderCalledUponCleanup() {
+    AtomicBoolean unloaderCalled = new AtomicBoolean();
+    PerThreadCache.Key<String> key = PerThreadCache.Key.create(String.class);
+    try (PerThreadCache ignored = PerThreadCache.create(null)) {
+      PerThreadCache cache = PerThreadCache.get();
+      cache.getWithLoader(key, () -> "value1", v -> unloaderCalled.set(true));
+      assertThat(unloaderCalled.get()).isFalse();
+    }
+    assertThat(unloaderCalled.get()).isTrue();
+  }
+
+  @Test
   public void doubleInstantiationFails() {
     try (PerThreadCache ignored = PerThreadCache.create(null)) {
       exception.expect(IllegalStateException.class);
@@ -124,14 +134,12 @@ public class PerThreadCacheTest {
     }
   }
 
+  @SuppressWarnings("deprecation")
   @Test
   public void enforceMaxSize() {
     try (PerThreadCache cache = PerThreadCache.create(null)) {
-      // Fill the cache
-      for (int i = 0; i < 50; i++) {
-        PerThreadCache.Key<String> key = PerThreadCache.Key.create(String.class, i);
-        cache.get(key, () -> "cached value");
-      }
+      fillTheCache(cache);
+
       // Assert that the value was not persisted
       PerThreadCache.Key<String> key = PerThreadCache.Key.create(String.class, 1000);
       cache.get(key, () -> "new value");
@@ -139,4 +147,34 @@ public class PerThreadCacheTest {
       assertThat(value).isEqualTo("directly served");
     }
   }
+
+  @Test
+  public void returnEmptyWhenCacheIsFull() {
+    try (PerThreadCache cache = PerThreadCache.create(null)) {
+      fillTheCache(cache);
+
+      PerThreadCache.Key<String> key = PerThreadCache.Key.create(String.class, 1000);
+      assertThat(cache.getWithLoader(key, () -> "new value", null)).isEmpty();
+    }
+  }
+
+  @Test
+  public void unloaderNotCalledUponCleanupIfCacheWasFull() {
+    AtomicBoolean unloaderCalled = new AtomicBoolean();
+    PerThreadCache.Key<String> key = PerThreadCache.Key.create(String.class);
+    try (PerThreadCache ignored = PerThreadCache.create(null)) {
+      PerThreadCache cache = PerThreadCache.get();
+      fillTheCache(cache);
+
+      cache.getWithLoader(key, () -> "value1", v -> unloaderCalled.set(true));
+    }
+    assertThat(unloaderCalled.get()).isFalse();
+  }
+
+  private void fillTheCache(PerThreadCache cache) {
+    for (int i = 0; i < 50; i++) {
+      PerThreadCache.Key<String> key = PerThreadCache.Key.create(String.class, i);
+      cache.getWithLoader(key, () -> "cached value", null);
+    }
+  }
 }