NO-ISSUE Add tests for ReqShield configuration and handling of update locks

core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/config/ReqShieldConfiguration.kt

@@ -52,12 +52,7 @@ data class ReqShieldConfiguration<T>(
     val isLocalLock: Boolean = true,
     val lockTimeoutMillis: Long = DEFAULT_LOCK_TIMEOUT_MILLIS,
     val decisionForUpdate: Int = DEFAULT_DECISION_FOR_UPDATE,
-    val keyLock: KeyLock =
-        if (isLocalLock) {
-            KeyLocalLock(lockTimeoutMillis)
-        } else {
-            KeyGlobalLock(globalLockFunction!!, globalUnLockFunction!!, lockTimeoutMillis)
-        },
+    val keyLock: KeyLock = defaultKeyLock(isLocalLock, globalLockFunction, globalUnLockFunction, lockTimeoutMillis),
     val maxAttemptGetCache: Int = MAX_ATTEMPT_GET_CACHE,
     val reqShieldWorkMode: ReqShieldWorkMode = ReqShieldWorkMode.CREATE_AND_UPDATE_CACHE,
     /**
@@ -100,6 +95,28 @@ data class ReqShieldConfiguration<T>(
     }
 }
 
+/**
+ * Builds the [KeyLock] used when the caller does not pass one.
+ *
+ * A default parameter expression is evaluated before the init block, so the global lock functions
+ * must be validated here as well to report a missing one as an [IllegalArgumentException].
+ */
+private fun defaultKeyLock(
+    isLocalLock: Boolean,
+    globalLockFunction: (suspend (String, String, Long) -> Boolean)?,
+    globalUnLockFunction: (suspend (String, String) -> Boolean)?,
+    lockTimeoutMillis: Long,
+): KeyLock =
+    if (isLocalLock) {
+        KeyLocalLock(lockTimeoutMillis)
+    } else {
+        KeyGlobalLock(
+            requireNotNull(globalLockFunction) { ErrorCode.DOES_NOT_EXIST_GLOBAL_LOCK_FUNCTION.message },
+            requireNotNull(globalUnLockFunction) { ErrorCode.DOES_NOT_EXIST_GLOBAL_UNLOCK_FUNCTION.message },
+            lockTimeoutMillis,
+        )
+    }
+
 enum class ReqShieldWorkMode {
     CREATE_AND_UPDATE_CACHE,
     ONLY_CREATE_CACHE,

core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLocalLockTest.kt

@@ -119,7 +119,7 @@ class KeyLocalLockTest : BaseKeyLockTest {
     override fun testConcurrencyWithOneKey() =
         runBlocking {
             val keyLock = KeyLocalLock(lockTimeoutMillis)
-            val key = "myKey"
+            val key = "myKey-concurrency-one"
             val lockType = LockType.CREATE
             val lockAcquiredCount = AtomicInteger(0)
             val tasksCompletedCount = AtomicInteger(0)
@@ -199,7 +199,7 @@ class KeyLocalLockTest : BaseKeyLockTest {
     override fun testLockExpiration() =
         runBlocking {
             val keyLock = KeyLocalLock(lockTimeoutMillis)
-            val key = "myKey"
+            val key = "myKey-lock-expiration"
             val lockType = LockType.CREATE
 
             assertNotNull(keyLock.tryLock(key, lockType))

core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/ReqShieldTest.kt

@@ -615,6 +615,27 @@ class ReqShieldTest : BaseReqShieldTest {
             coVerify { callable() }
         }
 
+    @Test
+    fun testSetMethodCacheExistsAndTheUpdateTargetButUpdateLockNotAcquired() =
+        runTest {
+            val timeToLiveMillis: Long = 1000
+            val isolatedKey = "update-lock-not-acquired-${java.util.UUID.randomUUID()}"
+            val reqShieldData = updateTargetData(oldValue, timeToLiveMillis)
+
+            coEvery { cacheGetter.invoke(isolatedKey) } returns reqShieldData
+            coEvery { keyLock.tryLock(isolatedKey, LockType.UPDATE) } returns null
+
+            val result = reqShield.getAndSetReqShieldData(isolatedKey, callable, timeToLiveMillis)
+            // Flush any (wrongly) queued refresh work on the background scope before asserting absence.
+            awaitBackgroundWrites()
+
+            assertSame(reqShieldData, result)
+            coVerify { keyLock.tryLock(isolatedKey, LockType.UPDATE) }
+            coVerify(inverse = true) { keyLock.unLock(isolatedKey, LockType.UPDATE, any()) }
+            coVerify(inverse = true) { callable() }
+            coVerify(inverse = true) { cacheSetter.invoke(isolatedKey, any(), any()) }
+        }
+
     @Test
     override fun testSetMethodCacheExistsAndTheUpdateTargetOnlyCreateCache() =
         runTest {

core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/config/ReqShieldConfigurationTest.kt

@@ -0,0 +1,53 @@
+/*
+ *  Copyright 2024 LY Corporation
+ *
+ *  LY Corporation 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:
+ *
+ *  https://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.linecorp.cse.reqshield.kotlin.coroutine.config
+
+import com.linecorp.cse.reqshield.support.exception.code.ErrorCode
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertThrows
+
+class ReqShieldConfigurationTest {
+    @Test
+    fun testGlobalLockWithoutLockFunctionAndWithoutExplicitKeyLock() {
+        val exception =
+            assertThrows<IllegalArgumentException> {
+                ReqShieldConfiguration<String>(
+                    setCacheFunction = { _, _, _ -> true },
+                    getCacheFunction = { null },
+                    isLocalLock = false,
+                )
+            }
+
+        assertEquals(ErrorCode.DOES_NOT_EXIST_GLOBAL_LOCK_FUNCTION.message, exception.message)
+    }
+
+    @Test
+    fun testGlobalLockWithoutUnLockFunctionAndWithoutExplicitKeyLock() {
+        val exception =
+            assertThrows<IllegalArgumentException> {
+                ReqShieldConfiguration<String>(
+                    setCacheFunction = { _, _, _ -> true },
+                    getCacheFunction = { null },
+                    globalLockFunction = { _, _, _ -> true },
+                    isLocalLock = false,
+                )
+            }
+
+        assertEquals(ErrorCode.DOES_NOT_EXIST_GLOBAL_UNLOCK_FUNCTION.message, exception.message)
+    }
+}

core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/ReqShield.kt

@@ -286,7 +286,9 @@ class ReqShield<T>(
         getFunction: (String) -> Mono<ReqShieldData<T>?>,
         key: String,
     ): Mono<ReqShieldData<T>?> =
-        getFunction(key)
+        // Deferred so a client function that throws synchronously fails as an onError signal.
+        Mono
+            .defer { getFunction(key) }
             .onErrorMap { e -> ClientException(ErrorCode.GET_CACHE_ERROR, cause = e) }
 
     private fun executeSetCacheFunction(
@@ -296,7 +298,10 @@ class ReqShield<T>(
         lockType: LockType,
         token: String?,
     ): Mono<Boolean> =
-        setFunction(key, value, value.timeToLiveMillis)
+        // Deferred so a client function that throws synchronously fails as an onError signal,
+        // which keeps the lock release in doFinally reachable.
+        Mono
+            .defer { setFunction(key, value, value.timeToLiveMillis) }
             .onErrorMap { e -> ClientException(ErrorCode.SET_CACHE_ERROR, cause = e) }
             .doFinally {
                 // Only the holder of a token took a lock, so only it may release one.
@@ -328,8 +333,10 @@ class ReqShield<T>(
         lockType: LockType?,
         token: String?,
     ): Mono<T?> =
-        callable
-            .call()
+        // Deferred so a supplier that throws synchronously fails as an onError signal,
+        // which keeps the lock release below reachable.
+        Mono
+            .defer { callable.call() }
             .doOnError { _ ->
                 // Only the holder of a token took a lock, so only it may release one.
                 if (lockType != null && token != null) {

core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/config/ReqShieldConfiguration.kt

@@ -51,12 +51,7 @@ data class ReqShieldConfiguration<T>(
     val lockTimeoutMillis: Long = DEFAULT_LOCK_TIMEOUT_MILLIS,
     val scheduler: Scheduler = Schedulers.boundedElastic(),
     val decisionForUpdate: Int = DEFAULT_DECISION_FOR_UPDATE,
-    val keyLock: KeyLock =
-        if (isLocalLock) {
-            KeyLocalLock(lockTimeoutMillis)
-        } else {
-            KeyGlobalLock(globalLockFunction!!, globalUnLockFunction!!, lockTimeoutMillis)
-        },
+    val keyLock: KeyLock = defaultKeyLock(isLocalLock, globalLockFunction, globalUnLockFunction, lockTimeoutMillis),
     val maxAttemptGetCache: Int = MAX_ATTEMPT_GET_CACHE,
     val reqShieldWorkMode: ReqShieldWorkMode = ReqShieldWorkMode.CREATE_AND_UPDATE_CACHE,
 ) {
@@ -72,6 +67,28 @@ data class ReqShieldConfiguration<T>(
     }
 }
 
+/**
+ * Builds the [KeyLock] used when the caller does not pass one.
+ *
+ * A default parameter expression is evaluated before the init block, so the global lock functions
+ * must be validated here as well to report a missing one as an [IllegalArgumentException].
+ */
+private fun defaultKeyLock(
+    isLocalLock: Boolean,
+    globalLockFunction: ((String, String, Long) -> Mono<Boolean>)?,
+    globalUnLockFunction: ((String, String) -> Mono<Boolean>)?,
+    lockTimeoutMillis: Long,
+): KeyLock =
+    if (isLocalLock) {
+        KeyLocalLock(lockTimeoutMillis)
+    } else {
+        KeyGlobalLock(
+            requireNotNull(globalLockFunction) { ErrorCode.DOES_NOT_EXIST_GLOBAL_LOCK_FUNCTION.message },
+            requireNotNull(globalUnLockFunction) { ErrorCode.DOES_NOT_EXIST_GLOBAL_UNLOCK_FUNCTION.message },
+            lockTimeoutMillis,
+        )
+    }
+
 enum class ReqShieldWorkMode {
     CREATE_AND_UPDATE_CACHE,
     ONLY_CREATE_CACHE,

core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/ReqShieldSyncThrowingClientFunctionTest.kt

@@ -0,0 +1,166 @@
+/*
+ *  Copyright 2024 LY Corporation
+ *
+ *  LY Corporation 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:
+ *
+ *  https://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.linecorp.cse.reqshield.reactor
+
+import com.linecorp.cse.reqshield.reactor.config.ReqShieldConfiguration
+import com.linecorp.cse.reqshield.reactor.config.ReqShieldWorkMode
+import com.linecorp.cse.reqshield.support.exception.ClientException
+import com.linecorp.cse.reqshield.support.exception.code.ErrorCode
+import com.linecorp.cse.reqshield.support.model.Product
+import com.linecorp.cse.reqshield.support.model.ReqShieldData
+import com.linecorp.cse.reqshield.support.utils.nowToEpochTime
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import reactor.core.publisher.Mono
+import reactor.core.scheduler.Schedulers
+import reactor.test.StepVerifier
+import java.util.UUID
+import java.util.concurrent.Callable
+import java.util.concurrent.atomic.AtomicInteger
+import kotlin.test.assertNotNull
+
+/**
+ * Client functions may fail before they ever return a publisher (e.g. a `require` at the top of the
+ * annotated method). Such a failure must become an onError signal so that the cleanup already wired
+ * into the chain - lock release and error mapping - actually runs.
+ */
+class ReqShieldSyncThrowingClientFunctionTest {
+    private val value = Product("testValue", "testValue")
+    private val cachedValue = Product("oldTestValue", "oldTestValue")
+    private val timeToLiveMillis = 10000L
+
+    // Long enough that a leaked lock cannot be released by its own expiration during the test.
+    private val lockTimeoutMillis = 60000L
+
+    private val throwingSupplier =
+        Callable<Mono<Product?>> { throw IllegalStateException("supplier failed before returning a Mono") }
+
+    /** Unique per test: KeyLocalLock keeps its lock map in a companion object shared by the whole JVM. */
+    private fun isolatedKey(name: String) = "$name-${UUID.randomUUID()}"
+
+    /** Cached entry that has passed the decisionForUpdate threshold (90% of its TTL). */
+    private fun updateTargetReqShieldData(): ReqShieldData<Product> =
+        ReqShieldData(
+            cachedValue,
+            ReqShieldData.Status.NEW,
+            nowToEpochTime() - (timeToLiveMillis * 0.9).toLong(),
+            timeToLiveMillis,
+        )
+
+    @Test
+    fun `should release the update lock when the supplier throws synchronously`() {
+        val key = isolatedKey("sync-throwing-supplier-update")
+        val keyLock = KeyLocalLock(lockTimeoutMillis)
+        val cached = updateTargetReqShieldData()
+        val reqShield =
+            ReqShield(
+                ReqShieldConfiguration<Product>(
+                    setCacheFunction = { _, _, _ -> Mono.just(true) },
+                    getCacheFunction = { Mono.just(cached) },
+                    keyLock = keyLock,
+                    scheduler = Schedulers.immediate(),
+                ),
+            )
+
+        StepVerifier
+            .create(reqShield.getAndSetReqShieldData(key, throwingSupplier, timeToLiveMillis))
+            .expectNext(cached)
+            .verifyComplete()
+
+        assertNotNull(
+            keyLock.tryLock(key, LockType.UPDATE).block(),
+            "the failed background refresh must have released its update lock",
+        )
+    }
+
+    @Test
+    fun `should still serve the cached value in only create cache mode when the supplier throws synchronously`() {
+        val key = isolatedKey("sync-throwing-supplier-only-create")
+        val cached = updateTargetReqShieldData()
+        val reqShield =
+            ReqShield(
+                ReqShieldConfiguration<Product>(
+                    setCacheFunction = { _, _, _ -> Mono.just(true) },
+                    getCacheFunction = { Mono.just(cached) },
+                    keyLock = KeyLocalLock(lockTimeoutMillis),
+                    scheduler = Schedulers.immediate(),
+                    reqShieldWorkMode = ReqShieldWorkMode.ONLY_CREATE_CACHE,
+                ),
+            )
+
+        // The refresh is fire-and-forget: its failure must not reach the caller of a cache hit.
+        StepVerifier
+            .create(reqShield.getAndSetReqShieldData(key, throwingSupplier, timeToLiveMillis))
+            .expectNext(cached)
+            .verifyComplete()
+    }
+
+    @Test
+    fun `should map a synchronously throwing get cache function to a get cache client exception`() {
+        val key = isolatedKey("sync-throwing-get-cache")
+        val getCacheInvocations = AtomicInteger(0)
+        val reqShield =
+            ReqShield(
+                ReqShieldConfiguration<Product>(
+                    setCacheFunction = { _, _, _ -> Mono.just(true) },
+                    getCacheFunction = {
+                        getCacheInvocations.incrementAndGet()
+                        throw IllegalStateException("cache read failed before returning a Mono")
+                    },
+                    keyLock = KeyLocalLock(lockTimeoutMillis),
+                    scheduler = Schedulers.immediate(),
+                ),
+            )
+
+        val result = reqShield.getAndSetReqShieldData(key, Callable { Mono.just<Product?>(value) }, timeToLiveMillis)
+
+        assertEquals(0, getCacheInvocations.get(), "the returned Mono must not read the cache before subscription")
+
+        StepVerifier
+            .create(result)
+            .expectErrorMatches { it is ClientException && it.errorCode == ErrorCode.GET_CACHE_ERROR }
+            .verify()
+
+        assertEquals(1, getCacheInvocations.get())
+    }
+
+    @Test
+    fun `should return the supplier value and release the create lock when the set cache function throws synchronously`() {
+        val key = isolatedKey("sync-throwing-set-cache")
+        val keyLock = KeyLocalLock(lockTimeoutMillis)
+        val reqShield =
+            ReqShield(
+                ReqShieldConfiguration<Product>(
+                    setCacheFunction = { _, _, _ -> throw IllegalStateException("cache write failed before returning a Mono") },
+                    getCacheFunction = { Mono.empty() },
+                    keyLock = keyLock,
+                    scheduler = Schedulers.immediate(),
+                ),
+            )
+
+        // The cache write is fire-and-forget, so its failure must not fail the request either.
+        StepVerifier
+            .create(reqShield.getAndSetReqShieldData(key, Callable { Mono.just<Product?>(value) }, timeToLiveMillis))
+            .assertNext { assertEquals(value, it.value) }
+            .verifyComplete()
+
+        assertNotNull(
+            keyLock.tryLock(key, LockType.CREATE).block(),
+            "the failed cache write must have released its create lock",
+        )
+    }
+}