[ISSUE #36] Fix ConcurrentHashMap memory leak and improve aspect cache performance

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

@@ -69,10 +69,11 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock, CoroutineScop
         lockType: LockType,
     ): Boolean {
         val completeKey = "${key}_${lockType.name}"
-        val lockInfo = lockMap[completeKey]
-        lockInfo?.let {
-            it.semaphore.release()
-            lockMap.remove(completeKey)
+        lockMap.compute(completeKey) { _, existingLockInfo ->
+            existingLockInfo?.let {
+                it.semaphore.release()
+                null // Remove the entry
+            }
         }
         return true
     }

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

@@ -66,10 +66,11 @@ class KeyLocalLock(
         Mono
             .fromCallable {
                 val completeKey = "${key}_${lockType.name}"
-                val lockInfo = lockMap[completeKey]
-                lockInfo?.let {
-                    it.semaphore.release()
-                    lockMap.remove(completeKey)
+                lockMap.compute(completeKey) { _, existingLockInfo ->
+                    existingLockInfo?.let {
+                        it.semaphore.release()
+                        null // Remove the entry
+                    }
                 }
             }.thenReturn(true)
 }

core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspect.kt

@@ -21,6 +21,8 @@ import com.linecorp.cse.reqshield.kotlin.coroutine.config.ReqShieldConfiguration
 import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.annotation.ReqShieldCacheEvict
 import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.annotation.ReqShieldCacheable
 import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.cache.AsyncCache
+import com.linecorp.cse.reqshield.support.config.ReqShieldAspectProperties
+import com.linecorp.cse.reqshield.support.utils.LRUCache
 import kotlinx.coroutines.reactor.awaitSingleOrNull
 import org.aspectj.lang.ProceedingJoinPoint
 import org.aspectj.lang.annotation.Around
@@ -42,21 +44,21 @@ import org.springframework.util.StringUtils
 import org.springframework.util.function.SingletonSupplier
 import reactor.core.publisher.Mono
 import java.lang.reflect.Method
-import java.util.concurrent.ConcurrentHashMap
 import kotlin.coroutines.Continuation
 
 @Aspect
 @Component
 class ReqShieldAspect<T>(
     private val asyncCache: AsyncCache<T>,
+    private val aspectProperties: ReqShieldAspectProperties = ReqShieldAspectProperties(),
 ) : BeanFactoryAware {
     private lateinit var beanFactory: BeanFactory
     private val springVersion = SpringVersion.getVersion()
     private val spelParser = SpelExpressionParser()
     private var defaultKeyGenerator = SingletonSupplier.of<KeyGenerator> { SimpleKeyGenerator() }
 
-    private val keyGeneratorMap = ConcurrentHashMap<String, KeyGenerator>()
-    internal val reqShieldMap = ConcurrentHashMap<String, ReqShield<T>>()
+    private val keyGeneratorMap = LRUCache<String, KeyGenerator>(aspectProperties.cacheMaxSize)
+    internal val reqShieldMap = LRUCache<String, ReqShield<T>>(aspectProperties.cacheMaxSize)
 
     @Around("execution(@com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.annotation.* * *(.., kotlin.coroutines.Continuation))")
     fun aroundTargetCacheable(joinPoint: ProceedingJoinPoint): Any? {

core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspect.kt

@@ -21,6 +21,8 @@ import com.linecorp.cse.reqshield.reactor.config.ReqShieldConfiguration
 import com.linecorp.cse.reqshield.spring.webflux.annotation.ReqShieldCacheEvict
 import com.linecorp.cse.reqshield.spring.webflux.annotation.ReqShieldCacheable
 import com.linecorp.cse.reqshield.spring.webflux.cache.AsyncCache
+import com.linecorp.cse.reqshield.support.config.ReqShieldAspectProperties
+import com.linecorp.cse.reqshield.support.utils.LRUCache
 import org.aspectj.lang.ProceedingJoinPoint
 import org.aspectj.lang.annotation.Around
 import org.aspectj.lang.annotation.Aspect
@@ -40,19 +42,19 @@ import org.springframework.util.StringUtils
 import org.springframework.util.function.SingletonSupplier
 import reactor.core.publisher.Mono
 import java.lang.reflect.Method
-import java.util.concurrent.ConcurrentHashMap
 
 @Aspect
 @Component
 class ReqShieldAspect<T>(
     private val asyncCache: AsyncCache<T>,
+    private val aspectProperties: ReqShieldAspectProperties = ReqShieldAspectProperties(),
 ) : BeanFactoryAware {
     private lateinit var beanFactory: BeanFactory
     private val spelParser = SpelExpressionParser()
     private var defaultKeyGenerator = SingletonSupplier.of<KeyGenerator> { SimpleKeyGenerator() }
 
-    private val keyGeneratorMap = ConcurrentHashMap<String, KeyGenerator>()
-    internal val reqShieldMap = ConcurrentHashMap<String, ReqShield<T>>()
+    private val keyGeneratorMap = LRUCache<String, KeyGenerator>(aspectProperties.cacheMaxSize)
+    internal val reqShieldMap = LRUCache<String, ReqShield<T>>(aspectProperties.cacheMaxSize)
 
     @Around("@annotation(com.linecorp.cse.reqshield.spring.webflux.annotation.ReqShieldCacheable)")
     fun aroundTargetCacheable(joinPoint: ProceedingJoinPoint): Mono<Any?> {

core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/aspect/ReqShieldAspect.kt

@@ -21,6 +21,8 @@ import com.linecorp.cse.reqshield.config.ReqShieldConfiguration
 import com.linecorp.cse.reqshield.spring.annotation.ReqShieldCacheEvict
 import com.linecorp.cse.reqshield.spring.annotation.ReqShieldCacheable
 import com.linecorp.cse.reqshield.spring.cache.ReqShieldCache
+import com.linecorp.cse.reqshield.support.config.ReqShieldAspectProperties
+import com.linecorp.cse.reqshield.support.utils.LRUCache
 import org.aspectj.lang.ProceedingJoinPoint
 import org.aspectj.lang.annotation.Around
 import org.aspectj.lang.annotation.Aspect
@@ -39,19 +41,19 @@ import org.springframework.stereotype.Component
 import org.springframework.util.StringUtils
 import org.springframework.util.function.SingletonSupplier
 import java.lang.reflect.Method
-import java.util.concurrent.ConcurrentHashMap
 
 @Aspect
 @Component
 class ReqShieldAspect<T>(
     private val reqShieldCache: ReqShieldCache<T>,
+    private val aspectProperties: ReqShieldAspectProperties = ReqShieldAspectProperties(),
 ) : BeanFactoryAware {
     private lateinit var beanFactory: BeanFactory
     private val spelParser = SpelExpressionParser()
     private val defaultKeyGenerator = SingletonSupplier.of<KeyGenerator> { SimpleKeyGenerator() }
 
-    private val keyGeneratorMap = ConcurrentHashMap<String, KeyGenerator>()
-    internal val reqShieldMap = ConcurrentHashMap<String, ReqShield<T>>()
+    private val keyGeneratorMap = LRUCache<String, KeyGenerator>(aspectProperties.cacheMaxSize)
+    internal val reqShieldMap = LRUCache<String, ReqShield<T>>(aspectProperties.cacheMaxSize)
 
     @Around("@annotation(com.linecorp.cse.reqshield.spring.annotation.ReqShieldCacheable)")
     fun aroundReqShieldCacheable(joinPoint: ProceedingJoinPoint): Any? {

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

@@ -130,10 +130,11 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock {
         lockType: LockType,
     ): Boolean {
         val completeKey = "${key}_${lockType.name}"
-        val lockInfo = lockMap[completeKey]
-        lockInfo?.let {
-            it.semaphore.release()
-            lockMap.remove(completeKey)
+        lockMap.compute(completeKey) { _, existingLockInfo ->
+            existingLockInfo?.let {
+                it.semaphore.release()
+                null // Remove the entry
+            }
         }
         return true
     }

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

@@ -194,10 +194,10 @@ class KeyLocalLockTest : BaseKeyLockTest {
 
         // Then - Instance2 should not be able to acquire the same lock
         val lock2Result = instance2.tryLock(key, lockType)
-        
+
         assertTrue(lock1Result)
         assertTrue(!lock2Result, "Instance2 should not acquire lock held by Instance1")
-        
+
         // Cleanup
         instance1.unLock(key, lockType)
         instance1.shutdown()
@@ -208,7 +208,7 @@ class KeyLocalLockTest : BaseKeyLockTest {
     fun `should maintain request collapsing across multiple instances`() {
         // Given
         val instance1 = KeyLocalLock(lockTimeoutMillis)
-        val instance2 = KeyLocalLock(lockTimeoutMillis) 
+        val instance2 = KeyLocalLock(lockTimeoutMillis)
         val instance3 = KeyLocalLock(lockTimeoutMillis)
         val key = "collapsing-key"
         val lockType = LockType.CREATE
@@ -220,11 +220,12 @@ class KeyLocalLockTest : BaseKeyLockTest {
         // When - Multiple instances try to acquire the same lock concurrently
         repeat(3) { index ->
             executor.submit {
-                val instance = when (index) {
-                    0 -> instance1
-                    1 -> instance2
-                    else -> instance3
-                }
+                val instance =
+                    when (index) {
+                        0 -> instance1
+                        1 -> instance2
+                        else -> instance3
+                    }
                 attemptCount.incrementAndGet()
                 if (instance.tryLock(key, lockType)) {
                     successCount.incrementAndGet()
@@ -241,7 +242,7 @@ class KeyLocalLockTest : BaseKeyLockTest {
         // Then - Only one should succeed in acquiring the lock
         assertEquals(3, attemptCount.get())
         assertEquals(1, successCount.get(), "Only one instance should acquire the lock")
-        
+
         // Cleanup
         instance1.shutdown()
         instance2.shutdown()
@@ -259,11 +260,11 @@ class KeyLocalLockTest : BaseKeyLockTest {
         // When - Instance1 acquires lock, Instance2 unlocks
         assertTrue(instance1.tryLock(key, lockType))
         instance2.unLock(key, lockType) // Should work even from different instance
-        
+
         // Then - New lock acquisition should succeed
         val newLockResult = instance2.tryLock(key, lockType)
         assertTrue(newLockResult, "Should be able to acquire lock after global unlock")
-        
+
         // Cleanup
         instance2.unLock(key, lockType)
         instance1.shutdown()
@@ -305,16 +306,20 @@ class KeyLocalLockTest : BaseKeyLockTest {
 
         // Then - Verify thread safety and concurrent operations handling
         assertEquals(0, errors.get(), "No errors should occur during concurrent operations")
-        
+
         // Due to sequential nature of ThreadPool(10) and brief work duration (10ms),
         // multiple operations can succeed on the same key at different times
-        assertTrue(operations.get() >= 10, 
-            "At least one operation per key should succeed (minimum 10)")
-        assertTrue(operations.get() <= 50, 
-            "No more operations than total attempts should succeed (maximum 50)")
-
+        assertTrue(
+            operations.get() >= 10,
+            "At least one operation per key should succeed (minimum 10)",
+        )
+        assertTrue(
+            operations.get() <= 50,
+            "No more operations than total attempts should succeed (maximum 50)",
+        )
+
         println("Successful operations: ${operations.get()}/50 total attempts")
-        
+
         // Cleanup
         instances.forEach { it.shutdown() }
     }

support/src/main/kotlin/com/linecorp/cse/reqshield/support/config/ReqShieldAspectProperties.kt

@@ -0,0 +1,35 @@
+/*
+ *  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.support.config
+
+/**
+ * Configuration properties for ReqShield aspect internal caches.
+ *
+ * Usage in Spring Boot:
+ * 1. Add @EnableConfigurationProperties(ReqShieldAspectProperties::class) to your config
+ * 2. Configure in application.yml:
+ *    req-shield:
+ *      aspect:
+ *        cache-max-size: 500
+ */
+data class ReqShieldAspectProperties(
+    /**
+     * Maximum size for aspect internal caches (keyGeneratorMap and reqShieldMap).
+     * Default: 1000
+     */
+    val cacheMaxSize: Int = 1000,
+)

support/src/main/kotlin/com/linecorp/cse/reqshield/support/utils/LRUCache.kt

@@ -0,0 +1,72 @@
+/*
+ *  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.support.utils
+
+/**
+ * High-performance LRU Cache with O(1) operations.
+ * Based on LinkedHashMap with access-order for optimal LRU behavior.
+ *
+ * This implementation provides better performance than the previous AtomicLong-based version
+ * by leveraging LinkedHashMap's built-in LRU behavior.
+ */
+class LRUCache<K, V>(private val maxSize: Int) {
+    private val cache =
+        object : LinkedHashMap<K, V>(maxSize + 1, 0.75f, true) {
+            override fun removeEldestEntry(eldest: MutableMap.MutableEntry<K, V>?): Boolean {
+                return size > maxSize
+            }
+        }
+
+    @Synchronized
+    fun computeIfAbsent(
+        key: K,
+        mappingFunction: (K) -> V,
+    ): V {
+        return cache.computeIfAbsent(key, mappingFunction)
+    }
+
+    @Synchronized
+    operator fun get(key: K): V? {
+        return cache[key]
+    }
+
+    @Synchronized
+    fun put(
+        key: K,
+        value: V,
+    ): V? {
+        return cache.put(key, value)
+    }
+
+    @Synchronized
+    fun remove(key: K): V? {
+        return cache.remove(key)
+    }
+
+    @Synchronized
+    fun clear() {
+        cache.clear()
+    }
+
+    @Synchronized
+    fun keys(): Set<K> = LinkedHashSet(cache.keys)
+
+    @Synchronized
+    fun size(): Int = cache.size
+
+    val size: Int get() = cache.size
+}