RATIS-2266. Use WeakValueCache instead of Guava cache in RaftId. (#1240)

Author: szetszwo

ratis-common/src/main/java/org/apache/ratis/protocol/ClientId.java

@@ -18,6 +18,7 @@
 package org.apache.ratis.protocol;
 
 import org.apache.ratis.thirdparty.com.google.protobuf.ByteString;
+import org.apache.ratis.util.WeakValueCache;
 
 import java.util.UUID;
 
@@ -26,13 +27,17 @@
  * to correctly identify retry requests from the same client.
  */
 public final class ClientId extends RaftId {
-  private static final Factory<ClientId> FACTORY = new Factory<ClientId>() {
+  private static final Factory<ClientId> FACTORY = new Factory<ClientId>(ClientId.class) {
     @Override
     ClientId newInstance(UUID uuid) {
       return new ClientId(uuid);
     }
   };
 
+  static WeakValueCache<UUID, ClientId> getCache() {
+    return FACTORY.getCache();
+  }
+
   public static ClientId emptyClientId() {
     return FACTORY.emptyId();
   }

ratis-common/src/main/java/org/apache/ratis/protocol/RaftGroupId.java

@@ -18,6 +18,7 @@
 package org.apache.ratis.protocol;
 
 import org.apache.ratis.thirdparty.com.google.protobuf.ByteString;
+import org.apache.ratis.util.WeakValueCache;
 
 import java.util.UUID;
 
@@ -27,13 +28,17 @@
  * This is a value-based class.
  */
 public final class RaftGroupId extends RaftId {
-  private static final Factory<RaftGroupId> FACTORY = new Factory<RaftGroupId>() {
+  private static final Factory<RaftGroupId> FACTORY = new Factory<RaftGroupId>(RaftGroupId.class) {
     @Override
     RaftGroupId newInstance(UUID uuid) {
       return new RaftGroupId(uuid);
     }
   };
 
+  static WeakValueCache<UUID, RaftGroupId> getCache() {
+    return FACTORY.getCache();
+  }
+
   public static RaftGroupId emptyGroupId() {
     return FACTORY.emptyId();
   }

ratis-common/src/main/java/org/apache/ratis/protocol/RaftId.java

@@ -17,17 +17,15 @@
  */
 package org.apache.ratis.protocol;
 
-import org.apache.ratis.thirdparty.com.google.common.cache.Cache;
-import org.apache.ratis.thirdparty.com.google.common.cache.CacheBuilder;
 import org.apache.ratis.thirdparty.com.google.protobuf.ByteString;
 import org.apache.ratis.thirdparty.com.google.protobuf.UnsafeByteOperations;
 import org.apache.ratis.util.JavaUtils;
 import org.apache.ratis.util.Preconditions;
+import org.apache.ratis.util.WeakValueCache;
 
 import java.nio.ByteBuffer;
 import java.util.Objects;
 import java.util.UUID;
-import java.util.concurrent.ExecutionException;
 import java.util.function.Supplier;
 
 /** Unique identifier implemented using {@link UUID}. */
@@ -53,18 +51,20 @@ static ByteString toByteString(UUID uuid) {
   }
 
   abstract static class Factory<ID extends RaftId> {
-    private final Cache<UUID, ID> cache = CacheBuilder.newBuilder()
-        .weakValues()
-        .build();
+    private final WeakValueCache<UUID, ID> cache;
+
+    Factory(Class<ID> clazz) {
+      this.cache = new WeakValueCache<>(clazz.getSimpleName() + "_UUID", this::newInstance);
+    }
 
     abstract ID newInstance(UUID uuid);
 
+    WeakValueCache<UUID, ID> getCache() {
+      return cache;
+    }
+
     final ID valueOf(UUID uuid) {
-      try {
-        return cache.get(uuid, () -> newInstance(uuid));
-      } catch (ExecutionException e) {
-        throw new IllegalStateException("Failed to valueOf(" + uuid + ")", e);
-      }
+      return cache.getOrCreate(uuid);
     }
 
     final ID valueOf(ByteString bytes) {

ratis-common/src/main/java/org/apache/ratis/util/BiWeakValueCache.java

@@ -33,13 +33,15 @@
  * Note that the cached values are weakly referenced.
  * A cached value could be garage-collected (i.e. evicted from the cache)
  * when there are no external (strong) references.
+ * <p>
+ * For key types with a component, use {@link WeakValueCache}.
  *
  * @param <OUTER> the type of the outer keys.
  * @param <INNER> the type of the inner keys.
  * @param <T> the type to be cached.
  */
 public final class BiWeakValueCache<OUTER, INNER, T> {
-  private static <K, V> ConcurrentMap<K, V> newMap() {
+  static <K, V> ConcurrentMap<K, V> newMap() {
     return new MapMaker().weakValues().makeMap();
   }
 
@@ -61,8 +63,8 @@ private static <K, V> ConcurrentMap<K, V> newMap() {
   /**
    * Create a cache for mapping ({@link OUTER}, {@link INNER}) keys to {@link T} values.
    *
-   * @param outerName the name of the outer long.
-   * @param innerName the name of the inner long.
+   * @param outerName the name of the outer keys.
+   * @param innerName the name of the inner keys.
    * @param constructor for constructing {@link T} values.
    */
   public BiWeakValueCache(String outerName, String innerName, BiFunction<OUTER, INNER, T> constructor) {

ratis-common/src/main/java/org/apache/ratis/util/WeakValueCache.java

@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.ratis.util;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
+
+import static org.apache.ratis.util.BiWeakValueCache.newMap;
+
+/**
+ * Weak Value Cache: {@link K} -> {@link V}.
+ * <p>
+ * Note that the cached values are weakly referenced.
+ * A cached value could be garage-collected (i.e. evicted from the cache)
+ * when there are no external (strong) references.
+ * <p>
+ * For key types with two components, use {@link BiWeakValueCache}.
+ *
+ * @param <K> the type of the keys.
+ * @param <V> the type to be cached values.
+ */
+public final class WeakValueCache<K, V> {
+  private final String keyName;
+  private final String name;
+
+  /** For constructing a value from a key. */
+  private final Function<K, V> constructor;
+  /** Count the number of values constructed. */
+  private final AtomicInteger constructionCount = new AtomicInteger(0);
+
+  /** Map: {@link K} -> {@link V}. */
+  private final ConcurrentMap<K, V> map = newMap();
+
+  /**
+   * Create a cache for mapping {@link K} keys to {@link V} values.
+   *
+   * @param keyName the name of the key.
+   * @param constructor for constructing {@link V} values.
+   */
+  public WeakValueCache(String keyName, Function<K, V> constructor) {
+    this.keyName = keyName;
+    this.name = keyName + "-cache";
+    this.constructor = constructor;
+  }
+
+  private V construct(K key) {
+    final V constructed = constructor.apply(key);
+    Objects.requireNonNull(constructed, "constructed == null");
+    constructionCount.incrementAndGet();
+    return constructed;
+  }
+
+  /**
+   * If the given key is in the cache, return its cached values.
+   * Otherwise, create a new value, put it in the cache and then return it.
+   */
+  public V getOrCreate(K key) {
+    Objects.requireNonNull(key, () -> keyName + " (key) == null");
+    return map.computeIfAbsent(key, this::construct);
+  }
+
+  List<V> getValues() {
+    return new ArrayList<>(map.values());
+  }
+
+  @Override
+  public String toString() {
+    return name;
+  }
+}

ratis-test/src/test/java/org/apache/ratis/protocol/TestRaftId.java

@@ -19,6 +19,7 @@
 
 import org.apache.ratis.BaseTest;
 import org.apache.ratis.thirdparty.com.google.protobuf.ByteString;
+import org.apache.ratis.util.WeakValueCache;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
@@ -27,6 +28,13 @@
 
 @Timeout(value = 1)
 public class TestRaftId extends BaseTest {
+  public static WeakValueCache<UUID, ClientId> getClientIdCache() {
+    return ClientId.getCache();
+  }
+
+  public static WeakValueCache<UUID, RaftGroupId> getRaftGroupIdCache() {
+    return RaftGroupId.getCache();
+  }
 
   @Test
   public void testRaftId() {

ratis-test/src/test/java/org/apache/ratis/util/TestRaftIdCache.java

@@ -0,0 +1,129 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.ratis.util;
+
+import org.apache.ratis.BaseTest;
+import org.apache.ratis.RaftTestUtil;
+import org.apache.ratis.protocol.ClientId;
+import org.apache.ratis.protocol.TestRaftId;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.ThreadLocalRandom;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
+
+/** Testing {@link WeakValueCache}. */
+public class TestRaftIdCache extends BaseTest {
+  static WeakValueCache<UUID, ClientId> CACHE = TestRaftId.getClientIdCache();
+
+  static String dumpCache() {
+    final List<ClientId> values = CACHE.getValues();
+    values.sort(Comparator.comparing(ClientId::getUuid));
+    String header = CACHE + ": " + values.size();
+    System.out.println(header);
+    System.out.println("  " + values);
+    return header;
+  }
+
+  static void assertCache(IDs expectedIDs) {
+    final List<ClientId> computed = CACHE.getValues();
+    computed.sort(Comparator.comparing(ClientId::getUuid));
+
+    final List<ClientId> expected = expectedIDs.getIds();
+    expected.sort(Comparator.comparing(ClientId::getUuid));
+
+    assertEquals(expected, computed, TestRaftIdCache::dumpCache);
+  }
+
+  void assertCacheSizeWithGC(IDs expectedIDs) throws Exception{
+    JavaUtils.attempt(() -> {
+      RaftTestUtil.gc();
+      assertCache(expectedIDs);
+    }, 5, HUNDRED_MILLIS, "assertCacheSizeWithGC", LOG);
+  }
+
+  class IDs {
+    private final List<ClientId> ids = new LinkedList<>();
+
+    List<ClientId> getIds() {
+      return new ArrayList<>(ids);
+    }
+
+    int size() {
+      return ids.size();
+    }
+
+    ClientId allocate() {
+      final ClientId id = ClientId.randomId();
+      LOG.info("allocate {}", id);
+      ids.add(id);
+      return id;
+    }
+
+    void release() {
+      final int r = ThreadLocalRandom.current().nextInt(size());
+      final ClientId removed = ids.remove(r);
+      LOG.info("release {}", removed);
+    }
+  }
+
+  @Test
+  public void testCaching() throws Exception {
+    final int n = 100;
+    final IDs ids = new IDs();
+    assertEquals(0, ids.size());
+    assertCache(ids);
+
+    for(int i = 0; i < n; i++) {
+      final ClientId id = ids.allocate();
+      assertSame(id, ClientId.valueOf(id.getUuid()));
+      assertCache(ids);
+    }
+
+    for(int i = 0; i < n/2; i++) {
+      ids.release();
+      if (ThreadLocalRandom.current().nextInt(10) == 0) {
+        assertCacheSizeWithGC(ids);
+      }
+    }
+    assertCacheSizeWithGC(ids);
+
+    for(int i = 0; i < n/2; i++) {
+      final ClientId id = ids.allocate();
+      assertSame(id, ClientId.valueOf(id.getUuid()));
+      assertCache(ids);
+    }
+
+
+    for(int i = 0; i < n; i++) {
+      ids.release();
+      if (ThreadLocalRandom.current().nextInt(10) == 0) {
+        assertCacheSizeWithGC(ids);
+      }
+    }
+    assertCacheSizeWithGC(ids);
+
+    assertEquals(0, ids.size());
+  }
+}