Add coroutine-based LoadingCache. (#232)

MODULE.bazel

@@ -196,6 +196,7 @@ maven.install(
         "org.reactivestreams:reactive-streams:1.0.4",
         "io.projectreactor:reactor-core:3.4.19",
         "com.google.crypto.tink:tink:" + TINK_VERSION,
+        "com.github.ben-manes.caffeine:caffeine:3.1.8",
 
         # AWS.
         "software.amazon.awssdk:auth:" + AWS_JAVA_SDK_VERSION,

imports/java/com/github/benmanes/caffeine/BUILD.bazel

@@ -0,0 +1,6 @@
+package(default_visibility = ["//visibility:public"])
+
+alias(
+    name = "caffeine",
+    actual = "@maven//:com_github_ben_manes_caffeine_caffeine",
+)

src/main/kotlin/org/wfanet/measurement/common/BUILD.bazel

@@ -6,6 +6,7 @@ kt_jvm_library(
     name = "common",
     srcs = glob(["*.kt"]),
     deps = [
+        "//imports/java/com/github/benmanes/caffeine",
         "//imports/java/com/google/common:guava",
         "//imports/java/com/google/devtools/build/runfiles",
         "//imports/java/com/google/gson",

src/main/kotlin/org/wfanet/measurement/common/LoadingCache.kt

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2024 The Cross-Media Measurement 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 org.wfanet.measurement.common
+
+import com.github.benmanes.caffeine.cache.AsyncCache
+import com.github.benmanes.caffeine.cache.AsyncLoadingCache
+import java.util.concurrent.CompletableFuture
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.asCoroutineDispatcher
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.future.await
+import kotlinx.coroutines.future.future
+
+/** Coroutine wrapper around [asyncCache] which loads values using [load]. */
+class LoadingCache<K, V>(
+  private val asyncCache: AsyncCache<K, V>,
+  private val load: suspend (K) -> V?,
+) {
+  init {
+    require(asyncCache !is AsyncLoadingCache) { "asyncCache already has an associated loader" }
+  }
+
+  /** Returns the value for [key], or `null` if no value could be loaded. */
+  suspend fun get(key: K): V? = coroutineScope { getAsync(key).await() }
+
+  /**
+   * Returns the value for [key].
+   *
+   * @throws NoSuchElementException if no value could be loaded for [key]
+   */
+  suspend fun getValue(key: K): V {
+    return get(key) ?: throw NoSuchElementException("No element with key $key")
+  }
+
+  private fun CoroutineScope.getAsync(key: K): CompletableFuture<V?> {
+    return asyncCache.get(key) { k, executor ->
+      future(executor.asCoroutineDispatcher()) { load(k) }
+    }
+  }
+}

src/test/kotlin/org/wfanet/measurement/common/BUILD.bazel

@@ -130,3 +130,17 @@ kt_jvm_test(
         "//src/main/proto/wfa/measurement/common/testing:depends_on_simple_kt_jvm_proto",
     ],
 )
+
+kt_jvm_test(
+    name = "LoadingCacheTest",
+    srcs = ["LoadingCacheTest.kt"],
+    test_class = "org.wfanet.measurement.common.LoadingCacheTest",
+    deps = [
+        "//imports/java/com/github/benmanes/caffeine",
+        "//imports/java/com/google/common/truth",
+        "//imports/java/org/junit",
+        "//imports/kotlin/kotlin/test",
+        "//imports/kotlin/kotlinx/coroutines:core",
+        "//src/main/kotlin/org/wfanet/measurement/common",
+    ],
+)

src/test/kotlin/org/wfanet/measurement/common/LoadingCacheTest.kt

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2024 The Cross-Media Measurement 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 org.wfanet.measurement.common
+
+import com.github.benmanes.caffeine.cache.Caffeine
+import com.google.common.truth.Truth.assertThat
+import kotlin.test.assertFailsWith
+import kotlinx.coroutines.runBlocking
+import org.junit.Test
+
+/** Test for [LoadingCache]. */
+class LoadingCacheTest {
+  @Test
+  fun `ctor throws IllegalArgumentException when asyncCache has loader`() {
+    val asyncCache = Caffeine.newBuilder().buildAsync<String, String> { key -> "$key-value" }
+    assertFailsWith<IllegalArgumentException> { LoadingCache(asyncCache) { "$it-value" } }
+  }
+
+  @Test
+  fun `get rethrows exception from load`() {
+    val asyncCache = Caffeine.newBuilder().buildAsync<String, String>()
+    val loadingCache = LoadingCache(asyncCache) { error("Loading error") }
+    assertFailsWith<IllegalStateException> { runBlocking { loadingCache.get("foo") } }
+  }
+
+  @Test
+  fun `get returns loaded value`() {
+    val asyncCache = Caffeine.newBuilder().buildAsync<String, String>()
+    val loadingCache = LoadingCache(asyncCache) { key -> "$key-value" }
+
+    val value: String? = runBlocking { loadingCache.get("foo") }
+
+    assertThat(value).isEqualTo("foo-value")
+  }
+
+  @Test
+  fun `getValue returns loaded value`() {
+    val asyncCache = Caffeine.newBuilder().buildAsync<String, String>()
+    val loadingCache = LoadingCache(asyncCache) { key -> "$key-value" }
+
+    val value: String = runBlocking { loadingCache.getValue("foo") }
+
+    assertThat(value).isEqualTo("foo-value")
+  }
+
+  @Test
+  fun `get returns null when load returns null`() {
+    val asyncCache = Caffeine.newBuilder().buildAsync<String, String>()
+    val loadingCache = LoadingCache(asyncCache) { null }
+
+    val value: String? = runBlocking { loadingCache.get("foo") }
+
+    assertThat(value).isNull()
+  }
+
+  @Test
+  fun `getValue throws NoSuchElementException when load returns null`() {
+    val asyncCache = Caffeine.newBuilder().buildAsync<String, String>()
+    val loadingCache = LoadingCache(asyncCache) { null }
+
+    assertFailsWith<NoSuchElementException> { runBlocking { loadingCache.getValue("foo") } }
+  }
+}