issue-36-new: Add LocalLockLimit to manage maximum entries in local l…

README.md

@@ -100,6 +100,44 @@ refresh dependencies with `./gradlew build --refresh-dependencies`.
 - The example modules contain working implementations for `RedisTemplate`, `ReactiveRedisTemplate` and the coroutine
   extensions.
 
+### Local lock map size
+
+- The local lock (`isLocalLock = true`, the default) keeps one entry per `(cache key, lock type)` it is currently
+  locking, in a map shared by the whole JVM. A cache key that is both created and refreshed therefore uses up to two
+  entries.
+- An entry is dropped as soon as the lock is released; an entry whose holder never released it is dropped by the
+  cleanup monitor once `lockTimeoutMillis` has passed, which runs on a 1 s interval and so can lag by up to that much.
+  The map's size therefore tracks how many keys are being locked **at the same time**, not how many distinct keys the
+  application sees. A service with 50 ms of supplier latency at 10k rps holds roughly 500 locks at once.
+- The map is uncapped by default. To bound it, set `req-shield.lock.max-entries`; `0` (the default) means uncapped.
+  Size it against concurrent lock ownership, not key cardinality.
+
+  ```yaml
+  # Spring modules: read from the Environment, so application.yml works
+  req-shield:
+    lock:
+      max-entries: 2000
+  ```
+
+  ```bash
+  # core / core-reactor / core-kotlin-coroutine used directly: system property
+  java -Dreq-shield.lock.max-entries=2000 -jar app.jar
+  ```
+
+  The Spring modules read the same key from the `Environment`, so a `-D` override still outranks the yml entry there,
+  and `REQ_SHIELD_LOCK_MAX_ENTRIES` works as an environment variable.
+- Once the map is full, a request for a **new** key is handed a permit that no map entry backs. It then runs exactly as
+  if it had taken the lock - it calls the supplier and writes the cache - so what the cap costs is request collapsing
+  for that key, and nothing else: no added latency, and the cache still gets populated. Keys whose entry is already in
+  the map are not subject to the cap at all.
+- The cap is a soft one. The size check is an estimate and is not atomic with the insertion it guards, so concurrent
+  callers can push the map slightly past the configured number.
+- The first refusal and every 1000th after it are logged at WARN.
+- A value that cannot be read as a non-negative number is logged at WARN and ignored, leaving the current cap
+  unchanged. This is deliberate in both directions: a class initializer that throws would poison the library for the
+  whole JVM, and an application context should not fail to start over a tuning knob typo. Setting the cap
+  programmatically (`LocalLockLimit.maxEntries = -1`) still fails fast, because there the stack trace is actionable.
+
 ### Cache eviction semantics
 
 - `@ReqShieldCacheEvict` evicts **after** the annotated method completes successfully (Spring's `@CacheEvict` default).

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

@@ -16,6 +16,7 @@
 
 package com.linecorp.cse.reqshield.kotlin.coroutine
 
+import com.linecorp.cse.reqshield.support.config.LocalLockLimit
 import com.linecorp.cse.reqshield.support.constant.ConfigValues.LOCK_MONITOR_INTERVAL_MILLIS
 import com.linecorp.cse.reqshield.support.utils.nowToEpochTime
 import kotlinx.coroutines.CancellationException
@@ -55,9 +56,11 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock, CoroutineScop
          */
         val isHeld: AtomicBoolean = AtomicBoolean(false),
         /**
-         * Token of the current owner, or null when the lock is not held.
+         * Token of the current owner.
          * Only the owner that presents this exact token may release the lock, so a holder whose
          * lock already expired and was reacquired by somebody else cannot release the new owner.
+         * It is null only while the entry is being built inside compute(): unLock and the monitor
+         * drop the whole entry instead of clearing the token, so every entry in the map is held.
          */
         @Volatile var token: String? = null,
     )
@@ -171,6 +174,15 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock, CoroutineScop
                 }
                 existing
             } else {
+                // At the cap, hand out a permit that no map entry backs instead of refusing.
+                // The caller then takes its normal path and writes the cache; making it wait
+                // would stall it on a holder that does not exist. Only this key's collapsing
+                // is lost, and the later unLock simply finds nothing to release.
+                if (LocalLockLimit.rejectsNewEntry(lockMap.mappingCount())) {
+                    acquiredToken.set(nextToken())
+                    return@compute null
+                }
+
                 // New entry: create and acquire
                 val token = nextToken()
                 val newLock = LockInfo(Semaphore(1), now + lockTimeoutMillis, token = token)
@@ -192,16 +204,19 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock, CoroutineScop
         val released = AtomicBoolean(false)
 
         // Release inside compute() so that it is atomic with respect to acquisition and cleanup
-        // of the same key: the token check and the semaphore release cannot be interleaved with a
-        // reacquisition. The entry itself is kept so a waiting caller can still acquire it.
+        // of the same key: the token check, the semaphore release and the removal cannot be
+        // interleaved with a reacquisition.
         lockMap.compute(completeKey) { _, existing ->
             if (existing == null) return@compute null
 
             // Only the current owner may release: a stale token belongs to an expired holder.
             if (existing.token == token && existing.isHeld.compareAndSet(true, false)) {
-                existing.token = null
                 existing.semaphore.release()
                 released.set(true)
+                // Drop the entry immediately rather than leaving it for the monitor. No reference
+                // to it escapes this lambda, and the next tryLock() just creates a fresh one, so
+                // keeping it would only pin the key until its lockTimeoutMillis runs out.
+                return@compute null
             }
             existing
         }

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

@@ -17,6 +17,8 @@
 package com.linecorp.cse.reqshield.kotlin.coroutine
 
 import com.linecorp.cse.reqshield.support.BaseKeyLockTest
+import com.linecorp.cse.reqshield.support.config.LocalLockLimit
+import com.linecorp.cse.reqshield.support.constant.ConfigValues.UNLIMITED_LOCK_ENTRIES
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
@@ -393,5 +395,114 @@ class KeyLocalLockTest : BaseKeyLockTest {
             keyLock.cancel()
         }
 
+    /**
+     * The lock map is private to the companion, so reflection is the only way to observe that
+     * unLock really drops the entry rather than leaving it for the expiry monitor.
+     */
+    @Suppress("UNCHECKED_CAST")
+    private fun readLockMap(): Map<String, Any> {
+        // The companion's private val is compiled as a static field on the outer class.
+        val field = KeyLocalLock::class.java.getDeclaredField("lockMap")
+        field.isAccessible = true
+        return field.get(null) as Map<String, Any>
+    }
+
+    @Test
+    fun `unLock removes the map entry right away instead of leaving it until expiry`() =
+        runBlocking {
+            // A long lock timeout rules the monitor out as the remover: were the entry kept by
+            // unLock, it would still be in the map for the next 60 seconds.
+            val keyLock = KeyLocalLock(60_000L)
+            val key = "unlock-removal-test-${java.util.UUID.randomUUID()}"
+            val mapKey = lockKeyOf(key, LockType.CREATE)
+
+            val token = assertNotNull(keyLock.tryLock(key, LockType.CREATE))
+            assertTrue(readLockMap().containsKey(mapKey), "A held lock must have an entry in the map")
+
+            assertTrue(keyLock.unLock(key, LockType.CREATE, token))
+            assertFalse(readLockMap().containsKey(mapKey), "unLock must drop the entry, not wait for the monitor")
+
+            keyLock.cancel()
+        }
+
+    @Test
+    fun `at the cap tryLock still grants a permit but stops adding map entries`() =
+        runBlocking {
+            val keyLock = KeyLocalLock(60_000L)
+            val first = "cap-test-first-${java.util.UUID.randomUUID()}"
+            val second = "cap-test-second-${java.util.UUID.randomUUID()}"
+
+            // Taken while still uncapped, so it always succeeds. Its 60s timeout keeps the entry
+            // in the map for the rest of the test, which is what makes a cap of one deterministic
+            // here: the map can only grow from this point, never shrink below one.
+            val firstToken = assertNotNull(keyLock.tryLock(first, LockType.CREATE))
+            assertTrue(readLockMap().containsKey(lockKeyOf(first, LockType.CREATE)))
+
+            LocalLockLimit.maxEntries = 1
+            try {
+                val secondToken =
+                    assertNotNull(
+                        keyLock.tryLock(second, LockType.CREATE),
+                        "past the cap the caller must still get a permit, so it writes the cache " +
+                            "instead of waiting for a holder that does not exist",
+                    )
+                assertFalse(
+                    readLockMap().containsKey(lockKeyOf(second, LockType.CREATE)),
+                    "a permit handed out past the cap must not add a map entry",
+                )
+                assertFalse(
+                    keyLock.unLock(second, LockType.CREATE, secondToken),
+                    "the permit is backed by no entry, so releasing it finds nothing",
+                )
+
+                // Losing collapsing for that key is exactly what the cap costs.
+                assertNotNull(
+                    keyLock.tryLock(second, LockType.CREATE),
+                    "past the cap a second caller for the same key is not collapsed either",
+                )
+            } finally {
+                LocalLockLimit.maxEntries = UNLIMITED_LOCK_ENTRIES
+            }
+
+            keyLock.unLock(first, LockType.CREATE, firstToken)
+            keyLock.cancel()
+        }
+
+    /**
+     * Only this module and the reactor one can pin this: reaching the existing-entry branch needs
+     * an expired entry that nothing sweeps, which means stopping the monitor. The core module has
+     * no equivalent hook, so the same case stays uncovered there.
+     */
+    @Test
+    fun `an entry already in the map is reacquired even when the map is at its cap`() =
+        runBlocking {
+            // Construct both locks first - each constructor restarts the monitor - then stop it,
+            // so the expired entry below survives for the rest of the test.
+            val shortLock = KeyLocalLock(1L)
+            val longLock = KeyLocalLock(60_000L)
+            KeyLocalLock.stopMonitoring()
+
+            val key = "cap-existing-${java.util.UUID.randomUUID()}"
+            assertNotNull(shortLock.tryLock(key, LockType.CREATE))
+            delay(20L) // the lock has timed out, and with the monitor stopped nothing removes it
+
+            LocalLockLimit.maxEntries = readLockMap().size.toLong() // exactly full
+            try {
+                val reacquired =
+                    assertNotNull(
+                        longLock.tryLock(key, LockType.CREATE),
+                        "an entry already in the map must not be subject to the cap",
+                    )
+                // A real lock, not a cap permit: it excludes the next caller and releases cleanly.
+                assertNull(longLock.tryLock(key, LockType.CREATE))
+                assertTrue(longLock.unLock(key, LockType.CREATE, reacquired))
+            } finally {
+                LocalLockLimit.maxEntries = UNLIMITED_LOCK_ENTRIES
+            }
+
+            shortLock.cancel()
+            longLock.cancel()
+        }
+
     private suspend fun doWork() = delay(1000)
 }

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

@@ -16,6 +16,7 @@
 
 package com.linecorp.cse.reqshield.reactor
 
+import com.linecorp.cse.reqshield.support.config.LocalLockLimit
 import com.linecorp.cse.reqshield.support.constant.ConfigValues.LOCK_KEY_PREFIX
 import com.linecorp.cse.reqshield.support.constant.ConfigValues.LOCK_MONITOR_INTERVAL_MILLIS
 import com.linecorp.cse.reqshield.support.utils.nowToEpochTime
@@ -56,9 +57,11 @@ class KeyLocalLock(
          */
         val isHeld: AtomicBoolean = AtomicBoolean(false),
         /**
-         * Ownership token of the current holder, null when the lock is not held.
+         * Ownership token of the current holder.
          * Only the holder that owns this token may release the lock, so a holder whose
          * lock already expired cannot release the lock of the next holder.
+         * It is null only while the entry is being built inside compute(): unLock and the monitor
+         * drop the whole entry instead of clearing the token, so every entry in the map is held.
          */
         @Volatile var token: String? = null,
     )
@@ -186,6 +189,15 @@ class KeyLocalLock(
                     }
                     existing
                 } else {
+                    // At the cap, hand out a permit that no map entry backs instead of refusing.
+                    // The caller then takes its normal path and writes the cache; making it wait
+                    // would stall it on a holder that does not exist. Only this key's collapsing
+                    // is lost, and the later unLock simply finds nothing to release.
+                    if (LocalLockLimit.rejectsNewEntry(lockMap.mappingCount())) {
+                        acquiredToken.set(nextToken())
+                        return@compute null
+                    }
+
                     // New entry: create and acquire
                     val token = nextToken()
                     val newLock = LockInfo(Semaphore(1), now + lockTimeoutMillis)
@@ -209,19 +221,23 @@ class KeyLocalLock(
             val completeKey = completeKey(key, lockType)
             val released = AtomicBoolean(false)
 
-            // Release inside compute() so that it is atomic with acquisition and cleanup:
-            // no other thread can reacquire this key while the ownership check runs.
+            // Release inside compute() so that the ownership check, the release and the removal
+            // are atomic with acquisition and cleanup: no other thread can reacquire this key
+            // while the ownership check runs.
             lockMap.compute(completeKey) { _, existing ->
                 if (existing == null) return@compute null
 
                 if (existing.token == token && existing.isHeld.compareAndSet(true, false)) {
                     existing.semaphore.release()
-                    existing.token = null
                     released.set(true)
-                } else {
-                    log.debug("Attempted to unlock key '{}' without holding its current token", completeKey)
+                    // Drop the entry immediately rather than leaving it for the monitor. No
+                    // reference to it escapes this lambda, and the next tryLock() just creates a
+                    // fresh one, so keeping it would only pin the key until its timeout runs out.
+                    return@compute null
                 }
-                existing // Keep the entry
+
+                log.debug("Attempted to unlock key '{}' without holding its current token", completeKey)
+                existing
             }
             released.get()
         }