[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,
 ) : 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-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/config/LibAutoConfiguration.kt

@@ -16,11 +16,22 @@
 
 package com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.config
 
+import com.linecorp.cse.reqshield.support.config.ReqShieldAspectProperties
+import org.springframework.beans.factory.annotation.Value
+import org.springframework.context.annotation.Bean
 import org.springframework.context.annotation.ComponentScan
 import org.springframework.context.annotation.Configuration
 import org.springframework.context.annotation.EnableAspectJAutoProxy
 
 @Configuration
 @EnableAspectJAutoProxy
 @ComponentScan(basePackages = ["com.linecorp.cse"])
-open class LibAutoConfiguration
+open class LibAutoConfiguration {
+
+    @Bean
+    open fun reqShieldAspectProperties(
+        @Value("\${req-shield.aspect.cache-max-size:1000}") cacheMaxSize: Int
+    ): ReqShieldAspectProperties {
+        return ReqShieldAspectProperties(cacheMaxSize = cacheMaxSize)
+    }
+}

core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectTest.kt

@@ -20,6 +20,7 @@ import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.annotation.Req
 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.BaseReqShieldModuleSupportTest
+import com.linecorp.cse.reqshield.support.config.ReqShieldAspectProperties
 import com.linecorp.cse.reqshield.support.model.Product
 import com.linecorp.cse.reqshield.support.model.ReqShieldData
 import io.mockk.coEvery
@@ -51,8 +52,9 @@ private val log = LoggerFactory.getLogger(ReqShieldAspectTest::class.java)
 @OptIn(ExperimentalCoroutinesApi::class)
 class ReqShieldAspectTest : BaseReqShieldModuleSupportTest {
     private val asyncCache: AsyncCache<Product> = mockk()
+    private val aspectProperties: ReqShieldAspectProperties = ReqShieldAspectProperties(cacheMaxSize = 1000)
     private val joinPoint: ProceedingJoinPoint = mockk<ProceedingJoinPoint>()
-    private val reqShieldAspect: ReqShieldAspect<Product> = spyk(ReqShieldAspect(asyncCache))
+    private val reqShieldAspect: ReqShieldAspect<Product> = spyk(ReqShieldAspect(asyncCache, aspectProperties))
     private val targetObject = spyk(TestBean())
     private val argument = mapOf("x" to "paramX", "y" to "paramY")
     private val mockContinuation = mockk<Continuation<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,
 ) : 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-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/config/LibAutoConfiguration.kt

@@ -16,11 +16,22 @@
 
 package com.linecorp.cse.reqshield.spring.webflux.config
 
+import com.linecorp.cse.reqshield.support.config.ReqShieldAspectProperties
+import org.springframework.beans.factory.annotation.Value
+import org.springframework.context.annotation.Bean
 import org.springframework.context.annotation.ComponentScan
 import org.springframework.context.annotation.Configuration
 import org.springframework.context.annotation.EnableAspectJAutoProxy
 
 @Configuration
 @EnableAspectJAutoProxy
 @ComponentScan(basePackages = ["com.linecorp.cse"])
-open class LibAutoConfiguration
+open class LibAutoConfiguration {
+
+    @Bean
+    open fun reqShieldAspectProperties(
+        @Value("\${req-shield.aspect.cache-max-size:1000}") cacheMaxSize: Int
+    ): ReqShieldAspectProperties {
+        return ReqShieldAspectProperties(cacheMaxSize = cacheMaxSize)
+    }
+}

core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectTest.kt

@@ -20,6 +20,7 @@ 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.BaseReqShieldModuleSupportTest
+import com.linecorp.cse.reqshield.support.config.ReqShieldAspectProperties
 import com.linecorp.cse.reqshield.support.model.Product
 import com.linecorp.cse.reqshield.support.model.ReqShieldData
 import io.mockk.every
@@ -44,7 +45,7 @@ import kotlin.test.assertTrue
 class ReqShieldAspectTest : BaseReqShieldModuleSupportTest {
     private val asyncCache: AsyncCache<Product> = mockk()
     private val joinPoint = mockk<ProceedingJoinPoint>()
-    private val reqShieldAspect = spyk(ReqShieldAspect(asyncCache))
+    private val reqShieldAspect = spyk(ReqShieldAspect(asyncCache, ReqShieldAspectProperties()))
     private val targetObject = spyk(TestBean())
     private val argument = mapOf("x" to "paramX", "y" to "paramY")
 

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,
 ) : 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-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/config/LibAutoConfiguration.kt

@@ -16,11 +16,22 @@
 
 package com.linecorp.cse.reqshield.spring.config
 
+import com.linecorp.cse.reqshield.support.config.ReqShieldAspectProperties
+import org.springframework.beans.factory.annotation.Value
+import org.springframework.context.annotation.Bean
 import org.springframework.context.annotation.ComponentScan
 import org.springframework.context.annotation.Configuration
 import org.springframework.context.annotation.EnableAspectJAutoProxy
 
 @Configuration
 @EnableAspectJAutoProxy
 @ComponentScan(basePackages = ["com.linecorp.cse"])
-open class LibAutoConfiguration
+open class LibAutoConfiguration {
+
+    @Bean
+    open fun reqShieldAspectProperties(
+        @Value("\${req-shield.aspect.cache-max-size:1000}") cacheMaxSize: Int
+    ): ReqShieldAspectProperties {
+        return ReqShieldAspectProperties(cacheMaxSize = cacheMaxSize)
+    }
+}

core-spring/src/test/kotlin/aspect/ReqShieldAspectTest.kt

@@ -22,6 +22,7 @@ import com.linecorp.cse.reqshield.spring.aspect.ReqShieldAspect
 import com.linecorp.cse.reqshield.spring.cache.ReqShieldCache
 import com.linecorp.cse.reqshield.support.BaseReqShieldModuleSupportTest
 import com.linecorp.cse.reqshield.support.BaseReqShieldTest
+import com.linecorp.cse.reqshield.support.config.ReqShieldAspectProperties
 import com.linecorp.cse.reqshield.support.model.Product
 import com.linecorp.cse.reqshield.support.model.ReqShieldData
 import io.mockk.every
@@ -49,7 +50,7 @@ private val log = LoggerFactory.getLogger(ReqShieldAspectTest::class.java)
 class ReqShieldAspectTest : BaseReqShieldModuleSupportTest {
     private val reqShieldCache: ReqShieldCache<Product> = mockk()
     private val joinPoint = mockk<ProceedingJoinPoint>()
-    private val reqShieldAspect = spyk(ReqShieldAspect(reqShieldCache))
+    private val reqShieldAspect = spyk(ReqShieldAspect(reqShieldCache, ReqShieldAspectProperties()))
     private val targetObject = spyk(TestBean())
     private val argument = mapOf("x" to "paramX", "y" to "paramY")
 

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,36 @@
+/*
+ *  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:
+ * Configure in application.yml:
+ *    req-shield:
+ *      aspect:
+ *        cache-max-size: 500
+ *
+ * The Spring modules automatically register this as a Bean using @Value injection.
+ */
+data class ReqShieldAspectProperties(
+    /**
+     * Maximum size for aspect internal caches (keyGeneratorMap and reqShieldMap).
+     * Default: 1000
+     */
+    val cacheMaxSize: Int = 1000,
+)