From 651b8eff4bee81fad22ad0b0b8cf3068a1568f58 Mon Sep 17 00:00:00 2001 From: "kanghyun.yang" Date: Wed, 16 Sep 2026 17:38:11 +0900 Subject: [PATCH] issue #45: Refactor cache locking mechanism and update dependencies for improved performance --- CLAUDE.md | 19 +- README.md | 68 +- build.gradle.kts | 20 + .../kotlin/coroutine/KeyGlobalLock.kt | 19 +- .../kotlin/coroutine/KeyLocalLock.kt | 82 +- .../cse/reqshield/kotlin/coroutine/KeyLock.kt | 13 +- .../reqshield/kotlin/coroutine/ReqShield.kt | 190 ++--- .../config/ReqShieldConfiguration.kt | 50 +- .../kotlin/coroutine/KeyGlobalLockTest.kt | 71 +- .../kotlin/coroutine/KeyLocalLockTest.kt | 162 ++-- .../kotlin/coroutine/ReqShieldTest.kt | 721 ++++++++++++++---- .../cse/reqshield/reactor/KeyGlobalLock.kt | 28 +- .../cse/reqshield/reactor/KeyLocalLock.kt | 80 +- .../linecorp/cse/reqshield/reactor/KeyLock.kt | 15 +- .../cse/reqshield/reactor/ReqShield.kt | 176 +++-- .../reactor/config/ReqShieldConfiguration.kt | 18 +- .../reqshield/reactor/KeyGlobalLockTest.kt | 79 +- .../cse/reqshield/reactor/KeyLocalLockTest.kt | 150 ++-- .../cse/reqshield/reactor/ReqShieldTest.kt | 390 ++++++---- .../annotation/ReqShieldCacheEvict.kt | 3 - .../annotation/ReqShieldCacheable.kt | 9 +- .../coroutine/aspect/CoroutineExtension.kt | 45 -- .../coroutine/aspect/ReqShieldAspect.kt | 120 ++- .../kotlin/coroutine/cache/AsyncCache.kt | 28 - .../coroutine/cache/GlobalLockSupport.kt | 33 + .../coroutine/config/LibAutoConfiguration.kt | 44 +- .../coroutine/aspect/InMemoryAsyncCache.kt | 47 +- .../aspect/ReqShieldAspectIntegrationTest.kt | 76 +- .../ReqShieldAspectRedisIntegrationTest.kt | 78 +- .../coroutine/aspect/ReqShieldAspectTest.kt | 288 +++++-- .../webflux/annotation/ReqShieldCacheEvict.kt | 6 +- .../webflux/annotation/ReqShieldCacheable.kt | 9 +- .../spring/webflux/aspect/ReqShieldAspect.kt | 166 ++-- .../spring/webflux/cache/AsyncCache.kt | 28 - .../spring/webflux/cache/GlobalLockSupport.kt | 35 + .../webflux/config/LibAutoConfiguration.kt | 8 +- .../webflux/aspect/InMemoryAsyncCache.kt | 47 +- .../webflux/aspect/LocalOnlyAsyncCache.kt | 28 + .../spring/webflux/aspect/RedisAsyncCache.kt | 75 ++ .../aspect/ReqShieldAspectIntegrationTest.kt | 135 +++- .../ReqShieldAspectRedisIntegrationTest.kt | 82 +- .../webflux/aspect/ReqShieldAspectTest.kt | 231 ++++-- .../spring/annotation/ReqShieldCacheEvict.kt | 3 - .../spring/annotation/ReqShieldCacheable.kt | 9 +- .../spring/aspect/ReqShieldAspect.kt | 81 +- .../spring/cache/GlobalLockSupport.kt | 33 + .../reqshield/spring/cache/ReqShieldCache.kt | 28 - .../spring/config/LibAutoConfiguration.kt | 31 +- .../aspect/ReqShieldAspectIntegrationTest.kt | 153 ++++ .../test/kotlin/aspect/ReqShieldAspectTest.kt | 317 ++++++-- .../linecorp/cse/reqshield/KeyGlobalLock.kt | 33 +- .../linecorp/cse/reqshield/KeyLocalLock.kt | 71 +- .../com/linecorp/cse/reqshield/KeyLock.kt | 20 +- .../com/linecorp/cse/reqshield/ReqShield.kt | 186 +++-- .../config/ReqShieldConfiguration.kt | 42 +- .../cse/reqshield/KeyGlobalLockTest.kt | 65 +- .../cse/reqshield/KeyLocalLockShutdownTest.kt | 32 +- .../cse/reqshield/KeyLocalLockTest.kt | 199 +++-- .../linecorp/cse/reqshield/ReqShieldTest.kt | 409 ++++++---- libs.versions.toml | 2 +- .../build.gradle.kts | 4 + .../mvc/example/cache/ReqShieldCacheImpl.kt | 24 +- .../configuration/RedisConfiguration.kt | 8 +- .../example/service/CacheAnnotationTest.kt | 14 +- .../mvc/example/service/GlobalLockTest.kt | 52 ++ .../build.gradle.kts | 4 + .../webflux/example/cache/AsyncCacheImpl.kt | 28 +- .../configuration/RedisConfiguration.kt | 7 +- .../example/service/CacheAnnotationTest.kt | 14 +- .../build.gradle.kts | 4 + .../coroutine/example/cache/AsyncCacheImpl.kt | 33 +- .../configuration/RedisConfiguration.kt | 7 +- .../coroutine/example/CacheAnnotationTest.kt | 151 +++- .../cse/reqshield/cache/ReqShieldCacheImpl.kt | 24 +- .../configuration/RedisConfiguration.kt | 8 +- .../spring/service/CacheAnnotationTest.kt | 14 +- .../spring/service/GlobalLockTest.kt | 52 ++ .../webflux/example/cache/AsyncCacheImpl.kt | 28 +- .../configuration/RedisConfiguration.kt | 7 +- .../example/service/CacheAnnotationTest.kt | 14 +- .../coroutine/example/cache/AsyncCacheImpl.kt | 36 +- .../configuration/RedisConfiguration.kt | 7 +- .../example/service/CacheAnnotationTest.kt | 59 +- support/build.gradle.kts | 3 - .../support/constant/ConfigValues.kt | 11 +- .../support/exception/ClientException.kt | 19 +- .../support/exception/ClientExceptionTest.kt | 95 +++ .../support/model/ReqShieldDataTest.kt | 76 ++ .../reqshield/support/utils/TimeUtilsTest.kt | 40 + .../utils/support/CommonUtilsTest.kt | 36 + 90 files changed, 4841 insertions(+), 1724 deletions(-) create mode 100644 core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/cache/GlobalLockSupport.kt create mode 100644 core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/cache/GlobalLockSupport.kt create mode 100644 core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/LocalOnlyAsyncCache.kt create mode 100644 core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/RedisAsyncCache.kt create mode 100644 core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/cache/GlobalLockSupport.kt create mode 100644 core-spring/src/test/kotlin/aspect/ReqShieldAspectIntegrationTest.kt create mode 100644 req-shield-spring-boot3-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/service/GlobalLockTest.kt create mode 100644 req-shield-spring-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/service/GlobalLockTest.kt create mode 100644 support/src/test/kotlin/com/linecorp/cse/reqshield/support/exception/ClientExceptionTest.kt create mode 100644 support/src/test/kotlin/com/linecorp/cse/reqshield/support/model/ReqShieldDataTest.kt create mode 100644 support/src/test/kotlin/com/linecorp/cse/reqshield/support/utils/TimeUtilsTest.kt diff --git a/CLAUDE.md b/CLAUDE.md index abe76ee..3d0147e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -21,10 +21,11 @@ The library is organized into several core modules: ### Key Components 1. **ReqShield**: Main orchestrator that manages cache operations and request collapsing -2. **KeyLock**: Locking mechanism (local or global) to prevent concurrent cache operations +2. **KeyLock**: Locking mechanism (local or global) to prevent concurrent cache operations. `tryLock` returns an ownership token and `unLock` only releases when the token matches, so a slow holder cannot release someone else's lock 3. **ReqShieldConfiguration**: Configuration object that defines cache functions, locking behavior, and timeouts 4. **ReqShieldData**: Wrapper for cached data with metadata (creation time, TTL) -5. **Spring Aspects**: AOP-based implementations that provide annotation-driven caching +5. **Spring Aspects**: AOP-based implementations that provide annotation-driven caching. One `ReqShield` instance per annotated method; cache and lock keys are namespaced as `"{cacheName}::{key}"`; eviction happens after the method returns successfully +6. **GlobalLockSupport**: Optional interface a cache bean implements to enable `isLocalLock = false` (Redis `SET NX PX` + compare-and-delete). Without it, `isLocalLock = false` fails fast ### Design Patterns @@ -88,29 +89,31 @@ Contains shared: ### ReqShieldConfiguration Parameters - `isLocalLock`: Use local vs distributed locking (default: true) +- `globalLockFunction` / `globalUnLockFunction`: `(lockKey, token, ttlMillis) -> Boolean` / `(lockKey, token) -> Boolean`, required when `isLocalLock = false` +- `executor` (core) / `scheduler` (reactor) / `scope` (coroutine): where background cache writes run; defaults are shared, the Spring adapters expose them as `reqShieldExecutor` / `reqShieldScheduler` / `reqShieldCoroutineScope` beans - `lockTimeoutMillis`: Lock acquisition timeout (default: 3000ms) - `decisionForUpdate`: Percentage of TTL after which to trigger async cache refresh (default: 80) -- `maxAttemptGetCache`: Max retry attempts when waiting for cache (default: 60) +- `maxAttemptGetCache`: Max retry attempts when waiting for cache (default: 60, 50ms apart). Three consecutive cache-read failures while waiting fall back to the supplier immediately; supplier failures propagate as `ClientException(SUPPLIER_ERROR)` - `reqShieldWorkMode`: CREATE_AND_UPDATE_CACHE | ONLY_CREATE_CACHE | ONLY_UPDATE_CACHE ### Work Modes - **CREATE_AND_UPDATE_CACHE**: Full functionality (default) -- **ONLY_CREATE_CACHE**: Never updates existing cache entries -- **ONLY_UPDATE_CACHE**: Never creates new cache entries +- **ONLY_CREATE_CACHE**: Request collapsing (lock) applies only when a cache entry is being created; refreshes of existing entries still happen but without a lock, so every request past the `decisionForUpdate` threshold may trigger a refresh +- **ONLY_UPDATE_CACHE**: Request collapsing (lock) applies only when an existing entry is being refreshed; on a cache miss every request calls the supplier and writes the cache without a lock ## Testing Guidelines ### Test Infrastructure - Uses JUnit 5 platform - MockK for Kotlin mocking -- Testcontainers for integration tests (Redis) +- Testcontainers for integration tests (Redis); set `TEST_REDIS_HOST`/`TEST_REDIS_PORT` to use an external Redis instead (CI does this) - Awaitility for asynchronous testing - Separate test fixtures in `support` module ### Test Coverage Requirements -- **Minimum test coverage**: 80% must be maintained across all modules +- **Minimum test coverage**: 80% line coverage for every library module (example modules are exempt) - Coverage reports generated via `./gradlew jacocoTestReport` -- Coverage enforced through Jacoco plugin configuration +- Enforced by `jacocoTestCoverageVerification`, which runs as part of `./gradlew check` / `./gradlew build` ### Code Quality Requirements - **Lint validation**: All code must pass ktlint checks before completion diff --git a/README.md b/README.md index aaa56cd..b651e87 100644 --- a/README.md +++ b/README.md @@ -39,27 +39,59 @@ A lib that regulates the cache-based requests an application receives in terms o - `EMIT_EMPTY` (default): map `null` to `Mono.empty()`. - `ERROR`: throw an `IllegalStateException` if a `null` value is produced. -### Global lock guidance - -- When `isLocalLock = false`, you must provide real global lock/unlock implementations. -- Recommended approach with Redis: - - Lock: `SETNX lock:{key} 1` + `PEXPIRE lock:{key} {ttlMillis}` - - Unlock: `DEL lock:{key}` -- The provided defaults return `true` and are only suitable for local/dev usage. - -### Reactor Scheduler tuning +### Cache key layout -- Reactor-based modules accept a `Scheduler` (e.g., `boundedElastic`) through configuration. -- Spring WebFlux adapter exposes a `reqShieldScheduler` bean you can override for tuning thread usage. +- The Spring adapters store every entry under `"{cacheName}::{key}"` (the same convention as Spring's `RedisCacheManager`), + where `key` is the SpEL result or the `KeyGenerator` output. `@ReqShieldCacheEvict` applies the same rule, so an evict + with the same `cacheName` and `key` always targets the entry written by `@ReqShieldCacheable`. +- Lock keys are derived from that namespaced key and prefixed with `reqshield:lock:`, so two caches that happen to use the + same raw key never share a lock or an entry. +- If you call `ReqShield` directly (core / core-reactor / core-kotlin-coroutine), the key you pass is used as is. -### Kotlin Coroutine Parallelism Configuration - -| Property | Default | Description | -|----------|---------|-------------| -| `reqshield.blocking.parallelism` | `availableProcessors * 2` (clamped 4-256) | Controls parallelism for blocking calls in the coroutine aspect | +### Global lock guidance -**Note**: This feature uses `Dispatchers.IO.limitedParallelism()` which is marked as `@ExperimentalCoroutinesApi`. -The API may change in future Kotlin Coroutines versions. +- When `isLocalLock = false`, your cache bean must also implement the module's `GlobalLockSupport` interface + (`com.linecorp.cse.reqshield.spring.cache.GlobalLockSupport`, `...spring.webflux.cache.GlobalLockSupport`, + `...spring.webflux.kotlin.coroutine.cache.GlobalLockSupport`). If it does not, the first call to the annotated method + fails with an `IllegalArgumentException` instead of silently running without request collapsing. +- Every lock acquisition carries an ownership token. Only the holder that acquired the lock can release it, so a slow + holder whose lock already expired can no longer release the lock of the next holder. +- Recommended Redis implementation: + - Lock: `SET {lockKey} {token} NX PX {ttlMillis}` (atomic; never `SETNX` followed by a separate `PEXPIRE`) + - Unlock: compare-and-delete in a Lua script, e.g. + `if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end` +- The example modules contain working implementations for `RedisTemplate`, `ReactiveRedisTemplate` and the coroutine + extensions. + +### Cache eviction semantics + +- `@ReqShieldCacheEvict` evicts **after** the annotated method completes successfully (Spring's `@CacheEvict` default). + If the method throws, nothing is evicted; if the eviction itself fails, that failure propagates to the caller. +- In WebFlux the eviction also runs when the method completes empty (for example a `Mono` handler). + +### Waiting for another request (lock not acquired) + +- A request that loses the lock polls the cache every 50 ms, up to `maxAttemptGetCache` times (default 60). +- A cache read failure while polling is logged and counted as a failed attempt. Three consecutive failures are treated as a + cache outage and the request falls back to calling the supplier itself right away. +- When the attempts are exhausted the request calls the supplier once. A supplier failure is propagated as a + `ClientException(SUPPLIER_ERROR)` with the original exception as `cause`; it is never turned into a cached `null`. + +### Thread pools and schedulers + +- `core` uses a `ScheduledExecutorService`. The default is a shared daemon pool; the Spring adapter exposes it as the + `reqShieldExecutor` bean, which you can override. +- `core-reactor` accepts a `Scheduler` (default `boundedElastic`). The Spring WebFlux adapter exposes it as the + `reqShieldScheduler` bean. +- `core-kotlin-coroutine` accepts a `CoroutineScope` for background cache writes (default: a shared supervisor scope on + `Dispatchers.IO`). The coroutine Spring adapter exposes it as the `reqShieldCoroutineScope` bean and cancels it on + context shutdown. + +### Kotlin coroutine adapter + +- `@ReqShieldCacheable` / `@ReqShieldCacheEvict` from `core-spring-webflux-kotlin-coroutine` require `suspend` + functions. Annotating a regular function fails fast with an `IllegalArgumentException`; use `core-spring` or + `core-spring-webflux` for blocking or `Mono`-returning methods. ## Contributing diff --git a/build.gradle.kts b/build.gradle.kts index 3778660..eb56210 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -63,6 +63,26 @@ allprojects { } } + // Enforce the minimum line coverage documented in CLAUDE.md for library modules. + // Example applications (req-shield-*-example) are demos and are not held to the threshold. + if (project != rootProject && !project.name.startsWith("req-shield-")) { + tasks.withType { + dependsOn(tasks.test) + violationRules { + rule { + limit { + counter = "LINE" + value = "COVEREDRATIO" + minimum = "0.80".toBigDecimal() + } + } + } + } + tasks.named("check") { + dependsOn(tasks.withType()) + } + } + jacoco { toolVersion = "0.8.12" } diff --git a/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyGlobalLock.kt b/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyGlobalLock.kt index c03c373..55ed8d4 100644 --- a/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyGlobalLock.kt +++ b/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyGlobalLock.kt @@ -16,24 +16,25 @@ package com.linecorp.cse.reqshield.kotlin.coroutine +import java.util.UUID + class KeyGlobalLock( - private val globalLockFunction: suspend (String, Long) -> Boolean, - private val globalUnLockFunction: suspend (String) -> Boolean, + private val globalLockFunction: suspend (String, String, Long) -> Boolean, + private val globalUnLockFunction: suspend (String, String) -> Boolean, private val lockTimeoutMillis: Long, ) : KeyLock { override suspend fun tryLock( key: String, lockType: LockType, - ): Boolean { - val completeKey = "${key}_${lockType.name}" - return globalLockFunction(completeKey, lockTimeoutMillis) + ): String? { + // The token must be unique across every process sharing the lock store. + val token = UUID.randomUUID().toString() + return if (globalLockFunction(lockKeyOf(key, lockType), token, lockTimeoutMillis)) token else null } override suspend fun unLock( key: String, lockType: LockType, - ): Boolean { - val completeKey = "${key}_${lockType.name}" - return globalUnLockFunction(completeKey) - } + token: String, + ): Boolean = globalUnLockFunction(lockKeyOf(key, lockType), token) } diff --git a/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLocalLock.kt b/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLocalLock.kt index 370befd..9439fcc 100644 --- a/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLocalLock.kt +++ b/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLocalLock.kt @@ -18,6 +18,7 @@ package com.linecorp.cse.reqshield.kotlin.coroutine import com.linecorp.cse.reqshield.support.constant.ConfigValues.LOCK_MONITOR_INTERVAL_MILLIS import com.linecorp.cse.reqshield.support.utils.nowToEpochTime +import kotlinx.coroutines.CancellationException import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Job @@ -28,6 +29,8 @@ import org.slf4j.LoggerFactory import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.Semaphore import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.atomic.AtomicLong +import java.util.concurrent.atomic.AtomicReference import kotlin.coroutines.CoroutineContext private val log = LoggerFactory.getLogger(KeyLocalLock::class.java) @@ -51,14 +54,28 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock, CoroutineScop * check vs unLock, or monitor cleanup vs unLock). */ val isHeld: AtomicBoolean = AtomicBoolean(false), + /** + * Token of the current owner, or null when the lock is not held. + * 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. + */ + @Volatile var token: String? = null, ) companion object { private val lockMap = ConcurrentHashMap() + /** + * Source of local ownership tokens. A monotonic counter is enough because tokens never + * leave this JVM, and it is far cheaper than a UUID on the lock acquisition path. + */ + private val tokenSequence = AtomicLong() + @Volatile private var monitorJob: Job? = null + private fun nextToken(): String = "local-${tokenSequence.incrementAndGet()}" + private fun ensureMonitorStarted() { if (monitorJob?.isActive == true) return synchronized(this) { @@ -66,7 +83,7 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock, CoroutineScop monitorJob = CoroutineScope(Dispatchers.IO).launch { while (isActive) { - runCatching { + try { val now = System.currentTimeMillis() // Remove expired locks using compute() for atomic check-and-remove. // This prevents TOCTOU race condition where removeIf's lambda returns true @@ -82,6 +99,7 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock, CoroutineScop // This handles the case where unlock() was missed due to exception. // CAS ensures safe release (no-op if already released). if (lockInfo.isHeld.compareAndSet(true, false)) { + lockInfo.token = null lockInfo.semaphore.release() } null // Atomic removal @@ -90,10 +108,15 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock, CoroutineScop } } } - delay(LOCK_MONITOR_INTERVAL_MILLIS) - }.onFailure { e -> + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { log.error("Error in lock lifecycle monitoring: {}", e.message, e) } + + // Delay outside the try/catch: a failed sweep must still wait for the next + // interval instead of turning the monitor into a hot loop. + delay(LOCK_MONITOR_INTERVAL_MILLIS) } } } @@ -119,13 +142,13 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock, CoroutineScop override suspend fun tryLock( key: String, lockType: LockType, - ): Boolean { - val completeKey = "${key}_${lockType.name}" + ): String? { + val completeKey = lockKeyOf(key, lockType) val now = nowToEpochTime() - val result = AtomicBoolean(false) + val acquiredToken = AtomicReference(null) // Use compute() for atomic lock acquisition. - // This ensures mutual exclusion with cleanup - they cannot race on the same key. + // This ensures mutual exclusion with cleanup and unLock - they cannot race on the same key. lockMap.compute(completeKey) { _, existing -> if (existing != null) { // Force-release expired locks to allow reacquisition. @@ -133,43 +156,60 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock, CoroutineScop // Without CAS, if unLock() executes between isHeld.get() and release(), // both threads would call release(), causing over-release (permits > 1). if (now > existing.expiresAt && existing.isHeld.compareAndSet(true, false)) { + // Drop the token as well so the timed-out owner cannot release the next one. + existing.token = null existing.semaphore.release() } // Existing entry: try to acquire semaphore if (existing.semaphore.tryAcquire()) { + val token = nextToken() existing.isHeld.set(true) + existing.token = token existing.expiresAt = now + lockTimeoutMillis - result.set(true) + acquiredToken.set(token) } existing } else { // New entry: create and acquire - val newLock = LockInfo(Semaphore(1), now + lockTimeoutMillis) + val token = nextToken() + val newLock = LockInfo(Semaphore(1), now + lockTimeoutMillis, token = token) newLock.semaphore.tryAcquire() // Always succeeds for new semaphore newLock.isHeld.set(true) - result.set(true) + acquiredToken.set(token) newLock } } - return result.get() + return acquiredToken.get() } override suspend fun unLock( key: String, lockType: LockType, + token: String, ): Boolean { - val completeKey = "${key}_${lockType.name}" - val lockInfo = lockMap[completeKey] ?: return false - - // Use CAS to prevent over-release: only release if we actually hold the lock - return if (lockInfo.isHeld.compareAndSet(true, false)) { - lockInfo.semaphore.release() - true - } else { - log.debug("Attempted to unlock key '{}' that is not held", completeKey) - false + val completeKey = lockKeyOf(key, lockType) + 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. + 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) + } + existing + } + + if (!released.get()) { + log.debug("Attempted to unlock key '{}' that is not held by token '{}'", completeKey, token) } + return released.get() } fun cancel() { diff --git a/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLock.kt b/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLock.kt index 7206451..396c9ac 100644 --- a/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLock.kt +++ b/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLock.kt @@ -16,15 +16,20 @@ package com.linecorp.cse.reqshield.kotlin.coroutine +import com.linecorp.cse.reqshield.support.constant.ConfigValues.LOCK_KEY_PREFIX + interface KeyLock { + /** Returns an opaque ownership token when acquired, or null when another holder owns the lock. */ suspend fun tryLock( key: String, lockType: LockType, - ): Boolean + ): String? + /** Releases only if [token] matches the current owner; false when not held or token mismatch. */ suspend fun unLock( key: String, lockType: LockType, + token: String, ): Boolean } @@ -32,3 +37,9 @@ enum class LockType { CREATE, UPDATE, } + +/** Lock keys are prefixed so that a lock entry can never collide with a cache entry. */ +internal fun lockKeyOf( + key: String, + lockType: LockType, +): String = "$LOCK_KEY_PREFIX${key}_${lockType.name}" diff --git a/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/ReqShield.kt b/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/ReqShield.kt index d1568c0..4bb625c 100644 --- a/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/ReqShield.kt +++ b/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/ReqShield.kt @@ -19,19 +19,23 @@ package com.linecorp.cse.reqshield.kotlin.coroutine import com.linecorp.cse.reqshield.kotlin.coroutine.config.ReqShieldConfiguration import com.linecorp.cse.reqshield.kotlin.coroutine.config.ReqShieldWorkMode import com.linecorp.cse.reqshield.support.constant.ConfigValues.GET_CACHE_INTERVAL_MILLIS -import com.linecorp.cse.reqshield.support.constant.ConfigValues.MAX_ATTEMPT_SET_CACHE -import com.linecorp.cse.reqshield.support.constant.ConfigValues.SET_CACHE_RETRY_INTERVAL_MILLIS +import com.linecorp.cse.reqshield.support.constant.ConfigValues.MAX_CONSECUTIVE_GET_CACHE_FAILURES import com.linecorp.cse.reqshield.support.exception.ClientException import com.linecorp.cse.reqshield.support.exception.code.ErrorCode import com.linecorp.cse.reqshield.support.model.ReqShieldData import com.linecorp.cse.reqshield.support.utils.decideToUpdateCache -import kotlinx.coroutines.CoroutineScope -import kotlinx.coroutines.Deferred -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.async +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CoroutineStart +import kotlinx.coroutines.ExperimentalCoroutinesApi +import kotlinx.coroutines.NonCancellable +import kotlinx.coroutines.currentCoroutineContext import kotlinx.coroutines.delay +import kotlinx.coroutines.ensureActive import kotlinx.coroutines.launch -import java.util.concurrent.atomic.AtomicInteger +import kotlinx.coroutines.withContext +import org.slf4j.LoggerFactory + +private val log = LoggerFactory.getLogger(ReqShield::class.java) class ReqShield( private val reqShieldConfig: ReqShieldConfiguration, @@ -55,34 +59,33 @@ class ReqShield( private fun shouldUpdateCache(reqShieldData: ReqShieldData): Boolean = decideToUpdateCache(reqShieldData.createdAt, reqShieldData.timeToLiveMillis, reqShieldConfig.decisionForUpdate) + @OptIn(ExperimentalCoroutinesApi::class) private suspend fun updateReqShieldData( key: String, callable: suspend () -> T?, timeToLiveMillis: Long, ) { val lockType = LockType.UPDATE + // ONLY_CREATE_CACHE collapses requests on creation only, so the update runs without a lock. + val onlyCreateCache = reqShieldConfig.reqShieldWorkMode == ReqShieldWorkMode.ONLY_CREATE_CACHE + val token = if (onlyCreateCache) null else reqShieldConfig.keyLock.tryLock(key, lockType) + + if (!onlyCreateCache && token == null) return - fun executeAsyncTask() { - CoroutineScope(Dispatchers.IO).launch { + reqShieldConfig.scope.launch(start = CoroutineStart.ATOMIC) { + try { val reqShieldData = buildReqShieldData( - executeCallable({ callable() }, true, key, lockType), + executeCallable(callable, key, lockType, token), timeToLiveMillis, ) - setReqShieldData( - reqShieldConfig.setCacheFunction, - key, - reqShieldData, - lockType, - ) + executeSetCacheFunction(reqShieldConfig.setCacheFunction, key, reqShieldData, lockType, token) + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { + log.error("[Req-Shield] failed to update the cache of key '{}'", key, e) } } - - if (reqShieldConfig.reqShieldWorkMode == ReqShieldWorkMode.ONLY_CREATE_CACHE || - reqShieldConfig.keyLock.tryLock(key, lockType) - ) { - return executeAsyncTask() - } } private suspend fun handleLockForCacheCreation( @@ -91,43 +94,79 @@ class ReqShield( timeToLiveMillis: Long, ): ReqShieldData { val lockType = LockType.CREATE + // ONLY_UPDATE_CACHE collapses requests on update only, so the creation runs without a lock. + val onlyUpdateCache = reqShieldConfig.reqShieldWorkMode == ReqShieldWorkMode.ONLY_UPDATE_CACHE + val token = if (onlyUpdateCache) null else reqShieldConfig.keyLock.tryLock(key, lockType) - return if (reqShieldConfig.reqShieldWorkMode == ReqShieldWorkMode.ONLY_UPDATE_CACHE || - reqShieldConfig.keyLock.tryLock(key, lockType) - ) { - createReqShieldData(key, callable, timeToLiveMillis, lockType) + return if (onlyUpdateCache || token != null) { + createReqShieldData(key, callable, timeToLiveMillis, lockType, token) } else { handleLockFailure(key, callable, timeToLiveMillis) } } + @OptIn(ExperimentalCoroutinesApi::class) private suspend fun createReqShieldData( key: String, callable: suspend () -> T?, timeToLiveMillis: Long, lockType: LockType, + token: String?, ): ReqShieldData { val reqShieldData = buildReqShieldData( - executeCallable({ callable() }, true, key, lockType), + executeCallable(callable, key, lockType, token), timeToLiveMillis, ) - CoroutineScope(Dispatchers.IO).launch { - setReqShieldData(reqShieldConfig.setCacheFunction, key, reqShieldData, lockType) + reqShieldConfig.scope.launch(start = CoroutineStart.ATOMIC) { + try { + executeSetCacheFunction(reqShieldConfig.setCacheFunction, key, reqShieldData, lockType, token) + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { + log.error("[Req-Shield] failed to create the cache of key '{}'", key, e) + } } return reqShieldData } + /** + * Another request owns the create lock, so wait for it to fill the cache. + * + * Polling runs on the caller's coroutine: [delay] is a cancellation point, so a cancelled caller + * stops polling immediately. Consecutive cache-read failures are treated as a cache outage and + * end the wait early, after which the supplier is called on this coroutine as the last resort. + */ private suspend fun handleLockFailure( key: String, callable: suspend () -> T?, timeToLiveMillis: Long, ): ReqShieldData { - val counter = createCounter() + var attempts = 0 + var consecutiveFailures = 0 + + while (attempts < reqShieldConfig.maxAttemptGetCache) { + val cachedData = + try { + executeGetCacheFunction(reqShieldConfig.getCacheFunction, key) + .also { consecutiveFailures = 0 } + } catch (e: CancellationException) { + throw e + } catch (e: ClientException) { + log.warn("[Req-Shield] failed to read the cache of key '{}' while waiting", key, e) + if (++consecutiveFailures >= MAX_CONSECUTIVE_GET_CACHE_FAILURES) break + null + } + + if (cachedData != null) return cachedData - val result = scheduleTask(counter, reqShieldConfig.getCacheFunction, callable, key).await() + attempts++ + delay(GET_CACHE_INTERVAL_MILLIS) + } - return buildReqShieldData(result, timeToLiveMillis) + // The other request never filled the cache: fall back to the supplier. No lock was acquired + // here, so there is nothing to release, and a supplier failure propagates as SUPPLIER_ERROR. + return buildReqShieldData(executeCallable(callable, key, null, null), timeToLiveMillis) } private fun buildReqShieldData( @@ -139,42 +178,16 @@ class ReqShield( timeToLiveMillis = timeToLiveMillis, ) - private suspend fun setReqShieldData( - cacheSetter: suspend (String, ReqShieldData, Long) -> Boolean, - key: String, - reqShieldData: ReqShieldData, - lockType: LockType, - ) { - executeSetCacheFunction(cacheSetter, key, reqShieldData, lockType) - } - - private fun createCounter(): AtomicInteger = AtomicInteger(0) - - private fun scheduleTask( - counter: AtomicInteger, - cacheGetter: suspend (String) -> ReqShieldData?, - callable: suspend () -> T?, - key: String, - ): Deferred = - CoroutineScope(Dispatchers.IO).async { - while (counter.incrementAndGet() <= reqShieldConfig.maxAttemptGetCache) { - executeGetCacheFunction(cacheGetter, key)?.let { - return@async it.value - } - delay(GET_CACHE_INTERVAL_MILLIS) - } - - return@async executeCallable({ callable() }, false) - } - private suspend fun executeGetCacheFunction( getFunction: suspend (String) -> ReqShieldData?, key: String, ): ReqShieldData? = - runCatching { + try { getFunction(key) - }.getOrElse { - throw ClientException(ErrorCode.GET_CACHE_ERROR, originErrorMessage = it.message) + } catch (e: CancellationException) { + throw e + } catch (e: Exception) { + throw ClientException(ErrorCode.GET_CACHE_ERROR, cause = e) } private suspend fun executeSetCacheFunction( @@ -182,47 +195,42 @@ class ReqShield( key: String, value: ReqShieldData, lockType: LockType, + token: String?, ) { try { + currentCoroutineContext().ensureActive() setFunction(key, value, value.timeToLiveMillis) + } catch (e: CancellationException) { + throw e } catch (e: Exception) { - throw ClientException(ErrorCode.SET_CACHE_ERROR, originErrorMessage = e.message) + throw ClientException(ErrorCode.SET_CACHE_ERROR, cause = e) } finally { - if (shouldAttemptUnlock(lockType)) { - unlockWithRetry(key, lockType) - } - } - } - - private suspend fun unlockWithRetry( - key: String, - lockType: LockType, - ) { - repeat(MAX_ATTEMPT_SET_CACHE) { attempt -> - if (reqShieldConfig.keyLock.unLock(key, lockType)) { - return - } else if (attempt < MAX_ATTEMPT_SET_CACHE - 1) { - delay(SET_CACHE_RETRY_INTERVAL_MILLIS) + // Only the holder of a token has a lock to release. + if (token != null) { + withContext(NonCancellable) { + reqShieldConfig.keyLock.unLock(key, lockType, token) + } } } } private suspend fun executeCallable( callable: suspend () -> T?, - isUnlockWhenException: Boolean, - key: String? = null, - lockType: LockType? = null, + key: String, + lockType: LockType?, + token: String?, ): T? = - runCatching { + try { + currentCoroutineContext().ensureActive() callable() - }.getOrElse { - if (isUnlockWhenException && key != null && lockType != null) { - reqShieldConfig.keyLock.unLock(key, lockType) + } catch (e: Exception) { + // Release the lock only when this call actually acquired one. + if (token != null && lockType != null) { + withContext(NonCancellable) { + reqShieldConfig.keyLock.unLock(key, lockType, token) + } } - throw ClientException(ErrorCode.SUPPLIER_ERROR, originErrorMessage = it.message) + if (e is CancellationException) throw e + throw ClientException(ErrorCode.SUPPLIER_ERROR, cause = e) } - - private fun shouldAttemptUnlock(lockType: LockType): Boolean = - (lockType == LockType.UPDATE && reqShieldConfig.reqShieldWorkMode != ReqShieldWorkMode.ONLY_CREATE_CACHE) || - (lockType == LockType.CREATE && reqShieldConfig.reqShieldWorkMode != ReqShieldWorkMode.ONLY_UPDATE_CACHE) } diff --git a/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/config/ReqShieldConfiguration.kt b/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/config/ReqShieldConfiguration.kt index 9d900ae..b310775 100644 --- a/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/config/ReqShieldConfiguration.kt +++ b/core-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/config/ReqShieldConfiguration.kt @@ -24,12 +24,31 @@ import com.linecorp.cse.reqshield.support.constant.ConfigValues.DEFAULT_LOCK_TIM import com.linecorp.cse.reqshield.support.constant.ConfigValues.MAX_ATTEMPT_GET_CACHE import com.linecorp.cse.reqshield.support.exception.code.ErrorCode import com.linecorp.cse.reqshield.support.model.ReqShieldData +import kotlinx.coroutines.CoroutineExceptionHandler +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.SupervisorJob +import org.slf4j.LoggerFactory data class ReqShieldConfiguration( val setCacheFunction: suspend (String, ReqShieldData, Long) -> Boolean, val getCacheFunction: suspend (String) -> ReqShieldData?, - val globalLockFunction: (suspend (String, Long) -> Boolean)? = null, - val globalUnLockFunction: (suspend (String) -> Boolean)? = null, + /** + * Acquires the global lock. Invoked with (lockKey, token, ttlMillis) and must return true only + * when this caller is the one that acquired the lock. + * + * Recommended Redis implementation: `SET lockKey token NX PX ttlMillis`. + */ + val globalLockFunction: (suspend (String, String, Long) -> Boolean)? = null, + /** + * Releases the global lock. Invoked with (lockKey, token) and must release the lock only when + * the stored value is still equal to the given token (compare-and-delete), so an owner whose + * lock already expired cannot release the lock of the next owner. + * + * Recommended Redis implementation: a Lua script that compares `GET lockKey` with the token and + * deletes the key only on a match. + */ + val globalUnLockFunction: (suspend (String, String) -> Boolean)? = null, val isLocalLock: Boolean = true, val lockTimeoutMillis: Long = DEFAULT_LOCK_TIMEOUT_MILLIS, val decisionForUpdate: Int = DEFAULT_DECISION_FOR_UPDATE, @@ -41,6 +60,14 @@ data class ReqShieldConfiguration( }, val maxAttemptGetCache: Int = MAX_ATTEMPT_GET_CACHE, val reqShieldWorkMode: ReqShieldWorkMode = ReqShieldWorkMode.CREATE_AND_UPDATE_CACHE, + /** + * Scope that runs every fire-and-forget cache write. + * + * The default is a process-wide shared scope, so creating a configuration never leaks a Job. + * Callers that need lifecycle control (e.g. draining pending writes on shutdown) should pass + * their own scope and cancel it themselves. + */ + val scope: CoroutineScope = defaultScope, ) { init { if (!isLocalLock) { @@ -52,6 +79,25 @@ data class ReqShieldConfiguration( } } } + + companion object { + private val log = LoggerFactory.getLogger(ReqShieldConfiguration::class.java) + + /** + * Shared scope for background cache writes, created on first use. + * + * SupervisorJob keeps one failed write from cancelling the others, and the exception handler + * is the last-resort backstop for anything the write path did not already log. + */ + private val defaultScope: CoroutineScope by lazy { + CoroutineScope( + SupervisorJob() + Dispatchers.IO + + CoroutineExceptionHandler { _, e -> + log.error("[Req-Shield] background task failed", e) + }, + ) + } + } } enum class ReqShieldWorkMode { diff --git a/core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyGlobalLockTest.kt b/core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyGlobalLockTest.kt index 2d7ab72..ba5b406 100644 --- a/core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyGlobalLockTest.kt +++ b/core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyGlobalLockTest.kt @@ -19,6 +19,8 @@ package com.linecorp.cse.reqshield.kotlin.coroutine import com.linecorp.cse.reqshield.support.BaseKeyLockTest import com.linecorp.cse.reqshield.support.redis.AbstractRedisTest import io.lettuce.core.RedisClient +import io.lettuce.core.ScriptOutputType +import io.lettuce.core.SetArgs import io.lettuce.core.api.async.RedisAsyncCommands import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.delay @@ -28,18 +30,27 @@ import kotlinx.coroutines.launch import kotlinx.coroutines.runBlocking import kotlinx.coroutines.withContext import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertFalse import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test import java.util.concurrent.atomic.AtomicInteger -import kotlin.test.Ignore +import kotlin.test.assertNotNull +import kotlin.test.assertNull + +/** + * Releases the lock only when the stored value still equals the token presented by the caller. + * This is the compare-and-delete the global unlock function is documented to implement. + */ +private const val UNLOCK_SCRIPT = + "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end" class KeyGlobalLockTest : AbstractRedisTest(), BaseKeyLockTest { private lateinit var redisCommands: RedisAsyncCommands - private lateinit var globalLockFunc: suspend (String, Long) -> Boolean - private lateinit var globalUnLockFunc: suspend (String) -> Boolean + private lateinit var globalLockFunc: suspend (String, String, Long) -> Boolean + private lateinit var globalUnLockFunc: suspend (String, String) -> Boolean @BeforeEach fun init() { @@ -53,13 +64,18 @@ class KeyGlobalLockTest : // Clean up all keys from previous tests for proper test isolation connection.sync().flushdb() - globalLockFunc = { key, timeToLiveMillis -> - redisCommands.setnx(key, key).toCompletableFuture().await() + globalLockFunc = { key, token, lockTimeoutMillis -> + redisCommands + .set(key, token, SetArgs.Builder.nx().px(lockTimeoutMillis)) + .toCompletableFuture() + .await() == "OK" } - globalUnLockFunc = { key -> - redisCommands.del(key).toCompletableFuture().await() - true + globalUnLockFunc = { key, token -> + redisCommands + .eval(UNLOCK_SCRIPT, ScriptOutputType.INTEGER, arrayOf(key), token) + .toCompletableFuture() + .await() == 1L } } @@ -76,7 +92,8 @@ class KeyGlobalLockTest : List(20) { launch { withContext(Dispatchers.IO) { - if (keyLock.tryLock(key, lockType)) { + val token = keyLock.tryLock(key, lockType) + if (token != null) { try { println("${Thread.currentThread().name} acquired the lock") lockAcquiredCount.incrementAndGet() @@ -84,7 +101,7 @@ class KeyGlobalLockTest : } catch (e: InterruptedException) { e.printStackTrace() } finally { - keyLock.unLock(key, lockType) + assertTrue(keyLock.unLock(key, lockType, token)) println("${Thread.currentThread().name} released the lock") } } else { @@ -101,7 +118,7 @@ class KeyGlobalLockTest : delay(100) - assertTrue(keyLock.tryLock(key, lockType)) + assertTrue(keyLock.tryLock(key, lockType) != null, "The lock must be free again") } @Test @@ -117,7 +134,8 @@ class KeyGlobalLockTest : val key = if (i % 2 == 0) "myKey1" else "myKey2" launch { withContext(Dispatchers.IO) { - if (keyLock.tryLock(key, lockType)) { + val token = keyLock.tryLock(key, lockType) + if (token != null) { try { println("${Thread.currentThread().name} acquired the lock") lockAcquiredCount.incrementAndGet() @@ -125,7 +143,7 @@ class KeyGlobalLockTest : } catch (e: InterruptedException) { e.printStackTrace() } finally { - keyLock.unLock(key, lockType) + assertTrue(keyLock.unLock(key, lockType, token)) println("${Thread.currentThread().name} released the lock") } } else { @@ -141,15 +159,34 @@ class KeyGlobalLockTest : delay(100) - assertTrue(keyLock.tryLock("myKey1", lockType)) - assertTrue(keyLock.tryLock("myKey2", lockType)) + assertTrue(keyLock.tryLock("myKey1", lockType) != null) + assertTrue(keyLock.tryLock("myKey2", lockType) != null) } @Test - @Ignore override fun testLockExpiration() = runBlocking { - // Global locks do not have an expiration + // The lock is written with `SET ... NX PX lockTimeoutMillis`, so Redis expires it for us. + val shortLockTimeout = 300L + val keyLock = KeyGlobalLock(globalLockFunc, globalUnLockFunc, shortLockTimeout) + val key = "expiring-key" + val lockType = LockType.CREATE + + val expiredToken = keyLock.tryLock(key, lockType) + assertNotNull(expiredToken) + assertNull(keyLock.tryLock(key, lockType), "The lock must stay held until it expires") + + delay(shortLockTimeout + 200L) + + val newToken = keyLock.tryLock(key, lockType) + assertNotNull(newToken, "The lock must be reacquirable once it expired") + + // The previous owner must not be able to release the lock of the new owner. + assertFalse(keyLock.unLock(key, lockType, expiredToken)) + assertNull(keyLock.tryLock(key, lockType), "The new owner must still hold the lock") + + assertTrue(keyLock.unLock(key, lockType, newToken)) + assertTrue(keyLock.tryLock(key, lockType) != null, "The released lock must be acquirable again") } private suspend fun doWork() = delay(100) diff --git a/core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLocalLockTest.kt b/core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLocalLockTest.kt index 8a4af79..6c3d70c 100644 --- a/core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLocalLockTest.kt +++ b/core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/KeyLocalLockTest.kt @@ -30,7 +30,10 @@ import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.Assertions.assertFalse import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.Test +import java.util.concurrent.ConcurrentLinkedQueue import java.util.concurrent.atomic.AtomicInteger +import kotlin.test.assertNotNull +import kotlin.test.assertNull class KeyLocalLockTest : BaseKeyLockTest { @AfterEach @@ -47,26 +50,69 @@ class KeyLocalLockTest : BaseKeyLockTest { val key = "shared-key" val lockType = LockType.CREATE - assertTrue(instance1.tryLock(key, lockType)) - assertTrue(!instance2.tryLock(key, lockType)) + val token = instance1.tryLock(key, lockType) + assertNotNull(token) + assertNull(instance2.tryLock(key, lockType)) - instance1.unLock(key, lockType) + // Any instance may release the lock as long as it presents the owning token. + assertTrue(instance2.unLock(key, lockType, token)) } @Test fun `should maintain request collapsing across multiple instances`() = runBlocking { - val instance1 = KeyLocalLock(lockTimeoutMillis) - val instance2 = KeyLocalLock(lockTimeoutMillis) - val instance3 = KeyLocalLock(lockTimeoutMillis) + val instances = List(3) { KeyLocalLock(lockTimeoutMillis) } val key = "collapsing-key" val lockType = LockType.CREATE - val acquired = listOf(instance1, instance2, instance3).map { it.tryLock(key, lockType) }.count { it } - assertEquals(1, acquired) + val tokens = instances.mapNotNull { it.tryLock(key, lockType) } + assertEquals(1, tokens.size) // cleanup whoever acquired - listOf(instance1, instance2, instance3).forEach { it.unLock(key, lockType) } + tokens.forEach { instances.first().unLock(key, lockType, it) } + } + + @Test + fun `should reject an unlock presenting a token of another owner`() = + runBlocking { + val keyLock = KeyLocalLock(lockTimeoutMillis) + val key = "foreign-token-test" + val lockType = LockType.CREATE + + val token = keyLock.tryLock(key, lockType) + assertNotNull(token) + + assertFalse(keyLock.unLock(key, lockType, "someone-elses-token"), "A foreign token must not release the lock") + // The lock is still held, so nobody else can acquire it. + assertNull(keyLock.tryLock(key, lockType)) + + assertTrue(keyLock.unLock(key, lockType, token), "The owning token must release the lock") + keyLock.cancel() + } + + @Test + fun `should not let a stale token release the lock of the next owner`() = + runBlocking { + val shortLockTimeout = 50L + val keyLock = KeyLocalLock(shortLockTimeout) + val key = "stale-token-test" + val lockType = LockType.CREATE + + val staleToken = keyLock.tryLock(key, lockType) + assertNotNull(staleToken) + + // Let the lock time out so that the next caller can force-acquire it. + delay(shortLockTimeout + 10L) + val newToken = keyLock.tryLock(key, lockType) + assertNotNull(newToken) + assertTrue(staleToken != newToken) + + // The timed-out owner must not be able to release the new owner's lock. + assertFalse(keyLock.unLock(key, lockType, staleToken)) + assertNull(keyLock.tryLock(key, lockType), "The new owner must still hold the lock") + + assertTrue(keyLock.unLock(key, lockType, newToken)) + keyLock.cancel() } @Test @@ -82,7 +128,8 @@ class KeyLocalLockTest : BaseKeyLockTest { List(20) { launch { withContext(Dispatchers.IO) { - if (keyLock.tryLock(key, lockType)) { + val token = keyLock.tryLock(key, lockType) + if (token != null) { try { println("${Thread.currentThread().name} acquired the lock") lockAcquiredCount.incrementAndGet() @@ -90,7 +137,7 @@ class KeyLocalLockTest : BaseKeyLockTest { } catch (e: InterruptedException) { e.printStackTrace() } finally { - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, token) println("${Thread.currentThread().name} released the lock") } } else { @@ -107,7 +154,7 @@ class KeyLocalLockTest : BaseKeyLockTest { delay(100) - assertTrue(keyLock.tryLock(key, lockType)) + assertTrue(keyLock.tryLock(key, lockType) != null, "The lock must be free again") } @Test @@ -123,14 +170,15 @@ class KeyLocalLockTest : BaseKeyLockTest { val key = if (i % 2 == 0) "myKey1" else "myKey2" launch { withContext(Dispatchers.IO) { - if (keyLock.tryLock(key, lockType)) { + val token = keyLock.tryLock(key, lockType) + if (token != null) { try { lockAcquiredCount.incrementAndGet() doWork() } catch (e: InterruptedException) { e.printStackTrace() } finally { - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, token) } } tasksCompletedCount.incrementAndGet() @@ -143,8 +191,8 @@ class KeyLocalLockTest : BaseKeyLockTest { delay(100) - assertTrue(keyLock.tryLock("myKey1", lockType)) - assertTrue(keyLock.tryLock("myKey2", lockType)) + assertTrue(keyLock.tryLock("myKey1", lockType) != null) + assertTrue(keyLock.tryLock("myKey2", lockType) != null) } @Test @@ -154,19 +202,19 @@ class KeyLocalLockTest : BaseKeyLockTest { val key = "myKey" val lockType = LockType.CREATE - assertTrue(keyLock.tryLock(key, lockType)) + assertNotNull(keyLock.tryLock(key, lockType)) // Wait for lock timeout + cleanup interval + buffer // lockTimeoutMillis = 3000ms, cleanup interval = 1000ms delay(lockTimeoutMillis + 1000L + 500L) // 4.5 seconds total - val result = + val token = withContext(Dispatchers.IO) { keyLock.tryLock(key, lockType) } - assertTrue(result) - assertTrue(keyLock.unLock(key, lockType)) + assertNotNull(token) + assertTrue(keyLock.unLock(key, lockType, token)) } @Test @@ -177,20 +225,22 @@ class KeyLocalLockTest : BaseKeyLockTest { val lockType = LockType.CREATE // Acquire lock - assertTrue(keyLock.tryLock(key, lockType)) + val token = keyLock.tryLock(key, lockType) + assertNotNull(token) // First unlock should succeed - assertTrue(keyLock.unLock(key, lockType), "First unlock should succeed") + assertTrue(keyLock.unLock(key, lockType, token), "First unlock should succeed") // Second unlock should return false (over-release prevention) - assertFalse(keyLock.unLock(key, lockType), "Second unlock should fail (over-release prevention)") + assertFalse(keyLock.unLock(key, lockType, token), "Second unlock should fail (over-release prevention)") // Verify semaphore is not over-released: can acquire once, not twice - assertTrue(keyLock.tryLock(key, lockType), "Should acquire lock after proper unlock") - assertFalse(keyLock.tryLock(key, lockType), "Should not acquire lock twice (semaphore intact)") + val reacquiredToken = keyLock.tryLock(key, lockType) + assertNotNull(reacquiredToken, "Should acquire lock after proper unlock") + assertNull(keyLock.tryLock(key, lockType), "Should not acquire lock twice (semaphore intact)") // Cleanup - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, reacquiredToken) keyLock.cancel() } @@ -200,31 +250,30 @@ class KeyLocalLockTest : BaseKeyLockTest { val keyLock = KeyLocalLock(lockTimeoutMillis) val key = "concurrent-over-release-test" val lockType = LockType.CREATE - val successfulAcquisitions = AtomicInteger(0) + val acquiredTokens = ConcurrentLinkedQueue() // Simulate over-release attempt - assertTrue(keyLock.tryLock(key, lockType)) - keyLock.unLock(key, lockType) + val token = keyLock.tryLock(key, lockType) + assertNotNull(token) + assertTrue(keyLock.unLock(key, lockType, token)) // Multiple unlock attempts should all return false (not over-release) - repeat(5) { assertFalse(keyLock.unLock(key, lockType)) } + repeat(5) { assertFalse(keyLock.unLock(key, lockType, token)) } // Try to acquire lock concurrently - only ONE should succeed val attempts = (1..10).map { async(Dispatchers.IO) { - if (keyLock.tryLock(key, lockType)) { - successfulAcquisitions.incrementAndGet() - } + keyLock.tryLock(key, lockType)?.let { acquiredTokens.add(it) } } } attempts.awaitAll() // Only one should have acquired the lock - assertEquals(1, successfulAcquisitions.get(), "Only one should acquire the lock") + assertEquals(1, acquiredTokens.size, "Only one should acquire the lock") // Cleanup - keyLock.unLock(key, lockType) + acquiredTokens.forEach { keyLock.unLock(key, lockType, it) } keyLock.cancel() } @@ -239,7 +288,8 @@ class KeyLocalLockTest : BaseKeyLockTest { repeat(100) { iteration -> // Step 1: Acquire lock - assertTrue(keyLock.tryLock(key, lockType), "Iteration $iteration: Initial lock should succeed") + val token = keyLock.tryLock(key, lockType) + assertNotNull(token, "Iteration $iteration: Initial lock should succeed") // Step 2: Wait for lock to expire (but not be cleaned up by monitor) delay(shortLockTimeout + 10L) @@ -254,34 +304,33 @@ class KeyLocalLockTest : BaseKeyLockTest { } val unLockResult = async(Dispatchers.IO) { - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, token) } - tryLockResult.await() + val raceToken = tryLockResult.await() unLockResult.await() // Step 4: Verify no over-release by checking lock behavior // If over-release occurred, permits > 1, allowing multiple acquisitions - val acquisitions = AtomicInteger(0) + val acquiredTokens = ConcurrentLinkedQueue() val attempts = (1..5).map { async(Dispatchers.IO) { - if (keyLock.tryLock(key, lockType)) { - acquisitions.incrementAndGet() - } + keyLock.tryLock(key, lockType)?.let { acquiredTokens.add(it) } } } attempts.awaitAll() // At most 1 should succeed (0 if tryLock already holds it, 1 if it released) assertTrue( - acquisitions.get() <= 1, + acquiredTokens.size <= 1, "Iteration $iteration: Over-release detected! " + - "Expected at most 1 acquisition, got ${acquisitions.get()}", + "Expected at most 1 acquisition, got ${acquiredTokens.size}", ) // Cleanup for next iteration - repeat(3) { keyLock.unLock(key, lockType) } + raceToken?.let { keyLock.unLock(key, lockType, it) } + acquiredTokens.forEach { keyLock.unLock(key, lockType, it) } } keyLock.cancel() @@ -298,38 +347,41 @@ class KeyLocalLockTest : BaseKeyLockTest { repeat(50) { iteration -> // Acquire lock and let it expire - assertTrue(keyLock.tryLock(key, lockType)) + val expiredToken = keyLock.tryLock(key, lockType) + assertNotNull(expiredToken, "Iteration $iteration: Initial lock should succeed") delay(shortLockTimeout + 5L) - // High contention: many concurrent tryLock and unLock calls + // High contention: many concurrent tryLock and unLock calls. + // Releases present whichever token is currently known - a stale one must be rejected. + val liveTokens = ConcurrentLinkedQueue() val jobs = (1..20).map { i -> if (i % 2 == 0) { - async(Dispatchers.IO) { keyLock.tryLock(key, lockType) } + async(Dispatchers.IO) { keyLock.tryLock(key, lockType)?.let { liveTokens.add(it) } } } else { - async(Dispatchers.IO) { keyLock.unLock(key, lockType) } + async(Dispatchers.IO) { + keyLock.unLock(key, lockType, liveTokens.poll() ?: expiredToken) + } } } jobs.awaitAll() // Verify: try to acquire lock multiple times concurrently - val acquisitions = AtomicInteger(0) + val acquiredTokens = ConcurrentLinkedQueue() val verifyJobs = (1..10).map { async(Dispatchers.IO) { - if (keyLock.tryLock(key, lockType)) { - acquisitions.incrementAndGet() - } + keyLock.tryLock(key, lockType)?.let { acquiredTokens.add(it) } } } verifyJobs.awaitAll() - if (acquisitions.get() > 1) { + if (acquiredTokens.size > 1) { overReleaseDetected.incrementAndGet() } // Cleanup - repeat(15) { keyLock.unLock(key, lockType) } + (liveTokens + acquiredTokens).forEach { keyLock.unLock(key, lockType, it) } } assertEquals( diff --git a/core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/ReqShieldTest.kt b/core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/ReqShieldTest.kt index e2ada9b..9ca0716 100644 --- a/core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/ReqShieldTest.kt +++ b/core-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/kotlin/coroutine/ReqShieldTest.kt @@ -19,21 +19,34 @@ package com.linecorp.cse.reqshield.kotlin.coroutine import com.linecorp.cse.reqshield.kotlin.coroutine.config.ReqShieldConfiguration import com.linecorp.cse.reqshield.kotlin.coroutine.config.ReqShieldWorkMode import com.linecorp.cse.reqshield.support.BaseReqShieldTest +import com.linecorp.cse.reqshield.support.constant.ConfigValues.GET_CACHE_INTERVAL_MILLIS +import com.linecorp.cse.reqshield.support.constant.ConfigValues.MAX_ATTEMPT_GET_CACHE +import com.linecorp.cse.reqshield.support.constant.ConfigValues.MAX_CONSECUTIVE_GET_CACHE_FAILURES 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 io.mockk.coEvery import io.mockk.coVerify -import io.mockk.every import io.mockk.mockk -import io.mockk.mockkStatic -import io.mockk.unmockkStatic +import io.mockk.slot +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.ExperimentalCoroutinesApi +import kotlinx.coroutines.SupervisorJob +import kotlinx.coroutines.TimeoutCancellationException +import kotlinx.coroutines.awaitCancellation +import kotlinx.coroutines.cancel +import kotlinx.coroutines.cancelAndJoin import kotlinx.coroutines.delay +import kotlinx.coroutines.job +import kotlinx.coroutines.launch import kotlinx.coroutines.runBlocking import kotlinx.coroutines.test.runTest -import kotlinx.coroutines.withTimeoutOrNull +import kotlinx.coroutines.withTimeout import org.junit.jupiter.api.AfterEach import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.Assertions.assertNotNull @@ -42,13 +55,16 @@ import org.junit.jupiter.api.Test import org.junit.jupiter.api.assertThrows import java.lang.reflect.InvocationTargetException import java.lang.reflect.Method -import java.time.LocalDateTime +import java.util.concurrent.atomic.AtomicInteger import kotlin.coroutines.Continuation import kotlin.coroutines.EmptyCoroutineContext import kotlin.test.assertFailsWith import kotlin.test.assertNull +import kotlin.test.assertSame import kotlin.test.assertTrue +private const val LOCAL_TOKEN = "local-token" + @OptIn(ExperimentalCoroutinesApi::class) class ReqShieldTest : BaseReqShieldTest { private lateinit var reqShield: ReqShield @@ -60,56 +76,37 @@ class ReqShieldTest : BaseReqShieldTest { private lateinit var cacheGetter: suspend (String) -> ReqShieldData? private lateinit var keyLock: KeyLock private lateinit var keyGlobalLock: KeyLock + + /** Runs the fire-and-forget cache writes, so tests can await them instead of sleeping. */ + private lateinit var backgroundScope: CoroutineScope private val key = "testKey" + private val createLockKey = lockKeyOf(key, LockType.CREATE) + private val updateLockKey = lockKeyOf(key, LockType.UPDATE) private val oldValue = Product("oldTestValue", "oldTestName") private val value = Product("testValue", "testName") private val callable: suspend () -> Product? = mockk() private var timeToLiveMillis: Long = 10000 - private lateinit var globalLockFunc: suspend (String, Long) -> Boolean - private lateinit var globalUnLockFunc: suspend (String) -> Boolean + private lateinit var globalLockFunc: suspend (String, String, Long) -> Boolean + private lateinit var globalUnLockFunc: suspend (String, String) -> Boolean @BeforeEach fun setup() { cacheSetter = mockk, Long) -> Boolean>() cacheGetter = mockk ReqShieldData?>() - globalLockFunc = mockk Boolean>() - globalUnLockFunc = mockk Boolean>() + globalLockFunc = mockk Boolean>() + globalUnLockFunc = mockk Boolean>() keyLock = mockk() keyGlobalLock = KeyGlobalLock(globalLockFunc, globalUnLockFunc, 3000) + backgroundScope = CoroutineScope(SupervisorJob() + Dispatchers.Default) coEvery { callable() } returns value - reqShield = - ReqShield( - ReqShieldConfiguration( - cacheSetter, - cacheGetter, - keyLock = keyLock, - ), - ) - - reqShieldOnlyUpdateCache = - ReqShield( - ReqShieldConfiguration( - cacheSetter, - cacheGetter, - keyLock = keyLock, - reqShieldWorkMode = ReqShieldWorkMode.ONLY_UPDATE_CACHE, - ), - ) - - reqShieldOnlyCreateCache = - ReqShield( - ReqShieldConfiguration( - cacheSetter, - cacheGetter, - keyLock = keyLock, - reqShieldWorkMode = ReqShieldWorkMode.ONLY_CREATE_CACHE, - ), - ) + reqShield = reqShieldOf() + reqShieldOnlyUpdateCache = reqShieldOf(workMode = ReqShieldWorkMode.ONLY_UPDATE_CACHE) + reqShieldOnlyCreateCache = reqShieldOf(workMode = ReqShieldWorkMode.ONLY_CREATE_CACHE) reqShieldForGlobalLock = ReqShield( @@ -120,74 +117,105 @@ class ReqShieldTest : BaseReqShieldTest { globalUnLockFunc, isLocalLock = false, keyLock = keyGlobalLock, + scope = backgroundScope, ), ) - - mockkStatic(LocalDateTime::class) - every { LocalDateTime.now() } returns LocalDateTime.of(2023, 11, 13, 12, 0, 0, 0) } @AfterEach fun tearDown() { - unmockkStatic(LocalDateTime::class) + backgroundScope.cancel() + } + + private fun reqShieldOf( + workMode: ReqShieldWorkMode = ReqShieldWorkMode.CREATE_AND_UPDATE_CACHE, + maxAttemptGetCache: Int = MAX_ATTEMPT_GET_CACHE, + ): ReqShield = + ReqShield( + ReqShieldConfiguration( + setCacheFunction = cacheSetter, + getCacheFunction = cacheGetter, + keyLock = keyLock, + maxAttemptGetCache = maxAttemptGetCache, + reqShieldWorkMode = workMode, + scope = backgroundScope, + ), + ) + + /** Awaits every cache write already submitted to [backgroundScope]. */ + private suspend fun awaitBackgroundWrites() { + backgroundScope.coroutineContext.job.children.toList().forEach { it.join() } } + private fun cachedData( + cachedValue: Product?, + ttl: Long, + createdAt: Long = nowToEpochTime(), + ): ReqShieldData = ReqShieldData(cachedValue, ReqShieldData.Status.NEW, createdAt, ttl) + + /** A cache entry that has consumed 90% of its TTL, i.e. past the default 80% update threshold. */ + private fun updateTargetData( + cachedValue: Product?, + ttl: Long, + ): ReqShieldData = cachedData(cachedValue, ttl, createdAt = nowToEpochTime() - (ttl * 0.9).toLong()) + @Test override fun testSetMethodCacheNotExistsAndLocalLockAcquired() = - runBlocking { + runTest { coEvery { cacheGetter.invoke(key) } returns null coEvery { cacheSetter.invoke(key, any(), any()) } returns true - coEvery { keyLock.tryLock(key, LockType.CREATE) } returns true - coEvery { keyLock.unLock(key, LockType.CREATE) } returns true + coEvery { keyLock.tryLock(key, LockType.CREATE) } returns LOCAL_TOKEN + coEvery { keyLock.unLock(key, LockType.CREATE, LOCAL_TOKEN) } returns true val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) - delay(100) + awaitBackgroundWrites() assertNotNull(result) + assertEquals(value, result.value) coVerify { cacheGetter.invoke(key) } coVerify { cacheSetter.invoke(key, result, timeToLiveMillis) } coVerify { keyLock.tryLock(key, LockType.CREATE) } - coVerify { keyLock.unLock(key, LockType.CREATE) } + coVerify { keyLock.unLock(key, LockType.CREATE, LOCAL_TOKEN) } coVerify { callable() } } @Test - override fun testSetMethodCacheNotExistsAndOnlyUpdateCache() { - runBlocking { + override fun testSetMethodCacheNotExistsAndOnlyUpdateCache() = + runTest { coEvery { cacheGetter.invoke(key) } returns null coEvery { cacheSetter.invoke(key, any(), any()) } returns true val result = reqShieldOnlyUpdateCache.getAndSetReqShieldData(key, callable, timeToLiveMillis) - delay(100) + awaitBackgroundWrites() assertNotNull(result) coVerify { cacheGetter.invoke(key) } coVerify { cacheSetter.invoke(key, result, timeToLiveMillis) } coVerify(inverse = true) { keyLock.tryLock(key, LockType.CREATE) } - coVerify(inverse = true) { keyLock.unLock(key, LockType.CREATE) } + coVerify(inverse = true) { keyLock.unLock(key, LockType.CREATE, any()) } coVerify { callable() } } - } @Test override fun testSetMethodCacheNotExistsAndGlobalLockAcquired() = - runBlocking { + runTest { coEvery { cacheGetter.invoke(key) } returns null coEvery { cacheSetter.invoke(key, any(), any()) } returns true - coEvery { globalLockFunc(any(), any()) } returns true - coEvery { globalUnLockFunc(any()) } returns true + coEvery { globalLockFunc(createLockKey, any(), any()) } returns true + coEvery { globalUnLockFunc(createLockKey, any()) } returns true val result = reqShieldForGlobalLock.getAndSetReqShieldData(key, callable, timeToLiveMillis) - delay(100) + awaitBackgroundWrites() assertNotNull(result) coVerify { cacheGetter.invoke(key) } coVerify { cacheSetter.invoke(key, result, timeToLiveMillis) } - coVerify { globalLockFunc(any(), any()) } - coVerify { globalUnLockFunc(any()) } - coVerify { keyGlobalLock.tryLock(key, LockType.CREATE) } - coVerify { keyGlobalLock.unLock(key, LockType.CREATE) } + + // The very token handed out by the lock function must be the one released. + val tokenSlot = slot() + coVerify { globalLockFunc(createLockKey, capture(tokenSlot), 3000) } + coVerify { globalUnLockFunc(createLockKey, tokenSlot.captured) } coVerify { callable() } } @@ -211,15 +239,15 @@ class ReqShieldTest : BaseReqShieldTest { @Test override fun testSetMethodCacheNotExistsAndLocalLockAcquiredAndCallableReturnNull() = - runBlocking { + runTest { coEvery { cacheGetter.invoke(key) } returns null coEvery { cacheSetter.invoke(key, any(), any()) } returns true - coEvery { keyLock.tryLock(key, LockType.CREATE) } returns true - coEvery { keyLock.unLock(key, LockType.CREATE) } returns true + coEvery { keyLock.tryLock(key, LockType.CREATE) } returns LOCAL_TOKEN + coEvery { keyLock.unLock(key, LockType.CREATE, LOCAL_TOKEN) } returns true coEvery { callable() } returns null val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) - delay(100) + awaitBackgroundWrites() assertNotNull(result) assertNull(result.value) @@ -227,81 +255,86 @@ class ReqShieldTest : BaseReqShieldTest { coVerify { cacheGetter.invoke(key) } coVerify { cacheSetter.invoke(key, result, timeToLiveMillis) } coVerify { keyLock.tryLock(key, LockType.CREATE) } - coVerify { keyLock.unLock(key, LockType.CREATE) } + coVerify { keyLock.unLock(key, LockType.CREATE, LOCAL_TOKEN) } coVerify { callable() } } @Test override fun testSetMethodCacheNotExistsAndGlobalLockAcquiredAndCallableReturnNull() = - runBlocking { + runTest { coEvery { cacheGetter.invoke(key) } returns null coEvery { cacheSetter.invoke(key, any(), any()) } returns true - coEvery { globalLockFunc(any(), any()) } returns true - coEvery { globalUnLockFunc(any()) } returns true + coEvery { globalLockFunc(createLockKey, any(), any()) } returns true + coEvery { globalUnLockFunc(createLockKey, any()) } returns true coEvery { callable() } returns null val result = reqShieldForGlobalLock.getAndSetReqShieldData(key, callable, timeToLiveMillis) - delay(100) + awaitBackgroundWrites() assertNotNull(result) assertNull(result.value) coVerify { cacheGetter.invoke(key) } coVerify { cacheSetter.invoke(key, result, timeToLiveMillis) } - coVerify { globalLockFunc(any(), any()) } - coVerify { globalUnLockFunc(any()) } - coVerify { keyGlobalLock.tryLock(key, LockType.CREATE) } - coVerify { keyGlobalLock.unLock(key, LockType.CREATE) } + + val tokenSlot = slot() + coVerify { globalLockFunc(createLockKey, capture(tokenSlot), 3000) } + coVerify { globalUnLockFunc(createLockKey, tokenSlot.captured) } coVerify { callable() } } @Test override fun testSetMethodCacheNotExistsAndLocalLockAcquiredAndThrowCallableClientException() = - runBlocking { + runTest { coEvery { cacheGetter.invoke(key) } returns null coEvery { cacheSetter.invoke(key, any(), any()) } returns true - coEvery { keyLock.tryLock(key, LockType.CREATE) } returns true - coEvery { keyLock.unLock(key, LockType.CREATE) } returns true + coEvery { keyLock.tryLock(key, LockType.CREATE) } returns LOCAL_TOKEN + coEvery { keyLock.unLock(key, LockType.CREATE, LOCAL_TOKEN) } returns true coEvery { callable() } throws Exception("callable error") val result = runCatching { reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) } - delay(100) + awaitBackgroundWrites() assertTrue(result.isFailure) - assertTrue(result.exceptionOrNull() is ClientException) - assertEquals(ErrorCode.SUPPLIER_ERROR, (result.exceptionOrNull() as? ClientException)?.errorCode) + val exception = result.exceptionOrNull() as? ClientException + assertNotNull(exception) + assertEquals(ErrorCode.SUPPLIER_ERROR, exception!!.errorCode) + assertEquals("callable error", exception.cause?.message) coVerify { cacheGetter.invoke(key) } coVerify { keyLock.tryLock(key, LockType.CREATE) } - coVerify { keyLock.unLock(key, LockType.CREATE) } + coVerify { keyLock.unLock(key, LockType.CREATE, LOCAL_TOKEN) } + coVerify(inverse = true) { cacheSetter.invoke(key, any(), any()) } coVerify { callable() } } @Test override fun testSetMethodCacheNotExistsAndGlobalLockAcquiredAndThrowCallableClientException() = - runBlocking { + runTest { coEvery { cacheGetter.invoke(key) } returns null coEvery { cacheSetter.invoke(key, any(), any()) } returns true - coEvery { globalLockFunc(any(), any()) } returns true - coEvery { globalUnLockFunc(any()) } returns true + coEvery { globalLockFunc(createLockKey, any(), any()) } returns true + coEvery { globalUnLockFunc(createLockKey, any()) } returns true coEvery { callable() } throws Exception("callable error") val result = runCatching { reqShieldForGlobalLock.getAndSetReqShieldData(key, callable, timeToLiveMillis) } - delay(100) + awaitBackgroundWrites() assertTrue(result.isFailure) - assertTrue(result.exceptionOrNull() is ClientException) - assertEquals(ErrorCode.SUPPLIER_ERROR, (result.exceptionOrNull() as? ClientException)?.errorCode) + val exception = result.exceptionOrNull() as? ClientException + assertNotNull(exception) + assertEquals(ErrorCode.SUPPLIER_ERROR, exception!!.errorCode) + assertEquals("callable error", exception.cause?.message) coVerify { cacheGetter.invoke(key) } - coVerify { globalLockFunc(any(), any()) } - coVerify { globalUnLockFunc(any()) } - coVerify { keyGlobalLock.tryLock(key, LockType.CREATE) } - coVerify { keyGlobalLock.unLock(key, LockType.CREATE) } + + val tokenSlot = slot() + coVerify { globalLockFunc(createLockKey, capture(tokenSlot), 3000) } + coVerify { globalUnLockFunc(createLockKey, tokenSlot.captured) } coVerify { callable() } } @@ -310,186 +343,181 @@ class ReqShieldTest : BaseReqShieldTest { runTest { coEvery { cacheGetter.invoke(key) } throws Exception("get cache error") coEvery { cacheSetter.invoke(key, any(), any()) } returns true - coEvery { keyLock.tryLock(key, LockType.CREATE) } returns true - coEvery { keyLock.unLock(key, LockType.CREATE) } returns true + coEvery { keyLock.tryLock(key, LockType.CREATE) } returns LOCAL_TOKEN + coEvery { keyLock.unLock(key, LockType.CREATE, LOCAL_TOKEN) } returns true val result = runCatching { reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) } - delay(100) assertTrue(result.isFailure) - assertTrue(result.exceptionOrNull() is ClientException) - assertEquals(ErrorCode.GET_CACHE_ERROR, (result.exceptionOrNull() as? ClientException)?.errorCode) + val exception = result.exceptionOrNull() as? ClientException + assertNotNull(exception) + assertEquals(ErrorCode.GET_CACHE_ERROR, exception!!.errorCode) + assertEquals("get cache error", exception.cause?.message) coVerify { cacheGetter.invoke(key) } coVerify(inverse = true) { keyLock.tryLock(key, LockType.CREATE) } - coVerify(inverse = true) { keyLock.unLock(key, LockType.CREATE) } + coVerify(inverse = true) { keyLock.unLock(key, LockType.CREATE, any()) } coVerify(inverse = true) { callable() } } @Test override fun testSetMethodCacheNotExistsAndGlobalLockAcquiredAndThrowGetCacheClientException() = - runBlocking { + runTest { coEvery { cacheGetter.invoke(key) } throws Exception("get cache error") coEvery { cacheSetter.invoke(key, any(), any()) } returns true - coEvery { globalLockFunc(any(), any()) } returns true - coEvery { globalUnLockFunc(any()) } returns true + coEvery { globalLockFunc(any(), any(), any()) } returns true + coEvery { globalUnLockFunc(any(), any()) } returns true val result = runCatching { reqShieldForGlobalLock.getAndSetReqShieldData(key, callable, timeToLiveMillis) } - delay(100) assertTrue(result.isFailure) - assertTrue(result.exceptionOrNull() is ClientException) - assertEquals(ErrorCode.GET_CACHE_ERROR, (result.exceptionOrNull() as? ClientException)?.errorCode) + val exception = result.exceptionOrNull() as? ClientException + assertNotNull(exception) + assertEquals(ErrorCode.GET_CACHE_ERROR, exception!!.errorCode) coVerify { cacheGetter.invoke(key) } - coVerify(inverse = true) { globalLockFunc(any(), any()) } - coVerify(inverse = true) { globalUnLockFunc(any()) } - coVerify(inverse = true) { keyLock.tryLock(key, LockType.CREATE) } - coVerify(inverse = true) { keyLock.unLock(key, LockType.CREATE) } + coVerify(inverse = true) { globalLockFunc(any(), any(), any()) } + coVerify(inverse = true) { globalUnLockFunc(any(), any()) } coVerify(inverse = true) { callable() } } @Test override fun testSetMethodCacheNotExistsAndLocalLockNotAcquired() = - runBlocking { - val timeToLiveMillis: Long = 10000 - + runTest { coEvery { cacheGetter.invoke(key) } returns null - coEvery { keyLock.tryLock(key, LockType.CREATE) } returns false + coEvery { keyLock.tryLock(key, LockType.CREATE) } returns null val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) - withTimeoutOrNull(1000L) { - while (result.value == null) { - delay(100L) // Check every 100 milliseconds - } - } - + // The cache was never filled by the lock owner, so the supplier is the last resort. assertNotNull(result) - coVerify { cacheGetter.invoke(key) } + assertEquals(value, result.value) + coVerify(exactly = 1 + MAX_ATTEMPT_GET_CACHE) { cacheGetter.invoke(key) } coVerify(inverse = true) { cacheSetter.invoke(key, any(), any()) } coVerify { keyLock.tryLock(key, LockType.CREATE) } + coVerify(inverse = true) { keyLock.unLock(key, LockType.CREATE, any()) } } @Test override fun testSetMethodCacheNotExistsAndGlobalLockNotAcquired() = - runBlocking { - val timeToLiveMillis: Long = 10000 - + runTest { coEvery { cacheGetter.invoke(key) } returns null - coEvery { globalLockFunc(any(), any()) } returns false + coEvery { globalLockFunc(createLockKey, any(), any()) } returns false val result = reqShieldForGlobalLock.getAndSetReqShieldData(key, callable, timeToLiveMillis) - withTimeoutOrNull(1000L) { - while (result.value == null) { - delay(100L) // Check every 100 milliseconds - } - } - assertNotNull(result) - coVerify { cacheGetter.invoke(key) } + assertEquals(value, result.value) coVerify(inverse = true) { cacheSetter.invoke(key, any(), any()) } - coVerify { globalLockFunc(any(), any()) } - coVerify { keyGlobalLock.tryLock(key, LockType.CREATE) } + coVerify { globalLockFunc(createLockKey, any(), any()) } + coVerify(inverse = true) { globalUnLockFunc(any(), any()) } } @Test override fun testSetMethodCacheExistsButNotTargetedForUpdate() = runTest { val timeToLiveMillis: Long = 10000 - val reqShieldData = ReqShieldData(value, timeToLiveMillis) + // Freshly created entry: far from the 80% update threshold. + val reqShieldData = cachedData(value, timeToLiveMillis) coEvery { cacheGetter.invoke(key) } returns reqShieldData - coEvery { cacheSetter.invoke(key, any(), any()) } returns true - coEvery { keyLock.tryLock(key, LockType.UPDATE) } returns false val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) assertEquals(reqShieldData, result) coVerify { cacheGetter.invoke(key) } + coVerify(inverse = true) { keyLock.tryLock(key, LockType.UPDATE) } + coVerify(inverse = true) { cacheSetter.invoke(key, any(), any()) } + coVerify(inverse = true) { callable() } } @Test override fun testSetMethodCacheExistsAndTheUpdateTarget() = - runBlocking { + runTest { val timeToLiveMillis: Long = 1000 - val reqShieldData = ReqShieldData(oldValue, timeToLiveMillis) - val newReqShieldData = ReqShieldData(value, timeToLiveMillis) + val reqShieldData = updateTargetData(oldValue, timeToLiveMillis) coEvery { cacheGetter.invoke(key) } returns reqShieldData - coEvery { cacheSetter.invoke(key, any(), any()) } coAnswers { true } - coEvery { keyLock.tryLock(key, LockType.UPDATE) } returns true - coEvery { keyLock.unLock(key, LockType.UPDATE) } returns true + coEvery { cacheSetter.invoke(key, any(), any()) } returns true + coEvery { keyLock.tryLock(key, LockType.UPDATE) } returns LOCAL_TOKEN + coEvery { keyLock.unLock(key, LockType.UPDATE, LOCAL_TOKEN) } returns true val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) + awaitBackgroundWrites() - delay(100) - + // The stale entry is served while the refresh happens in the background. assertEquals(reqShieldData, result) coVerify { cacheGetter.invoke(key) } - coVerify { cacheSetter.invoke(key, newReqShieldData, timeToLiveMillis) } + + val dataSlot = slot>() + coVerify { cacheSetter.invoke(key, capture(dataSlot), timeToLiveMillis) } + assertEquals(value, dataSlot.captured.value) + assertEquals(timeToLiveMillis, dataSlot.captured.timeToLiveMillis) + coVerify { keyLock.tryLock(key, LockType.UPDATE) } - coVerify { keyLock.unLock(key, LockType.UPDATE) } + coVerify { keyLock.unLock(key, LockType.UPDATE, LOCAL_TOKEN) } coVerify { callable() } } @Test - override fun testSetMethodCacheExistsAndTheUpdateTargetOnlyCreateCache() { - runBlocking { + override fun testSetMethodCacheExistsAndTheUpdateTargetOnlyCreateCache() = + runTest { val timeToLiveMillis: Long = 1000 - val reqShieldData = ReqShieldData(oldValue, timeToLiveMillis) - val newReqShieldData = ReqShieldData(value, timeToLiveMillis) + val reqShieldData = updateTargetData(oldValue, timeToLiveMillis) coEvery { cacheGetter.invoke(key) } returns reqShieldData - coEvery { cacheSetter.invoke(key, any(), any()) } coAnswers { true } + coEvery { cacheSetter.invoke(key, any(), any()) } returns true val result = reqShieldOnlyCreateCache.getAndSetReqShieldData(key, callable, timeToLiveMillis) - - delay(100) + awaitBackgroundWrites() assertEquals(reqShieldData, result) coVerify { cacheGetter.invoke(key) } - coVerify { cacheSetter.invoke(key, newReqShieldData, timeToLiveMillis) } + + val dataSlot = slot>() + coVerify { cacheSetter.invoke(key, capture(dataSlot), timeToLiveMillis) } + assertEquals(value, dataSlot.captured.value) + coVerify(inverse = true) { keyLock.tryLock(key, LockType.UPDATE) } - coVerify(inverse = true) { keyLock.unLock(key, LockType.UPDATE) } + coVerify(inverse = true) { keyLock.unLock(key, LockType.UPDATE, any()) } coVerify { callable() } } - } @Test override fun testSetMethodCacheExistsAndTheUpdateTargetAndCallableReturnNull() = - runBlocking { + runTest { timeToLiveMillis = 1000 - val reqShieldData = ReqShieldData(value, timeToLiveMillis) - val reqShieldDataNull = ReqShieldData(null, timeToLiveMillis) + val reqShieldData = updateTargetData(value, timeToLiveMillis) coEvery { cacheGetter.invoke(key) } returns reqShieldData - coEvery { cacheSetter.invoke(key, any(), any()) } coAnswers { true } - coEvery { keyLock.tryLock(key, LockType.UPDATE) } returns true - coEvery { keyLock.unLock(key, LockType.UPDATE) } returns true + coEvery { cacheSetter.invoke(key, any(), any()) } returns true + coEvery { keyLock.tryLock(key, LockType.UPDATE) } returns LOCAL_TOKEN + coEvery { keyLock.unLock(key, LockType.UPDATE, LOCAL_TOKEN) } returns true coEvery { callable() } returns null val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) - delay(100) + awaitBackgroundWrites() assertEquals(reqShieldData, result) coVerify { cacheGetter.invoke(key) } - coVerify { cacheSetter.invoke(key, reqShieldDataNull, timeToLiveMillis) } + + val dataSlot = slot>() + coVerify { cacheSetter.invoke(key, capture(dataSlot), timeToLiveMillis) } + assertNull(dataSlot.captured.value) + coVerify { keyLock.tryLock(key, LockType.UPDATE) } - coVerify { keyLock.unLock(key, LockType.UPDATE) } + coVerify { keyLock.unLock(key, LockType.UPDATE, LOCAL_TOKEN) } coVerify { callable() } } @Test override fun executeSetCacheFunctionShouldHandleExceptionFromCacheSetter() = runBlocking { - coEvery { keyLock.tryLock(any(), any()) } returns true - coEvery { keyLock.unLock(any(), any()) } returns true + coEvery { keyLock.unLock(any(), any(), any()) } returns true val key = "key" - val reqShieldData = ReqShieldData(value, 1000L) + val reqShieldData = cachedData(value, 1000L) val lockType = LockType.CREATE val method: Method = @@ -508,14 +536,375 @@ class ReqShieldTest : BaseReqShieldTest { coEvery { cacheSetter.invoke(any(), any(), any()) } throws Exception("set cache error") val exception = assertFailsWith { - method.invoke(reqShield, cacheSetter, key, reqShieldData, lockType, continuation) + method.invoke(reqShield, cacheSetter, key, reqShieldData, lockType, LOCAL_TOKEN, continuation) } val cause = exception.cause assertTrue(cause is ClientException) assertEquals(ErrorCode.SET_CACHE_ERROR, cause.errorCode) + assertEquals("set cache error", cause.cause?.message) coVerify { cacheSetter.invoke(key, reqShieldData, 1000L) } - coVerify { keyLock.unLock(any(), any()) } + coVerify { keyLock.unLock(key, lockType, LOCAL_TOKEN) } } + + @Test + fun `should not unlock when no lock was acquired and the cache setter fails`() = + runBlocking { + val key = "key" + val reqShieldData = cachedData(value, 1000L) + + val method: Method = + ReqShield::class.java.declaredMethods.first { it.name == "executeSetCacheFunction" } + method.isAccessible = true + + val continuation = + object : Continuation { + override val context = EmptyCoroutineContext + + override fun resumeWith(result: Result) { + result.getOrThrow() + } + } + coEvery { cacheSetter.invoke(any(), any(), any()) } throws Exception("set cache error") + + assertFailsWith { + method.invoke(reqShield, cacheSetter, key, reqShieldData, LockType.CREATE, null, continuation) + } + + // A null token means this call never held the lock, so it must not release anyone else's. + coVerify(inverse = true) { keyLock.unLock(any(), any(), any()) } + } + + @Test + fun `should return the entry another request cached while waiting for the lock`() = + runTest { + val cached = cachedData(value, timeToLiveMillis) + val reads = AtomicInteger(0) + + coEvery { cacheGetter.invoke(key) } coAnswers { + if (reads.incrementAndGet() >= 3) cached else null + } + coEvery { keyLock.tryLock(key, LockType.CREATE) } returns null + + val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) + + assertEquals(cached, result) + // 1 read for the initial miss + 2 polls until the lock owner filled the cache + assertEquals(3, reads.get()) + coVerify(inverse = true) { callable() } + coVerify(inverse = true) { cacheSetter.invoke(any(), any(), any()) } + } + + @Test + fun `should fall back to the supplier after consecutive cache read failures`() = + runTest { + val reads = AtomicInteger(0) + + coEvery { cacheGetter.invoke(key) } coAnswers { + if (reads.incrementAndGet() == 1) null else throw Exception("cache is down") + } + coEvery { keyLock.tryLock(key, LockType.CREATE) } returns null + + val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) + + assertEquals(value, result.value) + // The wait gives up as soon as the failures become consecutive, long before maxAttemptGetCache + assertEquals(1 + MAX_CONSECUTIVE_GET_CACHE_FAILURES, reads.get()) + coVerify(exactly = 1) { callable() } + coVerify(inverse = true) { cacheSetter.invoke(any(), any(), any()) } + } + + @Test + fun `should keep waiting when cache read failures are not consecutive`() = + runTest { + val maxAttemptGetCache = 6 + val reqShieldWithShortWait = reqShieldOf(maxAttemptGetCache = maxAttemptGetCache) + val reads = AtomicInteger(0) + + // Alternate failure / miss so the failure streak never reaches the threshold. + coEvery { cacheGetter.invoke(key) } coAnswers { + val read = reads.incrementAndGet() + if (read > 1 && read % 2 == 0) throw Exception("cache is flaky") else null + } + coEvery { keyLock.tryLock(key, LockType.CREATE) } returns null + + val result = reqShieldWithShortWait.getAndSetReqShieldData(key, callable, timeToLiveMillis) + + // Every attempt was used: an intermittent failure must not end the wait early + assertEquals(1 + maxAttemptGetCache, reads.get()) + assertEquals(value, result.value) + coVerify(exactly = 1) { callable() } + } + + @Test + fun `should raise a supplier error when the fallback supplier fails`() = + runTest { + val reqShieldWithShortWait = reqShieldOf(maxAttemptGetCache = 2) + + coEvery { cacheGetter.invoke(key) } returns null + coEvery { keyLock.tryLock(key, LockType.CREATE) } returns null + coEvery { callable() } throws IllegalStateException("supplier is down") + + val exception = + assertFailsWith { + reqShieldWithShortWait.getAndSetReqShieldData(key, callable, timeToLiveMillis) + } + + assertEquals(ErrorCode.SUPPLIER_ERROR, exception.errorCode) + assertNotNull(exception.cause) + assertEquals("supplier is down", exception.cause?.message) + // No lock was taken while waiting, so there is nothing to release. + coVerify(inverse = true) { keyLock.unLock(any(), any(), any()) } + } + + @Test + fun `should not fail the caller when the background cache write fails`() = + runTest { + coEvery { cacheGetter.invoke(key) } returns null + coEvery { cacheSetter.invoke(key, any(), any()) } throws Exception("set cache error") + coEvery { keyLock.tryLock(key, LockType.CREATE) } returns LOCAL_TOKEN + coEvery { keyLock.unLock(key, LockType.CREATE, LOCAL_TOKEN) } returns true + + val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) + awaitBackgroundWrites() + + // The write is fire-and-forget: its failure is logged, never handed to the caller, + // and the lock is released anyway. + assertEquals(value, result.value) + coVerify { cacheSetter.invoke(key, result, timeToLiveMillis) } + coVerify { keyLock.unLock(key, LockType.CREATE, LOCAL_TOKEN) } + } + + @Test + fun `should release the lock when the background cache update fails`() = + runTest { + val timeToLiveMillis: Long = 1000 + val reqShieldData = updateTargetData(oldValue, timeToLiveMillis) + + coEvery { cacheGetter.invoke(key) } returns reqShieldData + coEvery { keyLock.tryLock(key, LockType.UPDATE) } returns LOCAL_TOKEN + coEvery { keyLock.unLock(key, LockType.UPDATE, LOCAL_TOKEN) } returns true + coEvery { callable() } throws Exception("callable error") + + val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) + awaitBackgroundWrites() + + // The stale entry keeps being served, and the failing supplier still releases the lock. + assertEquals(reqShieldData, result) + coVerify { keyLock.unLock(key, LockType.UPDATE, LOCAL_TOKEN) } + coVerify(inverse = true) { cacheSetter.invoke(any(), any(), any()) } + } + + @Test + fun `should stop polling the cache when the caller is cancelled`() = + runBlocking { + val reads = AtomicInteger(0) + + coEvery { cacheGetter.invoke(key) } coAnswers { + reads.incrementAndGet() + null + } + coEvery { keyLock.tryLock(key, LockType.CREATE) } returns null + + val waiter = + launch(Dispatchers.Default) { + reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) + } + + withTimeout(2000L) { + while (reads.get() < 3) { + delay(10L) + } + } + waiter.cancelAndJoin() + + val readsAtCancellation = reads.get() + delay(GET_CACHE_INTERVAL_MILLIS * 5) + + assertEquals(readsAtCancellation, reads.get(), "Polling must stop once the caller is cancelled") + coVerify(inverse = true) { callable() } + } + + @Test + fun shouldReleaseLocalCreateLockWhenSupplierTimesOut() = + runTest { + val localLock = KeyLocalLock(60_000) + val uniqueKey = "supplier-timeout-${System.nanoTime()}" + val shield = + ReqShield( + ReqShieldConfiguration( + setCacheFunction = cacheSetter, + getCacheFunction = { null }, + keyLock = localLock, + scope = backgroundScope, + ), + ) + + assertFailsWith { + withTimeout(100) { + shield.getAndSetReqShieldData(uniqueKey, { awaitCancellation() }, timeToLiveMillis) + } + } + + val token = localLock.tryLock(uniqueKey, LockType.CREATE) + assertNotNull(token, "Cancellation must release the lock before its TTL") + assertTrue(localLock.unLock(uniqueKey, LockType.CREATE, token!!)) + coVerify(exactly = 0) { cacheSetter.invoke(any(), any(), any()) } + } + + @Test + fun shouldReleaseGlobalCreateLockWhenSupplierIsCancelled() = + verifyGlobalLockReleasedOnCancellation(LockType.CREATE, cancelSetter = false) + + @Test + fun shouldReleaseGlobalUpdateLockWhenSupplierIsCancelled() = + verifyGlobalLockReleasedOnCancellation(LockType.UPDATE, cancelSetter = false) + + @Test + fun shouldReleaseGlobalLockWhenCacheSetterIsCancelled() = verifyGlobalLockReleasedOnCancellation(LockType.CREATE, cancelSetter = true) + + @Test + fun shouldReleaseGlobalCreateLockWhenScopeIsCancelledBeforeWriteStarts() = + verifyGlobalLockReleasedBeforeBackgroundTaskStarts(LockType.CREATE) + + @Test + fun shouldReleaseGlobalUpdateLockWhenScopeIsCancelledBeforeSupplierStarts() = + verifyGlobalLockReleasedBeforeBackgroundTaskStarts(LockType.UPDATE) + + private fun verifyGlobalLockReleasedBeforeBackgroundTaskStarts(lockType: LockType) = + runTest { + val locks = mutableMapOf() + var releases = 0 + val globalLock = + KeyGlobalLock( + globalLockFunction = { lockKey, token, _ -> locks.putIfAbsent(lockKey, token) == null }, + globalUnLockFunction = { lockKey, token -> + // Suspending cleanup must finish even after cancellation. + delay(1) + releases++ + locks.remove(lockKey, token) + }, + lockTimeoutMillis = 60_000, + ) + val taskScope = CoroutineScope(coroutineContext + SupervisorJob()) + var supplierCalls = 0 + var writes = 0 + val shield = + ReqShield( + ReqShieldConfiguration( + setCacheFunction = { _, _, _ -> + writes++ + true + }, + getCacheFunction = { + if (lockType == LockType.UPDATE) updateTargetData(oldValue, timeToLiveMillis) else null + }, + keyLock = globalLock, + scope = taskScope, + ), + ) + + try { + val result = + shield.getAndSetReqShieldData( + key, + { + supplierCalls++ + value + }, + timeToLiveMillis, + ) + val expectedSupplierCalls = if (lockType == LockType.CREATE) 1 else 0 + assertEquals(if (lockType == LockType.CREATE) value else oldValue, result.value) + assertEquals(expectedSupplierCalls, supplierCalls) + assertEquals(0, writes, "The background task must still be queued") + assertNull(globalLock.tryLock(key, lockType), "The queued task already owns the lock") + + val owner = taskScope.coroutineContext.job.children.single() + val cancellation = CancellationException("scope stopped before dispatch") + var propagated: Throwable? = null + owner.invokeOnCompletion { propagated = it } + // Cancel before yielding to the test dispatcher, so the task has not started. + taskScope.cancel(cancellation) + owner.join() + + assertSame(cancellation, propagated) + assertEquals(expectedSupplierCalls, supplierCalls, "Cancellation must not start the supplier") + assertEquals(0, writes, "Cancellation must not start the cache write") + assertEquals(1, releases, "The owned token must be released exactly once") + val token = globalLock.tryLock(key, lockType) + assertNotNull(token, "Pre-start cancellation must release the lock before its TTL") + assertTrue(globalLock.unLock(key, lockType, token!!)) + } finally { + taskScope.cancel() + } + } + + private fun verifyGlobalLockReleasedOnCancellation( + lockType: LockType, + cancelSetter: Boolean, + ) = runTest { + val locks = mutableMapOf() + val globalLock = + KeyGlobalLock( + globalLockFunction = { lockKey, token, _ -> locks.putIfAbsent(lockKey, token) == null }, + globalUnLockFunction = { lockKey, token -> + // A suspending unlock must finish even after its caller was cancelled. + delay(1) + locks.remove(lockKey, token) + }, + lockTimeoutMillis = 60_000, + ) + val started = CompletableDeferred() + val cancellation = CancellationException("request cancelled") + var propagated: Throwable? = null + val taskScope = CoroutineScope(coroutineContext + SupervisorJob()) + val shield = + ReqShield( + ReqShieldConfiguration( + setCacheFunction = { _, _, _ -> + check(cancelSetter) { "A cancelled supplier must not write the cache" } + started.complete(Unit) + awaitCancellation() + }, + getCacheFunction = { + if (lockType == LockType.UPDATE) updateTargetData(oldValue, timeToLiveMillis) else null + }, + keyLock = globalLock, + scope = taskScope, + ), + ) + + try { + val caller = + taskScope.launch { + shield.getAndSetReqShieldData( + key, + { + if (!cancelSetter) { + started.complete(Unit) + awaitCancellation() + } + value + }, + timeToLiveMillis, + ) + } + started.await() + val owner = if (lockType == LockType.CREATE && !cancelSetter) caller else taskScope.coroutineContext.job.children.single() + owner.invokeOnCompletion { propagated = it } + assertNull(globalLock.tryLock(key, lockType), "The supplier or setter must hold the lock") + + owner.cancel(cancellation) + owner.join() + + assertSame(cancellation, propagated) + val token = globalLock.tryLock(key, lockType) + assertNotNull(token, "Suspending cleanup must release the owned lock") + assertTrue(globalLock.unLock(key, lockType, token!!)) + } finally { + taskScope.cancel() + } + } } diff --git a/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/KeyGlobalLock.kt b/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/KeyGlobalLock.kt index 3803b13..d74d335 100644 --- a/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/KeyGlobalLock.kt +++ b/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/KeyGlobalLock.kt @@ -16,26 +16,34 @@ package com.linecorp.cse.reqshield.reactor +import com.linecorp.cse.reqshield.support.constant.ConfigValues.LOCK_KEY_PREFIX import reactor.core.publisher.Mono +import java.util.UUID class KeyGlobalLock( - private val globalLockFunction: (String, Long) -> Mono, - private val globalUnLockFunction: (String) -> Mono, + private val globalLockFunction: (String, String, Long) -> Mono, + private val globalUnLockFunction: (String, String) -> Mono, private val lockTimeoutMillis: Long, ) : KeyLock { override fun tryLock( key: String, lockType: LockType, - ): Mono { - val completeKey = "${key}_${lockType.name}" - return globalLockFunction(completeKey, lockTimeoutMillis) - } + ): Mono = + Mono.defer { + // A fresh token per attempt: only this attempt may release the lock it acquired. + val token = UUID.randomUUID().toString() + globalLockFunction(completeKey(key, lockType), token, lockTimeoutMillis) + .flatMap { acquired -> if (acquired) Mono.just(token) else Mono.empty() } + } override fun unLock( key: String, lockType: LockType, - ): Mono { - val completeKey = "${key}_${lockType.name}" - return globalUnLockFunction(completeKey) - } + token: String, + ): Mono = globalUnLockFunction(completeKey(key, lockType), token) + + private fun completeKey( + key: String, + lockType: LockType, + ): String = "$LOCK_KEY_PREFIX${key}_${lockType.name}" } diff --git a/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/KeyLocalLock.kt b/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/KeyLocalLock.kt index 3ea4b06..b5aa768 100644 --- a/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/KeyLocalLock.kt +++ b/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.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 import org.slf4j.LoggerFactory @@ -28,6 +29,8 @@ import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.Semaphore import java.util.concurrent.atomic.AtomicBoolean import java.util.concurrent.atomic.AtomicInteger +import java.util.concurrent.atomic.AtomicLong +import java.util.concurrent.atomic.AtomicReference private val log = LoggerFactory.getLogger(KeyLocalLock::class.java) @@ -52,20 +55,32 @@ class KeyLocalLock( * check vs unLock, or monitor cleanup vs unLock). */ val isHeld: AtomicBoolean = AtomicBoolean(false), + /** + * Ownership token of the current holder, null when the lock is not held. + * 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. + */ + @Volatile var token: String? = null, ) companion object { private val lockMap = ConcurrentHashMap() + // Monotonic counter backing the local ownership tokens. A counter is enough because + // the tokens never leave this JVM, and it is far cheaper than UUID generation. + private val tokenSequence = AtomicLong(0) + @Volatile private var monitoringStarted: Boolean = false @Volatile private var monitorDisposable: Disposable? = null - // Track consecutive failures for backoff logging + // Track consecutive failures for rate-limited logging private val consecutiveFailures = AtomicInteger(0) + private fun nextToken(): String = "local-${tokenSequence.incrementAndGet()}" + private fun startMonitoringOnce() { if (monitoringStarted) return synchronized(this) { @@ -73,7 +88,8 @@ class KeyLocalLock( monitorDisposable = Flux .interval(Duration.ofMillis(LOCK_MONITOR_INTERVAL_MILLIS), Schedulers.single()) - .flatMap { + // Concurrency of 1: a slow cleanup must not fan out into overlapping runs. + .flatMap({ Mono .fromRunnable { val now = System.currentTimeMillis() @@ -93,6 +109,7 @@ class KeyLocalLock( if (lockInfo.isHeld.compareAndSet(true, false)) { lockInfo.semaphore.release() } + lockInfo.token = null null // Atomic removal } else { lockInfo // Keep the entry @@ -110,11 +127,9 @@ class KeyLocalLock( e.message, ) } - // Backoff: delay on failure (max 5 seconds) - val backoffMs = minOf(failures * LOCK_MONITOR_INTERVAL_MILLIS, 5000L) - Mono.delay(Duration.ofMillis(backoffMs)).then(Mono.empty()) + Mono.empty() } - }.subscribe( + }, 1).subscribe( { /* success - no action needed */ }, { e -> log.error("Fatal error in lock lifecycle monitoring: {}", e.message, e) }, ) @@ -140,11 +155,12 @@ class KeyLocalLock( override fun tryLock( key: String, lockType: LockType, - ): Mono = + ): Mono = Mono.fromCallable { - val completeKey = "${key}_${lockType.name}" + val completeKey = completeKey(key, lockType) val now = nowToEpochTime() - val result = AtomicBoolean(false) + // Holds the token handed out by this attempt, or null when the lock could not be acquired. + val acquiredToken = AtomicReference(null) // Use compute() for atomic lock acquisition. // This ensures mutual exclusion with cleanup - they cannot race on the same key. @@ -156,42 +172,62 @@ class KeyLocalLock( // both threads would call release(), causing over-release (permits > 1). if (now > existing.expiresAt && existing.isHeld.compareAndSet(true, false)) { existing.semaphore.release() + // The previous holder lost ownership: its token must no longer release the lock. + existing.token = null } // Existing entry: try to acquire semaphore if (existing.semaphore.tryAcquire()) { + val token = nextToken() existing.isHeld.set(true) existing.expiresAt = now + lockTimeoutMillis - result.set(true) + existing.token = token + acquiredToken.set(token) } existing } else { // New entry: create and acquire + val token = nextToken() val newLock = LockInfo(Semaphore(1), now + lockTimeoutMillis) newLock.semaphore.tryAcquire() // Always succeeds for new semaphore newLock.isHeld.set(true) - result.set(true) + newLock.token = token + acquiredToken.set(token) newLock } } - result.get() + // Mono.fromCallable completes empty on a null result, which signals "not acquired". + acquiredToken.get() } override fun unLock( key: String, lockType: LockType, + token: String, ): Mono = Mono.fromCallable { - val completeKey = "${key}_${lockType.name}" - val lockInfo = lockMap[completeKey] ?: return@fromCallable false - - // Use CAS to prevent over-release: only release if we actually hold the lock - if (lockInfo.isHeld.compareAndSet(true, false)) { - lockInfo.semaphore.release() - true - } else { - log.debug("Attempted to unlock key '{}' that is not held", completeKey) - false + 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. + 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) + } + existing // Keep the entry } + released.get() } + + private fun completeKey( + key: String, + lockType: LockType, + ): String = "$LOCK_KEY_PREFIX${key}_${lockType.name}" } diff --git a/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/KeyLock.kt b/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/KeyLock.kt index 19c0cf5..4c4ab88 100644 --- a/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/KeyLock.kt +++ b/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/KeyLock.kt @@ -19,14 +19,27 @@ package com.linecorp.cse.reqshield.reactor import reactor.core.publisher.Mono interface KeyLock { + /** + * Tries to acquire the lock for [key] and [lockType]. + * + * Emits an opaque ownership token when the lock was acquired, and completes EMPTY + * when another holder currently owns the lock. + */ fun tryLock( key: String, lockType: LockType, - ): Mono + ): Mono + /** + * Releases the lock for [key] and [lockType] only if [token] matches the current owner. + * + * Emits false when the lock is not held or the token does not match, so an expired lock + * that was already handed to another holder can never be released by a stale owner. + */ fun unLock( key: String, lockType: LockType, + token: String, ): Mono } diff --git a/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/ReqShield.kt b/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/ReqShield.kt index 7353314..35962ee 100644 --- a/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/ReqShield.kt +++ b/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/ReqShield.kt @@ -19,18 +19,26 @@ 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.constant.ConfigValues.GET_CACHE_INTERVAL_MILLIS +import com.linecorp.cse.reqshield.support.constant.ConfigValues.MAX_CONSECUTIVE_GET_CACHE_FAILURES import com.linecorp.cse.reqshield.support.exception.ClientException import com.linecorp.cse.reqshield.support.exception.code.ErrorCode import com.linecorp.cse.reqshield.support.model.ReqShieldData import com.linecorp.cse.reqshield.support.utils.decideToUpdateCache import org.slf4j.LoggerFactory -import reactor.core.publisher.Flux import reactor.core.publisher.Mono import java.time.Duration import java.util.concurrent.Callable +import java.util.concurrent.atomic.AtomicInteger private val log = LoggerFactory.getLogger(ReqShield::class.java) +/** + * Internal signal telling the lock waiter to stop polling the cache because the cache + * itself looks unavailable. Never leaves [ReqShield]: it is always resumed into the + * supplier fallback. Stack trace is disabled because it carries no diagnostic value. + */ +private class GetCacheUnavailableException : RuntimeException(null, null, false, false) + class ReqShield( private val reqShieldConfig: ReqShieldConfiguration, ) { @@ -46,7 +54,7 @@ class ReqShield( if (shouldUpdateCache(reqShieldData)) { updateReqShieldData(key, callable, timeToLiveMillis) } - Mono.justOrEmpty(reqShieldData!!) + Mono.justOrEmpty(reqShieldData) }.switchIfEmpty( Mono.defer { handleLockForCacheCreation(key, callable, timeToLiveMillis) @@ -69,8 +77,8 @@ class ReqShield( ) { val lockType = LockType.UPDATE - fun processMono(): Mono> = - executeCallable({ callable.call() }, true, key, lockType) + fun processMono(token: String?): Mono> = + executeCallable({ callable.call() }, key, lockType, token) .map { data -> buildReqShieldData(data, timeToLiveMillis) } .flatMap { reqShieldData -> setReqShieldData( @@ -78,6 +86,7 @@ class ReqShield( key, reqShieldData, lockType, + token, ).thenReturn(reqShieldData) }.switchIfEmpty( Mono.defer { @@ -87,28 +96,28 @@ class ReqShield( key, reqShieldData, lockType, + token, ).thenReturn(reqShieldData) }, ) - if (reqShieldConfig.reqShieldWorkMode == ReqShieldWorkMode.ONLY_CREATE_CACHE) { - processMono() - .subscribeOn(reqShieldConfig.scheduler) - .subscribe( - { /* success - no action needed */ }, - { e -> log.error("Failed to update cache for key '{}': {}", key, e.message, e) }, - ) - } else { - reqShieldConfig.keyLock - .tryLock(key, lockType) - .filter { it } - .flatMap { processMono() } - .subscribeOn(reqShieldConfig.scheduler) - .subscribe( - { /* success - no action needed */ }, - { e -> log.error("Failed to update cache for key '{}': {}", key, e.message, e) }, - ) - } + val updateMono = + if (reqShieldConfig.reqShieldWorkMode == ReqShieldWorkMode.ONLY_CREATE_CACHE) { + // This mode never refreshes an existing entry through the lock, so no token is taken. + processMono(null) + } else { + // Empty means another request already holds the update lock: nothing to do here. + reqShieldConfig.keyLock + .tryLock(key, lockType) + .flatMap { token -> processMono(token) } + } + + updateMono + .subscribeOn(reqShieldConfig.scheduler) + .subscribe( + { /* success - no action needed */ }, + { e -> log.error("Failed to update cache for key '{}': {}", key, e.message, e) }, + ) } private fun handleLockForCacheCreation( @@ -119,18 +128,18 @@ class ReqShield( val lockType = LockType.CREATE if (reqShieldConfig.reqShieldWorkMode == ReqShieldWorkMode.ONLY_UPDATE_CACHE) { - return createReqShieldData(key, callable, timeToLiveMillis, lockType) + // This mode never creates a cache entry through the lock, so no token is taken. + return createReqShieldData(key, callable, timeToLiveMillis, lockType, null) } return reqShieldConfig.keyLock .tryLock(key, lockType) - .flatMap { acquired -> - if (acquired) { - createReqShieldData(key, callable, timeToLiveMillis, lockType) - } else { + .flatMap { token -> createReqShieldData(key, callable, timeToLiveMillis, lockType, token) } + .switchIfEmpty( + Mono.defer { handleLockFailure(key, callable, timeToLiveMillis) - } - } + }, + ) } private fun createReqShieldData( @@ -138,8 +147,9 @@ class ReqShield( callable: Callable>, timeToLiveMillis: Long, lockType: LockType, + token: String?, ): Mono> = - executeCallable({ callable.call() }, true, key, lockType) + executeCallable({ callable.call() }, key, lockType, token) .map { data -> buildReqShieldData(data, timeToLiveMillis) } .doOnNext { reqShieldData -> // Async fire-and-forget cache storage (matches coroutine implementation) @@ -148,6 +158,7 @@ class ReqShield( key, reqShieldData, lockType, + token, ).subscribeOn(reqShieldConfig.scheduler) .subscribe( { /* success - no action needed */ }, @@ -162,6 +173,7 @@ class ReqShield( key, reqShieldData, lockType, + token, ).subscribeOn(reqShieldConfig.scheduler) .subscribe( { /* success - no action needed */ }, @@ -171,36 +183,66 @@ class ReqShield( }, ) + /** + * Waits for the request that owns the lock to fill the cache. + * + * The cache is polled up to `maxAttemptGetCache` times. A read failure is logged and counted; + * [MAX_CONSECUTIVE_GET_CACHE_FAILURES] consecutive failures are treated as a cache outage and + * stop the polling immediately. Once polling gives up, the supplier is called directly, and a + * failing supplier surfaces as `ClientException(SUPPLIER_ERROR)` instead of a null-valued entry. + */ private fun handleLockFailure( key: String, callable: Callable>, timeToLiveMillis: Long, - ): Mono> = - reqShieldConfig - .getCacheFunction(key) - .repeatWhenEmpty { - Flux - .range(1, reqShieldConfig.maxAttemptGetCache) + ): Mono> { + val consecutiveGetCacheFailures = AtomicInteger(0) + + return Mono + .defer { getCacheWhileWaitingForLock(key, consecutiveGetCacheFailures) } + .repeatWhenEmpty { companion -> + companion + .take(reqShieldConfig.maxAttemptGetCache.toLong()) .delayElements(Duration.ofMillis(GET_CACHE_INTERVAL_MILLIS)) - }.flatMap { reqShieldData -> - if (reqShieldData != null) { - Mono.just(reqShieldData) + }.onErrorResume(GetCacheUnavailableException::class.java) { Mono.empty() } + .switchIfEmpty( + Mono.defer { + executeCallable({ callable.call() }, key, null, null) + .map { data -> buildReqShieldData(data, timeToLiveMillis) } + .switchIfEmpty(Mono.fromSupplier { buildReqShieldData(null, timeToLiveMillis) }) + }, + ).subscribeOn(reqShieldConfig.scheduler) + } + + /** + * Reads the cache once while waiting for the lock holder. + * + * Completes empty while the cache is not filled yet, which drives the retry loop. + * A transient read failure also completes empty so the loop continues, but + * [MAX_CONSECUTIVE_GET_CACHE_FAILURES] consecutive failures raise [GetCacheUnavailableException] + * to bail out of the loop. Any successful read resets the failure counter. + */ + private fun getCacheWhileWaitingForLock( + key: String, + consecutiveGetCacheFailures: AtomicInteger, + ): Mono> = + executeGetCacheFunction(reqShieldConfig.getCacheFunction, key) + .doOnSuccess { consecutiveGetCacheFailures.set(0) } + .flatMap { reqShieldData -> Mono.justOrEmpty(reqShieldData) } + .onErrorResume { e -> + val failures = consecutiveGetCacheFailures.incrementAndGet() + log.warn( + "Failed to read cache for key '{}' while waiting for the lock holder (consecutive failures: {}): {}", + key, + failures, + e.message, + ) + if (failures >= MAX_CONSECUTIVE_GET_CACHE_FAILURES) { + Mono.error(GetCacheUnavailableException()) } else { Mono.empty() } - }.switchIfEmpty( - executeCallable({ callable.call() }, false) - .map { - buildReqShieldData(it, timeToLiveMillis) - }.onErrorResume { - Mono.just(buildReqShieldData(null, timeToLiveMillis)) - }.switchIfEmpty( - Mono.defer { - val reqShieldData = buildReqShieldData(null, timeToLiveMillis) - Mono.just(reqShieldData) - }, - ), - ).subscribeOn(reqShieldConfig.scheduler) + } private fun buildReqShieldData( value: T?, @@ -216,28 +258,31 @@ class ReqShield( key: String, reqShieldData: ReqShieldData, lockType: LockType, - ): Mono = executeSetCacheFunction(cacheSetter, key, reqShieldData, lockType) + token: String?, + ): Mono = executeSetCacheFunction(cacheSetter, key, reqShieldData, lockType, token) private fun executeGetCacheFunction( getFunction: (String) -> Mono?>, key: String, ): Mono?> = getFunction(key) - .onErrorMap { e -> ClientException(ErrorCode.GET_CACHE_ERROR, originErrorMessage = e.message) } + .onErrorMap { e -> ClientException(ErrorCode.GET_CACHE_ERROR, cause = e) } private fun executeSetCacheFunction( setFunction: (String, ReqShieldData, Long) -> Mono, key: String, value: ReqShieldData, lockType: LockType, + token: String?, ): Mono = setFunction(key, value, value.timeToLiveMillis) - .onErrorMap { e -> ClientException(ErrorCode.SET_CACHE_ERROR, originErrorMessage = e.message) } + .onErrorMap { e -> ClientException(ErrorCode.SET_CACHE_ERROR, cause = e) } .doFinally { - if (shouldAttemptUnlock(lockType)) { + // Only the holder of a token took a lock, so only it may release one. + if (token != null) { // No retry needed: false means lock already released or expired (not an error) reqShieldConfig.keyLock - .unLock(key, lockType) + .unLock(key, lockType, token) .doOnNext { unlocked -> if (!unlocked) { log.debug("Lock already released or expired for key '{}'", key) @@ -251,16 +296,17 @@ class ReqShield( private fun executeCallable( callable: Callable>, - isUnlockWhenException: Boolean, - key: String? = null, - lockType: LockType? = null, + key: String, + lockType: LockType?, + token: String?, ): Mono = callable .call() .doOnError { _ -> - if (isUnlockWhenException && key != null && lockType != null) { + // Only the holder of a token took a lock, so only it may release one. + if (lockType != null && token != null) { reqShieldConfig.keyLock - .unLock(key, lockType) + .unLock(key, lockType, token) .subscribe( { /* success - no action needed */ }, { unlockError -> @@ -274,10 +320,6 @@ class ReqShield( ) } }.onErrorMap { e -> - ClientException(ErrorCode.SUPPLIER_ERROR, originErrorMessage = e.message) + ClientException(ErrorCode.SUPPLIER_ERROR, cause = e) } - - private fun shouldAttemptUnlock(lockType: LockType): Boolean = - (lockType == LockType.UPDATE && reqShieldConfig.reqShieldWorkMode != ReqShieldWorkMode.ONLY_CREATE_CACHE) || - (lockType == LockType.CREATE && reqShieldConfig.reqShieldWorkMode != ReqShieldWorkMode.ONLY_UPDATE_CACHE) } diff --git a/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/config/ReqShieldConfiguration.kt b/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/config/ReqShieldConfiguration.kt index abfe4ab..332a3b4 100644 --- a/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/config/ReqShieldConfiguration.kt +++ b/core-reactor/src/main/kotlin/com/linecorp/cse/reqshield/reactor/config/ReqShieldConfiguration.kt @@ -31,8 +31,22 @@ import reactor.core.scheduler.Schedulers data class ReqShieldConfiguration( val setCacheFunction: (String, ReqShieldData, Long) -> Mono, val getCacheFunction: (String) -> Mono?>, - val globalLockFunction: ((String, Long) -> Mono)? = null, - val globalUnLockFunction: ((String) -> Mono)? = null, + /** + * Acquires the distributed lock. Called with (lockKey, token, ttlMillis) and must emit + * true only when this caller acquired the lock. + * + * The token identifies the owner, so the lock must be stored together with it - + * `SET key token NX PX ttl` is the recommended implementation. + */ + val globalLockFunction: ((String, String, Long) -> Mono)? = null, + /** + * Releases the distributed lock. Called with (lockKey, token). + * + * It must release the lock only when the stored value still equals the token + * (compare-and-delete, e.g. a Lua script doing `get` + `del`), otherwise a caller whose + * lock already expired could release the lock of the next owner. + */ + val globalUnLockFunction: ((String, String) -> Mono)? = null, val isLocalLock: Boolean = true, val lockTimeoutMillis: Long = DEFAULT_LOCK_TIMEOUT_MILLIS, val scheduler: Scheduler = Schedulers.boundedElastic(), diff --git a/core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/KeyGlobalLockTest.kt b/core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/KeyGlobalLockTest.kt index 613bf53..f81ef67 100644 --- a/core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/KeyGlobalLockTest.kt +++ b/core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/KeyGlobalLockTest.kt @@ -19,8 +19,11 @@ package com.linecorp.cse.reqshield.reactor import com.linecorp.cse.reqshield.support.BaseKeyLockTest import com.linecorp.cse.reqshield.support.redis.AbstractRedisTest import io.lettuce.core.RedisClient +import io.lettuce.core.ScriptOutputType +import io.lettuce.core.SetArgs import io.lettuce.core.api.async.RedisAsyncCommands import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNotEquals import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test @@ -29,14 +32,14 @@ import reactor.core.scheduler.Schedulers import reactor.test.StepVerifier import java.time.Duration import java.util.concurrent.atomic.AtomicInteger -import kotlin.test.Ignore +import kotlin.test.assertNotNull class KeyGlobalLockTest : AbstractRedisTest(), BaseKeyLockTest { private lateinit var redisCommands: RedisAsyncCommands - private lateinit var globalLockFunc: (String, Long) -> Mono - private lateinit var globalUnLockFunc: (String) -> Mono + private lateinit var globalLockFunc: (String, String, Long) -> Mono + private lateinit var globalUnLockFunc: (String, String) -> Mono @BeforeEach fun init() { @@ -50,12 +53,23 @@ class KeyGlobalLockTest : // Clean up all keys from previous tests for proper test isolation connection.sync().flushdb() - globalLockFunc = { key, timeToLiveMillis -> - Mono.fromFuture { redisCommands.setnx(key, key).toCompletableFuture() } + // Recommended lock implementation: the token is stored as the value so that + // ownership can be checked on release, and PX makes the lock self-expiring. + globalLockFunc = { lockKey, token, ttlMillis -> + Mono + .fromFuture { + redisCommands.set(lockKey, token, SetArgs.Builder.nx().px(ttlMillis)).toCompletableFuture() + }.map { it == "OK" } } - globalUnLockFunc = { key -> - Mono.fromFuture { redisCommands.del(key).toCompletableFuture() }.map { true } + // Compare-and-delete: only the owner of the stored token may release the lock. + globalUnLockFunc = { lockKey, token -> + Mono + .fromFuture { + redisCommands + .eval(UNLOCK_SCRIPT, ScriptOutputType.INTEGER, arrayOf(lockKey), token) + .toCompletableFuture() + }.map { it == 1L } } } @@ -72,13 +86,12 @@ class KeyGlobalLockTest : tasksCompletedCount.incrementAndGet() keyLock .tryLock(key, lockType) - .filter { it } - .flatMap { + .flatMap { token -> lockAcquiredCount.incrementAndGet() doWork() .publishOn(Schedulers.boundedElastic()) .doFinally { _ -> - keyLock.unLock(key, lockType).subscribe() + keyLock.unLock(key, lockType, token).subscribe() } }.onErrorResume { Mono.just(Unit) } } @@ -97,7 +110,7 @@ class KeyGlobalLockTest : Mono .delay(Duration.ofMillis(100)) .then(keyLock.tryLock(key, lockType)), - ).expectNext(true) + ).expectNextCount(1) .verifyComplete() } @@ -114,13 +127,12 @@ class KeyGlobalLockTest : val key = if (i % 2 == 0) "myKey1" else "myKey2" keyLock .tryLock(key, lockType) - .filter { it } - .flatMap { + .flatMap { token -> lockAcquiredCount.incrementAndGet() doWork() .publishOn(Schedulers.boundedElastic()) .doFinally { _ -> - keyLock.unLock(key, lockType).subscribe() + keyLock.unLock(key, lockType, token).subscribe() } }.onErrorResume { Mono.just(Unit) } } @@ -139,7 +151,7 @@ class KeyGlobalLockTest : Mono .delay(Duration.ofMillis(100)) .then(keyLock.tryLock("myKey1", lockType)), - ).expectNext(true) + ).expectNextCount(1) .verifyComplete() StepVerifier @@ -147,14 +159,40 @@ class KeyGlobalLockTest : Mono .delay(Duration.ofMillis(100)) .then(keyLock.tryLock("myKey2", lockType)), - ).expectNext(true) + ).expectNextCount(1) .verifyComplete() } @Test - @Ignore override fun testLockExpiration() { - // Global locks do not have an expiration + val shortTimeoutMillis = 500L + val keyLock = KeyGlobalLock(globalLockFunc, globalUnLockFunc, shortTimeoutMillis) + val key = "expirationKey" + val lockType = LockType.CREATE + + val expiredToken = keyLock.tryLock(key, lockType).block() + assertNotNull(expiredToken) + + // While the lock is held, nobody else can acquire it + StepVerifier.create(keyLock.tryLock(key, lockType)).verifyComplete() + + // Wait for the Redis key TTL to elapse + Thread.sleep(shortTimeoutMillis + 300L) + + val newToken = keyLock.tryLock(key, lockType).block() + assertNotNull(newToken) + assertNotEquals(expiredToken, newToken) + + // The expired holder must not release the lock that now belongs to someone else + StepVerifier + .create(keyLock.unLock(key, lockType, expiredToken)) + .expectNext(false) + .verifyComplete() + + StepVerifier + .create(keyLock.unLock(key, lockType, newToken)) + .expectNext(true) + .verifyComplete() } private fun doWork(): Mono = @@ -162,4 +200,9 @@ class KeyGlobalLockTest : .delay(Duration.ofSeconds(1)) .then(Mono.just(Unit)) .subscribeOn(Schedulers.boundedElastic()) + + companion object { + private const val UNLOCK_SCRIPT = + "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end" + } } diff --git a/core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/KeyLocalLockTest.kt b/core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/KeyLocalLockTest.kt index f245284..37df328 100644 --- a/core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/KeyLocalLockTest.kt +++ b/core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/KeyLocalLockTest.kt @@ -19,6 +19,7 @@ package com.linecorp.cse.reqshield.reactor import com.linecorp.cse.reqshield.support.BaseKeyLockTest import org.junit.jupiter.api.AfterEach import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNotEquals import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.Test import reactor.core.publisher.Mono @@ -26,6 +27,7 @@ import reactor.core.scheduler.Schedulers import reactor.test.StepVerifier import java.time.Duration import java.util.concurrent.atomic.AtomicInteger +import kotlin.test.assertNotNull class KeyLocalLockTest : BaseKeyLockTest { @AfterEach @@ -41,10 +43,13 @@ class KeyLocalLockTest : BaseKeyLockTest { val key = "shared-key" val lockType = LockType.CREATE - StepVerifier.create(instance1.tryLock(key, lockType)).expectNext(true).verifyComplete() - StepVerifier.create(instance2.tryLock(key, lockType)).expectNext(false).verifyComplete() + val token = instance1.tryLock(key, lockType).block() + assertNotNull(token) - StepVerifier.create(instance1.unLock(key, lockType)).expectNext(true).verifyComplete() + // Another instance sees the same lock as held: empty means "not acquired" + StepVerifier.create(instance2.tryLock(key, lockType)).verifyComplete() + + StepVerifier.create(instance1.unLock(key, lockType, token)).expectNext(true).verifyComplete() } @Test @@ -57,22 +62,64 @@ class KeyLocalLockTest : BaseKeyLockTest { val attempts = listOf(instance1, instance2, instance3).map { inst -> - inst.tryLock(key, lockType).map { acquired -> if (acquired) 1 else 0 } + inst + .tryLock(key, lockType) + .map { 1 } + .defaultIfEmpty(0) } StepVerifier .create(Mono.zip(attempts) { arr -> arr.sumOf { it as Int } }) .expectNextMatches { it == 1 } .verifyComplete() + } + + @Test + fun `should not release a lock held by another owner`() { + val keyLock = KeyLocalLock(lockTimeoutMillis) + val key = "foreign-token-test" + val lockType = LockType.CREATE + + val token = keyLock.tryLock(key, lockType).block() + assertNotNull(token) + + // A token that never owned this lock must not release it + StepVerifier + .create(keyLock.unLock(key, lockType, "someone-elses-token")) + .expectNext(false) + .verifyComplete() - // cleanup by unlocking whoever acquired - listOf(instance1, instance2, instance3).forEach { inst -> inst.unLock(key, lockType).subscribe() } + // The lock is therefore still held + StepVerifier.create(keyLock.tryLock(key, lockType)).verifyComplete() + + StepVerifier.create(keyLock.unLock(key, lockType, token)).expectNext(true).verifyComplete() + } + + @Test + fun `should not let a stale token release the lock of the next owner`() { + val shortTimeoutMillis = 300L + val keyLock = KeyLocalLock(shortTimeoutMillis) + val key = "stale-token-test" + val lockType = LockType.CREATE + + val staleToken = keyLock.tryLock(key, lockType).block() + assertNotNull(staleToken) + + // Let the lock expire so that it can be force-released and handed to a new owner + Thread.sleep(shortTimeoutMillis + 100L) + + val newToken = keyLock.tryLock(key, lockType).block() + assertNotNull(newToken) + assertNotEquals(staleToken, newToken) + + StepVerifier.create(keyLock.unLock(key, lockType, staleToken)).expectNext(false).verifyComplete() + StepVerifier.create(keyLock.unLock(key, lockType, newToken)).expectNext(true).verifyComplete() } @Test override fun testConcurrencyWithOneKey() { val keyLock = KeyLocalLock(lockTimeoutMillis) - val key = "myKey" + val key = "one-key-test" val lockType = LockType.CREATE val lockAcquiredCount = AtomicInteger(0) val tasksCompletedCount = AtomicInteger(0) @@ -82,13 +129,12 @@ class KeyLocalLockTest : BaseKeyLockTest { tasksCompletedCount.incrementAndGet() keyLock .tryLock(key, lockType) - .filter { it } - .flatMap { + .flatMap { token -> lockAcquiredCount.incrementAndGet() doWork() .publishOn(Schedulers.boundedElastic()) .doFinally { _ -> - keyLock.unLock(key, lockType).subscribe() + keyLock.unLock(key, lockType, token).subscribe() } }.onErrorResume { Mono.just(Unit) } } @@ -107,7 +153,7 @@ class KeyLocalLockTest : BaseKeyLockTest { Mono .delay(Duration.ofMillis(100)) .then(keyLock.tryLock(key, lockType)), - ).expectNext(true) + ).expectNextCount(1) .verifyComplete() } @@ -121,16 +167,15 @@ class KeyLocalLockTest : BaseKeyLockTest { val tasks = (0 until 20).map { i -> tasksCompletedCount.incrementAndGet() - val key = if (i % 2 == 0) "myKey1" else "myKey2" + val key = if (i % 2 == 0) "two-key-test1" else "two-key-test2" keyLock .tryLock(key, lockType) - .filter { it } - .flatMap { + .flatMap { token -> lockAcquiredCount.incrementAndGet() doWork() .publishOn(Schedulers.boundedElastic()) .doFinally { _ -> - keyLock.unLock(key, lockType).subscribe() + keyLock.unLock(key, lockType, token).subscribe() } }.onErrorResume { Mono.just(Unit) } } @@ -148,46 +193,39 @@ class KeyLocalLockTest : BaseKeyLockTest { .create( Mono .delay(Duration.ofMillis(100)) - .then(keyLock.tryLock("myKey1", lockType)), - ).expectNext(true) + .then(keyLock.tryLock("two-key-test1", lockType)), + ).expectNextCount(1) .verifyComplete() StepVerifier .create( Mono .delay(Duration.ofMillis(100)) - .then(keyLock.tryLock("myKey2", lockType)), - ).expectNext(true) + .then(keyLock.tryLock("two-key-test2", lockType)), + ).expectNextCount(1) .verifyComplete() } @Test override fun testLockExpiration() { val keyLock = KeyLocalLock(lockTimeoutMillis) - val key = "myKey" + val key = "expiration-test" val lockType = LockType.CREATE - StepVerifier - .create( - keyLock.tryLock(key, lockType), - ).expectNext(true) - .verifyComplete() + val expiredToken = keyLock.tryLock(key, lockType).block() + assertNotNull(expiredToken) // Wait for lock timeout + cleanup interval + buffer // lockTimeoutMillis = 3000ms, cleanup interval = 1000ms Thread.sleep(lockTimeoutMillis + 1000L + 500L) // 4.5 seconds total - StepVerifier - .create( - keyLock.tryLock(key, lockType), - ).expectNext(true) - .verifyComplete() + val newToken = keyLock.tryLock(key, lockType).block() + assertNotNull(newToken) - StepVerifier - .create( - keyLock.unLock(key, lockType), - ).expectNext(true) - .verifyComplete() + // The expired holder must not be able to release the lock of the new holder + StepVerifier.create(keyLock.unLock(key, lockType, expiredToken)).expectNext(false).verifyComplete() + + StepVerifier.create(keyLock.unLock(key, lockType, newToken)).expectNext(true).verifyComplete() } @Test @@ -197,31 +235,31 @@ class KeyLocalLockTest : BaseKeyLockTest { val lockType = LockType.CREATE // Acquire lock - StepVerifier.create(keyLock.tryLock(key, lockType)) - .expectNext(true) - .verifyComplete() + val token = keyLock.tryLock(key, lockType).block() + assertNotNull(token) // First unlock should succeed - StepVerifier.create(keyLock.unLock(key, lockType)) + StepVerifier + .create(keyLock.unLock(key, lockType, token)) .expectNext(true) .verifyComplete() // Second unlock should return false (over-release prevention) - StepVerifier.create(keyLock.unLock(key, lockType)) + StepVerifier + .create(keyLock.unLock(key, lockType, token)) .expectNext(false) .verifyComplete() // Verify semaphore is not over-released: can acquire once, not twice - StepVerifier.create(keyLock.tryLock(key, lockType)) - .expectNext(true) - .verifyComplete() + val reacquiredToken = keyLock.tryLock(key, lockType).block() + assertNotNull(reacquiredToken) - StepVerifier.create(keyLock.tryLock(key, lockType)) - .expectNext(false) + StepVerifier + .create(keyLock.tryLock(key, lockType)) .verifyComplete() // Cleanup - keyLock.unLock(key, lockType).subscribe() + keyLock.unLock(key, lockType, reacquiredToken).subscribe() } @Test @@ -232,17 +270,18 @@ class KeyLocalLockTest : BaseKeyLockTest { val successfulAcquisitions = AtomicInteger(0) // Simulate over-release attempt - StepVerifier.create(keyLock.tryLock(key, lockType)) - .expectNext(true) - .verifyComplete() + val token = keyLock.tryLock(key, lockType).block() + assertNotNull(token) - StepVerifier.create(keyLock.unLock(key, lockType)) + StepVerifier + .create(keyLock.unLock(key, lockType, token)) .expectNext(true) .verifyComplete() // Multiple unlock attempts should all return false repeat(5) { - StepVerifier.create(keyLock.unLock(key, lockType)) + StepVerifier + .create(keyLock.unLock(key, lockType, token)) .expectNext(false) .verifyComplete() } @@ -250,8 +289,10 @@ class KeyLocalLockTest : BaseKeyLockTest { // Try to acquire lock concurrently - only ONE should succeed val attempts = (1..10).map { - keyLock.tryLock(key, lockType) - .map { acquired -> if (acquired) successfulAcquisitions.incrementAndGet() else 0 } + keyLock + .tryLock(key, lockType) + .map { successfulAcquisitions.incrementAndGet() } + .defaultIfEmpty(0) } StepVerifier @@ -261,9 +302,6 @@ class KeyLocalLockTest : BaseKeyLockTest { // Only one should have acquired the lock assertEquals(1, successfulAcquisitions.get(), "Only one should acquire the lock") - - // Cleanup - keyLock.unLock(key, lockType).subscribe() } private fun doWork(): Mono = diff --git a/core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/ReqShieldTest.kt b/core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/ReqShieldTest.kt index be82f40..aa79dcc 100644 --- a/core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/ReqShieldTest.kt +++ b/core-reactor/src/test/kotlin/com/linecorp/cse/reqshield/reactor/ReqShieldTest.kt @@ -19,16 +19,16 @@ 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.BaseReqShieldTest +import com.linecorp.cse.reqshield.support.constant.ConfigValues.LOCK_KEY_PREFIX 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 io.mockk.every import io.mockk.mockk -import io.mockk.mockkStatic -import io.mockk.unmockkStatic +import io.mockk.slot import io.mockk.verify -import org.junit.jupiter.api.AfterEach import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.BeforeEach @@ -39,8 +39,8 @@ import reactor.test.StepVerifier import java.lang.reflect.InvocationTargetException import java.lang.reflect.Method import java.time.Duration -import java.time.LocalDateTime import java.util.concurrent.Callable +import java.util.concurrent.atomic.AtomicInteger import kotlin.test.assertNotNull import kotlin.test.assertNull @@ -58,18 +58,19 @@ class ReqShieldTest : BaseReqShieldTest { private val oldValue = Product("oldTestValue", "oldTestValue") private val value = Product("testValue", "testValue") private val callable: Callable> = mockk() + private val token = "token" private var timeToLiveMillis: Long = 10000 - private lateinit var globalLockFunc: (String, Long) -> Mono - private lateinit var globalUnLockFunc: (String) -> Mono + private lateinit var globalLockFunc: (String, String, Long) -> Mono + private lateinit var globalUnLockFunc: (String, String) -> Mono @BeforeEach fun setup() { cacheSetter = mockk<(String, ReqShieldData, Long) -> Mono>() cacheGetter = mockk<(String) -> Mono?>>() - globalLockFunc = mockk<(String, Long) -> Mono>() - globalUnLockFunc = mockk<(String) -> Mono>() + globalLockFunc = mockk<(String, String, Long) -> Mono>() + globalUnLockFunc = mockk<(String, String) -> Mono>() keyLock = mockk() keyGlobalLock = KeyGlobalLock(globalLockFunc, globalUnLockFunc, 3000) @@ -116,22 +117,38 @@ class ReqShieldTest : BaseReqShieldTest { keyLock = keyGlobalLock, ), ) - - mockkStatic(LocalDateTime::class) - every { LocalDateTime.now() } returns LocalDateTime.of(2023, 11, 13, 12, 0, 0, 0) } - @AfterEach - fun tearDown() { - unmockkStatic(LocalDateTime::class) - } + /** ReqShield instance that gives up waiting for the lock holder quickly. */ + private fun reqShieldWithMaxAttempt(maxAttemptGetCache: Int): ReqShield = + ReqShield( + ReqShieldConfiguration( + cacheSetter, + cacheGetter, + keyLock = keyLock, + maxAttemptGetCache = maxAttemptGetCache, + ), + ) + + /** Cached entry that is not yet old enough to be refreshed. */ + private fun freshReqShieldData(cachedValue: Product?): ReqShieldData = + ReqShieldData(cachedValue, ReqShieldData.Status.NEW, nowToEpochTime(), timeToLiveMillis) + + /** Cached entry that has passed the decisionForUpdate threshold (90% of its TTL). */ + private fun updateTargetReqShieldData(cachedValue: Product?): ReqShieldData = + ReqShieldData( + cachedValue, + ReqShieldData.Status.NEW, + nowToEpochTime() - (timeToLiveMillis * 0.9).toLong(), + timeToLiveMillis, + ) @Test override fun testSetMethodCacheNotExistsAndLocalLockAcquired() { every { cacheGetter.invoke(key) } returns Mono.empty() every { cacheSetter.invoke(key, any(), any()) } returns Mono.just(true) - every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.just(true) - every { keyLock.unLock(key, LockType.CREATE) } returns Mono.just(true) + every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.just(token) + every { keyLock.unLock(key, LockType.CREATE, token) } returns Mono.just(true) val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) @@ -141,17 +158,12 @@ class ReqShieldTest : BaseReqShieldTest { assertNotNull(it) }.verifyComplete() - StepVerifier - .create(Mono.delay(Duration.ofMillis(100))) - .expectSubscription() - .thenAwait(Duration.ofMillis(100)) - .expectNextCount(1) - .verifyComplete() + awaitFireAndForget() verify { cacheGetter.invoke(key) } verify { cacheSetter.invoke(key, any(), any()) } verify { keyLock.tryLock(key, LockType.CREATE) } - verify { keyLock.unLock(key, LockType.CREATE) } + verify { keyLock.unLock(key, LockType.CREATE, token) } verify { callable.call() } } @@ -168,17 +180,12 @@ class ReqShieldTest : BaseReqShieldTest { assertNotNull(it) }.verifyComplete() - StepVerifier - .create(Mono.delay(Duration.ofMillis(100))) - .expectSubscription() - .thenAwait(Duration.ofMillis(100)) - .expectNextCount(1) - .verifyComplete() + awaitFireAndForget() verify { cacheGetter.invoke(key) } verify { cacheSetter.invoke(key, any(), any()) } verify(inverse = true) { keyLock.tryLock(key, LockType.CREATE) } - verify(inverse = true) { keyLock.unLock(key, LockType.CREATE) } + verify(inverse = true) { keyLock.unLock(key, LockType.CREATE, any()) } verify { callable.call() } } @@ -187,8 +194,8 @@ class ReqShieldTest : BaseReqShieldTest { every { cacheGetter.invoke(key) } returns Mono.empty() every { cacheSetter.invoke(key, any(), any()) } returns Mono.just(true) - every { globalLockFunc(any(), any()) } returns Mono.just(true) - every { globalUnLockFunc(any()) } returns Mono.just(true) + every { globalLockFunc(any(), any(), any()) } returns Mono.just(true) + every { globalUnLockFunc(any(), any()) } returns Mono.just(true) val result = reqShieldForGlobalLock.getAndSetReqShieldData(key, callable, timeToLiveMillis) @@ -198,19 +205,16 @@ class ReqShieldTest : BaseReqShieldTest { assertNotNull(it) }.verifyComplete() - StepVerifier - .create(Mono.delay(Duration.ofMillis(100))) - .expectSubscription() - .thenAwait(Duration.ofMillis(100)) - .expectNextCount(1) - .verifyComplete() + awaitFireAndForget() + + val lockKey = "$LOCK_KEY_PREFIX${key}_${LockType.CREATE.name}" + val tokenSlot = slot() verify { cacheGetter.invoke(key) } verify { cacheSetter.invoke(key, any(), any()) } - verify { globalLockFunc(any(), any()) } - verify { globalUnLockFunc(any()) } - verify { keyGlobalLock.tryLock(key, LockType.CREATE) } - verify { keyGlobalLock.unLock(key, LockType.CREATE) } + verify { globalLockFunc(lockKey, capture(tokenSlot), 3000) } + // The very token handed out by tryLock must be the one used to release the lock + verify { globalUnLockFunc(lockKey, tokenSlot.captured) } verify { callable.call() } } @@ -236,8 +240,8 @@ class ReqShieldTest : BaseReqShieldTest { override fun testSetMethodCacheNotExistsAndLocalLockAcquiredAndCallableReturnNull() { every { cacheGetter.invoke(key) } returns Mono.empty() every { cacheSetter.invoke(key, any(), any()) } returns Mono.just(true) - every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.just(true) - every { keyLock.unLock(key, LockType.CREATE) } returns Mono.empty() + every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.just(token) + every { keyLock.unLock(key, LockType.CREATE, token) } returns Mono.empty() every { callable.call() } returns Mono.empty() val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) @@ -249,17 +253,12 @@ class ReqShieldTest : BaseReqShieldTest { assertNull(it.value) }.verifyComplete() - StepVerifier - .create(Mono.delay(Duration.ofMillis(100))) - .expectSubscription() - .thenAwait(Duration.ofMillis(100)) - .expectNextCount(1) - .verifyComplete() + awaitFireAndForget() verify { cacheGetter.invoke(key) } verify { cacheSetter.invoke(key, any(), any()) } verify { keyLock.tryLock(key, LockType.CREATE) } - verify { keyLock.unLock(key, LockType.CREATE) } + verify { keyLock.unLock(key, LockType.CREATE, token) } verify { callable.call() } } @@ -268,8 +267,8 @@ class ReqShieldTest : BaseReqShieldTest { every { cacheGetter.invoke(key) } returns Mono.empty() every { cacheSetter.invoke(key, any(), any()) } returns Mono.just(true) - every { globalLockFunc(any(), any()) } returns Mono.just(true) - every { globalUnLockFunc(any()) } returns Mono.just(true) + every { globalLockFunc(any(), any(), any()) } returns Mono.just(true) + every { globalUnLockFunc(any(), any()) } returns Mono.just(true) every { callable.call() } returns Mono.empty() @@ -282,19 +281,12 @@ class ReqShieldTest : BaseReqShieldTest { assertNull(it.value) }.verifyComplete() - StepVerifier - .create(Mono.delay(Duration.ofMillis(100))) - .expectSubscription() - .thenAwait(Duration.ofMillis(100)) - .expectNextCount(1) - .verifyComplete() + awaitFireAndForget() verify { cacheGetter.invoke(key) } verify { cacheSetter.invoke(key, any(), any()) } - verify { globalLockFunc(any(), any()) } - verify { globalUnLockFunc(any()) } - verify { keyGlobalLock.tryLock(key, LockType.CREATE) } - verify { keyGlobalLock.unLock(key, LockType.CREATE) } + verify { globalLockFunc(any(), any(), any()) } + verify { globalUnLockFunc(any(), any()) } verify { callable.call() } } @@ -302,8 +294,8 @@ class ReqShieldTest : BaseReqShieldTest { override fun testSetMethodCacheNotExistsAndLocalLockAcquiredAndThrowCallableClientException() { every { cacheGetter.invoke(key) } returns Mono.empty() every { cacheSetter.invoke(key, any(), any()) } returns Mono.just(true) - every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.just(true) - every { keyLock.unLock(key, LockType.CREATE) } returns Mono.empty() + every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.just(token) + every { keyLock.unLock(key, LockType.CREATE, token) } returns Mono.empty() every { callable.call() } returns Mono.error(Exception("callable error")) StepVerifier @@ -312,16 +304,11 @@ class ReqShieldTest : BaseReqShieldTest { throwable is ClientException && throwable.errorCode == ErrorCode.SUPPLIER_ERROR }.verify() - StepVerifier - .create(Mono.delay(Duration.ofMillis(100))) - .expectSubscription() - .thenAwait(Duration.ofMillis(100)) - .expectNextCount(1) - .verifyComplete() + awaitFireAndForget() verify { cacheGetter.invoke(key) } verify { keyLock.tryLock(key, LockType.CREATE) } - verify { keyLock.unLock(key, LockType.CREATE) } + verify { keyLock.unLock(key, LockType.CREATE, token) } verify { callable.call() } } @@ -330,8 +317,8 @@ class ReqShieldTest : BaseReqShieldTest { every { cacheGetter.invoke(key) } returns Mono.empty() every { cacheSetter.invoke(key, any(), any()) } returns Mono.just(true) - every { globalLockFunc(any(), any()) } returns Mono.just(true) - every { globalUnLockFunc(any()) } returns Mono.just(true) + every { globalLockFunc(any(), any(), any()) } returns Mono.just(true) + every { globalUnLockFunc(any(), any()) } returns Mono.just(true) every { callable.call() } returns Mono.error(Exception("callable error")) @@ -341,18 +328,11 @@ class ReqShieldTest : BaseReqShieldTest { throwable is ClientException && throwable.errorCode == ErrorCode.SUPPLIER_ERROR }.verify() - StepVerifier - .create(Mono.delay(Duration.ofMillis(100))) - .expectSubscription() - .thenAwait(Duration.ofMillis(100)) - .expectNextCount(1) - .verifyComplete() + awaitFireAndForget() verify { cacheGetter.invoke(key) } - verify { globalLockFunc(any(), any()) } - verify { globalUnLockFunc(any()) } - verify { keyGlobalLock.tryLock(key, LockType.CREATE) } - verify { keyGlobalLock.unLock(key, LockType.CREATE) } + verify { globalLockFunc(any(), any(), any()) } + verify { globalUnLockFunc(any(), any()) } verify { callable.call() } } @@ -360,8 +340,6 @@ class ReqShieldTest : BaseReqShieldTest { override fun testSetMethodCacheNotExistsAndLocalLockAcquiredAndThrowGetCacheClientException() { every { cacheGetter.invoke(key) } returns Mono.error(Exception("get cache error")) every { cacheSetter.invoke(key, any(), any()) } returns Mono.just(true) - every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.just(true) - every { keyLock.unLock(key, LockType.CREATE) } returns Mono.just(true) StepVerifier .create(reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis)) @@ -369,16 +347,11 @@ class ReqShieldTest : BaseReqShieldTest { throwable is ClientException && throwable.errorCode == ErrorCode.GET_CACHE_ERROR }.verify() - StepVerifier - .create(Mono.delay(Duration.ofMillis(100))) - .expectSubscription() - .thenAwait(Duration.ofMillis(100)) - .expectNextCount(1) - .verifyComplete() + awaitFireAndForget() verify { cacheGetter.invoke(key) } verify(inverse = true) { keyLock.tryLock(key, LockType.CREATE) } - verify(inverse = true) { keyLock.unLock(key, LockType.CREATE) } + verify(inverse = true) { keyLock.unLock(key, LockType.CREATE, any()) } verify(inverse = true) { callable.call() } } @@ -387,41 +360,33 @@ class ReqShieldTest : BaseReqShieldTest { every { cacheGetter.invoke(key) } returns Mono.error(Exception("get cache error")) every { cacheSetter.invoke(key, any(), any()) } returns Mono.just(true) - every { globalLockFunc(any(), any()) } returns Mono.just(true) - every { globalUnLockFunc(any()) } returns Mono.just(true) - StepVerifier .create(reqShieldForGlobalLock.getAndSetReqShieldData(key, callable, timeToLiveMillis)) .expectErrorMatches { throwable -> throwable is ClientException && throwable.errorCode == ErrorCode.GET_CACHE_ERROR }.verify() - StepVerifier - .create(Mono.delay(Duration.ofMillis(100))) - .expectSubscription() - .thenAwait(Duration.ofMillis(100)) - .expectNextCount(1) - .verifyComplete() + awaitFireAndForget() verify { cacheGetter.invoke(key) } - verify(inverse = true) { globalLockFunc(any(), any()) } - verify(inverse = true) { globalUnLockFunc(any()) } - verify(inverse = true) { keyGlobalLock.tryLock(key, LockType.CREATE) } - verify(inverse = true) { keyGlobalLock.unLock(key, LockType.CREATE) } + verify(inverse = true) { globalLockFunc(any(), any(), any()) } + verify(inverse = true) { globalUnLockFunc(any(), any()) } verify(inverse = true) { callable.call() } } @Test override fun testSetMethodCacheNotExistsAndLocalLockNotAcquired() { every { cacheGetter.invoke(key) } returns Mono.empty() - every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.just(false) + every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.empty() - val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) + val reqShieldWithFewAttempts = reqShieldWithMaxAttempt(3) + val result = reqShieldWithFewAttempts.getAndSetReqShieldData(key, callable, timeToLiveMillis) StepVerifier .create(result) .assertNext { assertNotNull(it) + assertEquals(value, it.value) }.verifyComplete() verify { cacheGetter.invoke(key) } @@ -433,30 +398,43 @@ class ReqShieldTest : BaseReqShieldTest { @Test override fun testSetMethodCacheNotExistsAndGlobalLockNotAcquired() { every { cacheGetter.invoke(key) } returns Mono.empty() - every { globalLockFunc(any(), any()) } returns Mono.just(false) + every { globalLockFunc(any(), any(), any()) } returns Mono.just(false) - val result = reqShieldForGlobalLock.getAndSetReqShieldData(key, callable, timeToLiveMillis) + val reqShieldWithFewAttempts = + ReqShield( + ReqShieldConfiguration( + cacheSetter, + cacheGetter, + globalLockFunc, + globalUnLockFunc, + isLocalLock = false, + keyLock = keyGlobalLock, + maxAttemptGetCache = 3, + ), + ) + + val result = reqShieldWithFewAttempts.getAndSetReqShieldData(key, callable, timeToLiveMillis) StepVerifier .create(result) .assertNext { assertNotNull(it) + assertEquals(value, it.value) }.verifyComplete() verify { cacheGetter.invoke(key) } verify(inverse = true) { cacheSetter.invoke(key, any(), any()) } - verify { globalLockFunc(any(), any()) } - verify { keyGlobalLock.tryLock(key, LockType.CREATE) } + verify { globalLockFunc(any(), any(), any()) } + verify(inverse = true) { globalUnLockFunc(any(), any()) } verify { callable.call() } } @Test override fun testSetMethodCacheExistsButNotTargetedForUpdate() { timeToLiveMillis = 1000 - val reqShieldData = ReqShieldData(value, timeToLiveMillis) + val reqShieldData = freshReqShieldData(value) every { cacheGetter.invoke(key) } returns Mono.just(reqShieldData) - every { keyLock.tryLock(key, LockType.UPDATE) } returns Mono.just(false) val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) @@ -468,20 +446,21 @@ class ReqShieldTest : BaseReqShieldTest { }.expectComplete() .verify() - verify { keyLock.tryLock(key, LockType.UPDATE) } + // A fresh entry is not refreshed, so the update lock must never be taken + verify(inverse = true) { keyLock.tryLock(key, LockType.UPDATE) } + verify(inverse = true) { callable.call() } verify { cacheGetter.invoke(key) } } @Test override fun testSetMethodCacheExistsAndTheUpdateTarget() { timeToLiveMillis = 1000 - val reqShieldData = ReqShieldData(oldValue, timeToLiveMillis) - val newReqShieldData = ReqShieldData(value, timeToLiveMillis) + val reqShieldData = updateTargetReqShieldData(oldValue) every { cacheGetter.invoke(key) } returns Mono.just(reqShieldData) every { cacheSetter.invoke(key, any(), any()) } answers { Mono.just(true) } - every { keyLock.tryLock(key, LockType.UPDATE) } returns Mono.just(true) - every { keyLock.unLock(key, LockType.UPDATE) } returns Mono.empty() + every { keyLock.tryLock(key, LockType.UPDATE) } returns Mono.just(token) + every { keyLock.unLock(key, LockType.UPDATE, token) } returns Mono.empty() val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) @@ -493,25 +472,25 @@ class ReqShieldTest : BaseReqShieldTest { }.expectComplete() .verify() - StepVerifier - .create(Mono.delay(Duration.ofMillis(100))) - .expectSubscription() - .thenAwait(Duration.ofMillis(100)) - .expectNextCount(1) - .verifyComplete() + awaitFireAndForget() verify { keyLock.tryLock(key, LockType.UPDATE) } verify { cacheGetter.invoke(key) } - verify { cacheSetter.invoke(key, newReqShieldData, timeToLiveMillis) } - verify { keyLock.unLock(key, LockType.UPDATE) } + verify { + cacheSetter.invoke( + key, + match { it.value == value && it.timeToLiveMillis == timeToLiveMillis }, + timeToLiveMillis, + ) + } + verify { keyLock.unLock(key, LockType.UPDATE, token) } verify { callable.call() } } @Test override fun testSetMethodCacheExistsAndTheUpdateTargetOnlyCreateCache() { timeToLiveMillis = 1000 - val reqShieldData = ReqShieldData(oldValue, timeToLiveMillis) - val newReqShieldData = ReqShieldData(value, timeToLiveMillis) + val reqShieldData = updateTargetReqShieldData(oldValue) every { cacheGetter.invoke(key) } returns Mono.just(reqShieldData) every { cacheSetter.invoke(key, any(), any()) } answers { Mono.just(true) } @@ -526,30 +505,30 @@ class ReqShieldTest : BaseReqShieldTest { }.expectComplete() .verify() - StepVerifier - .create(Mono.delay(Duration.ofMillis(100))) - .expectSubscription() - .thenAwait(Duration.ofMillis(100)) - .expectNextCount(1) - .verifyComplete() + awaitFireAndForget() verify(inverse = true) { keyLock.tryLock(key, LockType.UPDATE) } verify { cacheGetter.invoke(key) } - verify { cacheSetter.invoke(key, newReqShieldData, timeToLiveMillis) } - verify(inverse = true) { keyLock.unLock(key, LockType.UPDATE) } + verify { + cacheSetter.invoke( + key, + match { it.value == value && it.timeToLiveMillis == timeToLiveMillis }, + timeToLiveMillis, + ) + } + verify(inverse = true) { keyLock.unLock(key, LockType.UPDATE, any()) } verify { callable.call() } } @Test override fun testSetMethodCacheExistsAndTheUpdateTargetAndCallableReturnNull() { timeToLiveMillis = 1000 - val reqShieldData = ReqShieldData(value, timeToLiveMillis) - val reqShieldDataNull = ReqShieldData(null, timeToLiveMillis) + val reqShieldData = updateTargetReqShieldData(value) every { cacheGetter.invoke(key) } returns Mono.just(reqShieldData) every { cacheSetter.invoke(key, any(), any()) } answers { Mono.just(true) } - every { keyLock.tryLock(key, LockType.UPDATE) } returns Mono.just(true) - every { keyLock.unLock(key, LockType.UPDATE) } returns Mono.empty() + every { keyLock.tryLock(key, LockType.UPDATE) } returns Mono.just(token) + every { keyLock.unLock(key, LockType.UPDATE, token) } returns Mono.empty() every { callable.call() } returns Mono.empty() val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) @@ -561,24 +540,24 @@ class ReqShieldTest : BaseReqShieldTest { .expectComplete() .verify(Duration.ofSeconds(1)) - StepVerifier - .create(Mono.delay(Duration.ofMillis(100))) - .expectSubscription() - .thenAwait(Duration.ofMillis(100)) - .expectNextCount(1) - .verifyComplete() + awaitFireAndForget() verify { cacheGetter.invoke(key) } - verify { cacheSetter.invoke(key, reqShieldDataNull, timeToLiveMillis) } + verify { + cacheSetter.invoke( + key, + match { it.value == null && it.timeToLiveMillis == timeToLiveMillis }, + timeToLiveMillis, + ) + } verify { keyLock.tryLock(key, LockType.UPDATE) } - verify { keyLock.unLock(key, LockType.UPDATE) } + verify { keyLock.unLock(key, LockType.UPDATE, token) } verify { callable.call() } } @Test override fun executeSetCacheFunctionShouldHandleExceptionFromCacheSetter() { - every { keyLock.tryLock(any(), any()) } returns Mono.just(true) - every { keyLock.unLock(any(), any()) } returns Mono.just(true) + every { keyLock.unLock(any(), any(), any()) } returns Mono.just(true) val key = "key" val reqShieldData = ReqShieldData(value, 1000L) @@ -595,7 +574,8 @@ class ReqShieldTest : BaseReqShieldTest { val mono = Mono.defer { try { - method.invoke(reqShield, cacheSetter, key, reqShieldData, lockType) as Mono + @Suppress("UNCHECKED_CAST") + method.invoke(reqShield, cacheSetter, key, reqShieldData, lockType, token) as Mono } catch (e: InvocationTargetException) { Mono.error(e.cause ?: e) } @@ -606,9 +586,111 @@ class ReqShieldTest : BaseReqShieldTest { .expectErrorSatisfies { throwable -> assertTrue(throwable is ClientException) assertEquals(ErrorCode.SET_CACHE_ERROR, (throwable as ClientException).errorCode) + assertNotNull(throwable.cause) }.verify() verify { cacheSetter.invoke(key, reqShieldData, 1000L) } - verify { keyLock.unLock(any(), any()) } + verify { keyLock.unLock(key, lockType, token) } + } + + @Test + fun `should return the value another request wrote to the cache while waiting for the lock`() { + // 1 initial read + 3 polls: the lock holder fills the cache on the 3rd poll + val cachedData = freshReqShieldData(oldValue) + val getCacheInvocations = AtomicInteger(0) + + every { cacheGetter.invoke(key) } answers { + if (getCacheInvocations.incrementAndGet() <= 3) Mono.empty() else Mono.just(cachedData) + } + every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.empty() + + StepVerifier + .create(reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis)) + .expectNext(cachedData) + .verifyComplete() + + assertEquals(4, getCacheInvocations.get()) + // The waiter must not call the supplier when the cache gets filled + verify(inverse = true) { callable.call() } + verify(inverse = true) { cacheSetter.invoke(key, any(), any()) } + } + + @Test + fun `should fall back to the supplier after consecutive get cache failures`() { + val getCacheInvocations = AtomicInteger(0) + + every { cacheGetter.invoke(key) } answers { + // 1st call is the initial read, the 3 following polls all fail + if (getCacheInvocations.incrementAndGet() == 1) { + Mono.empty() + } else { + Mono.error(RuntimeException("cache down")) + } + } + every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.empty() + + // maxAttemptGetCache is far higher than the failure threshold, so bailing out is what stops the polling + StepVerifier + .create(reqShieldWithMaxAttempt(30).getAndSetReqShieldData(key, callable, timeToLiveMillis)) + .assertNext { assertEquals(value, it.value) } + .verifyComplete() + + // 1 initial read + MAX_CONSECUTIVE_GET_CACHE_FAILURES(3) failing polls + assertEquals(4, getCacheInvocations.get()) + verify(exactly = 1) { callable.call() } + verify(inverse = true) { keyLock.unLock(key, LockType.CREATE, any()) } + } + + @Test + fun `should keep polling when get cache failures are not consecutive`() { + val maxAttemptGetCache = 6 + val getCacheInvocations = AtomicInteger(0) + + every { cacheGetter.invoke(key) } answers { + // Initial read, then polls alternating between two failures and an empty (successful) read. + // Without resetting the counter on a successful read, polling would bail out on the 4th poll. + when (getCacheInvocations.incrementAndGet()) { + 1, 4, 7 -> Mono.empty() + else -> Mono.error(RuntimeException("cache down")) + } + } + every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.empty() + + StepVerifier + .create( + reqShieldWithMaxAttempt(maxAttemptGetCache).getAndSetReqShieldData(key, callable, timeToLiveMillis), + ).assertNext { assertEquals(value, it.value) } + .verifyComplete() + + // Every poll was attempted: 1 initial read + first poll + maxAttemptGetCache retries + assertEquals(2 + maxAttemptGetCache, getCacheInvocations.get()) + verify(exactly = 1) { callable.call() } + } + + @Test + fun `should surface supplier error when the supplier fails after waiting for the lock`() { + every { cacheGetter.invoke(key) } returns Mono.empty() + every { keyLock.tryLock(key, LockType.CREATE) } returns Mono.empty() + every { callable.call() } returns Mono.error(IllegalStateException("supplier down")) + + StepVerifier + .create(reqShieldWithMaxAttempt(2).getAndSetReqShieldData(key, callable, timeToLiveMillis)) + .expectErrorMatches { + it is ClientException && it.errorCode == ErrorCode.SUPPLIER_ERROR && it.cause != null + }.verify() + + verify(exactly = 1) { callable.call() } + // The waiter holds no lock, so it must not try to release one + verify(inverse = true) { keyLock.unLock(key, LockType.CREATE, any()) } + } + + /** Gives the fire-and-forget cache writes / unlocks a chance to run before verifying them. */ + private fun awaitFireAndForget() { + StepVerifier + .create(Mono.delay(Duration.ofMillis(100))) + .expectSubscription() + .thenAwait(Duration.ofMillis(100)) + .expectNextCount(1) + .verifyComplete() } } diff --git a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/annotation/ReqShieldCacheEvict.kt b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/annotation/ReqShieldCacheEvict.kt index b6e09af..aff0108 100644 --- a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/annotation/ReqShieldCacheEvict.kt +++ b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/annotation/ReqShieldCacheEvict.kt @@ -25,7 +25,4 @@ annotation class ReqShieldCacheEvict( val cacheName: String, val key: String = "", val keyGenerator: String = "", - val isLocalLock: Boolean = true, - val lockTimeoutMillis: Long = 3000, - val condition: String = "", ) diff --git a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/annotation/ReqShieldCacheable.kt b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/annotation/ReqShieldCacheable.kt index ff101b7..7c0cf9a 100644 --- a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/annotation/ReqShieldCacheable.kt +++ b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/annotation/ReqShieldCacheable.kt @@ -17,6 +17,9 @@ package com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.annotation import com.linecorp.cse.reqshield.kotlin.coroutine.config.ReqShieldWorkMode +import com.linecorp.cse.reqshield.support.constant.ConfigValues.DEFAULT_DECISION_FOR_UPDATE +import com.linecorp.cse.reqshield.support.constant.ConfigValues.DEFAULT_LOCK_TIMEOUT_MILLIS +import com.linecorp.cse.reqshield.support.constant.ConfigValues.DEFAULT_TIME_TO_LIVE_MILLIS import com.linecorp.cse.reqshield.support.constant.ConfigValues.MAX_ATTEMPT_GET_CACHE import java.lang.annotation.Inherited @@ -28,9 +31,9 @@ annotation class ReqShieldCacheable( val key: String = "", val keyGenerator: String = "", val isLocalLock: Boolean = true, - val lockTimeoutMillis: Long = 3000, - val decisionForUpdate: Int = 90, + val lockTimeoutMillis: Long = DEFAULT_LOCK_TIMEOUT_MILLIS, + val decisionForUpdate: Int = DEFAULT_DECISION_FOR_UPDATE, val maxAttemptGetCache: Int = MAX_ATTEMPT_GET_CACHE, - val timeToLiveMillis: Long = 10 * 60 * 1000, + val timeToLiveMillis: Long = DEFAULT_TIME_TO_LIVE_MILLIS, val reqShieldWorkMode: ReqShieldWorkMode = ReqShieldWorkMode.CREATE_AND_UPDATE_CACHE, ) diff --git a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/CoroutineExtension.kt b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/CoroutineExtension.kt index 5ace0a5..6513ef1 100644 --- a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/CoroutineExtension.kt +++ b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/CoroutineExtension.kt @@ -18,10 +18,6 @@ package com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.aspect -import kotlinx.coroutines.CoroutineDispatcher -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.ExperimentalCoroutinesApi -import kotlinx.coroutines.withContext import org.aspectj.lang.ProceedingJoinPoint import kotlin.coroutines.Continuation import kotlin.coroutines.intrinsics.startCoroutineUninterceptedOrReturn @@ -43,44 +39,3 @@ suspend fun ProceedingJoinPoint.proceedCoroutine(args: Array = this.corout fun ProceedingJoinPoint.runCoroutine(block: suspend () -> Any?): Any? = block.startCoroutineUninterceptedOrReturn(this.coroutineContinuation) - -/** - * Bounded dispatcher for non-suspend join point execution to prevent IO dispatcher exhaustion. - * Limits concurrent blocking calls to prevent thread pool saturation under heavy load. - * - * Parallelism can be configured via system property: - * - `reqshield.blocking.parallelism`: explicit parallelism value (1-1024) - * - Default: availableProcessors * 2, clamped to [4, 256] - * - * Examples: - * - `-Dreqshield.blocking.parallelism=64` for high-throughput environments - * - `-Dreqshield.blocking.parallelism=8` for resource-constrained environments - */ -@OptIn(ExperimentalCoroutinesApi::class) -private val boundedBlockingDispatcher: CoroutineDispatcher by lazy { - val defaultParallelism = - (Runtime.getRuntime().availableProcessors() * 2) - .coerceIn(4, 256) // Min 4, max 256 - - val parallelism = - System.getProperty("reqshield.blocking.parallelism") - ?.toIntOrNull() - ?.coerceIn(1, 1024) // Configured value also bounded - ?: defaultParallelism - - Dispatchers.IO.limitedParallelism(parallelism) -} - -/** - * Proceed supporting both suspend and non-suspend join points. - * If the last argument is a Continuation, treat as suspend; otherwise proceed normally. - * Uses a bounded dispatcher for non-suspend calls to prevent IO thread pool exhaustion. - */ -suspend fun ProceedingJoinPoint.proceedSmart(): Any? = - if (this.args.isNotEmpty() && this.args.last() is Continuation<*>) { - this.proceedCoroutine() - } else { - // Use bounded dispatcher to prevent IO thread pool exhaustion - // when many synchronous methods are proxied concurrently - withContext(boundedBlockingDispatcher) { this@proceedSmart.proceed() } - } diff --git a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspect.kt b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspect.kt index 8cf2302..b95909d 100644 --- a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspect.kt +++ b/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.spring.webflux.kotlin.coroutine.cache.GlobalLockSupport +import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.reactor.awaitSingleOrNull import org.aspectj.lang.ProceedingJoinPoint import org.aspectj.lang.annotation.Around @@ -28,6 +30,7 @@ import org.aspectj.lang.annotation.Aspect import org.aspectj.lang.reflect.MethodSignature import org.springframework.beans.factory.BeanFactory import org.springframework.beans.factory.BeanFactoryAware +import org.springframework.beans.factory.annotation.Qualifier import org.springframework.cache.interceptor.KeyGenerator import org.springframework.cache.interceptor.SimpleKeyGenerator import org.springframework.context.expression.MethodBasedEvaluationContext @@ -37,7 +40,6 @@ import org.springframework.core.annotation.AnnotationUtils import org.springframework.expression.EvaluationContext import org.springframework.expression.Expression import org.springframework.expression.spel.standard.SpelExpressionParser -import org.springframework.stereotype.Component import org.springframework.util.StringUtils import org.springframework.util.function.SingletonSupplier import reactor.core.publisher.Mono @@ -46,21 +48,32 @@ import java.util.concurrent.ConcurrentHashMap import kotlin.coroutines.Continuation @Aspect -@Component open class ReqShieldAspect( private val asyncCache: AsyncCache, + @Qualifier("reqShieldCoroutineScope") private val scope: CoroutineScope, ) : BeanFactoryAware { private lateinit var beanFactory: BeanFactory private val springVersion = SpringVersion.getVersion() private val spelParser = SpelExpressionParser() + private val parameterNameDiscoverer = DefaultParameterNameDiscoverer() private var defaultKeyGenerator = SingletonSupplier.of { SimpleKeyGenerator() } + /** Global locking is opt-in: only `isLocalLock = false` needs the cache to support it. */ + private val lockSupport = asyncCache as? GlobalLockSupport + private val keyGeneratorMap = ConcurrentHashMap() - internal val reqShieldMap = ConcurrentHashMap>() + + /** Parsing a SpEL expression is expensive, so each annotation key is parsed only once. */ + private val expressionMap = ConcurrentHashMap() + + /** One ReqShield per annotated method: the cache key varies per call, the configuration does not. */ + internal val reqShieldMap = ConcurrentHashMap>() @Around("@annotation(com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.annotation.ReqShieldCacheable)") - fun aroundReqShieldCacheable(joinPoint: ProceedingJoinPoint): Any? = - joinPoint.runCoroutine { + fun aroundReqShieldCacheable(joinPoint: ProceedingJoinPoint): Any? { + requireSuspendTarget(joinPoint, "ReqShieldCacheable") + + return joinPoint.runCoroutine { val annotation = getCacheableAnnotation(joinPoint) val reqShield = getOrCreateReqShield(joinPoint) val cacheKey = getCacheableCacheKey(joinPoint) @@ -69,7 +82,8 @@ open class ReqShieldAspect( .getAndSetReqShieldData( cacheKey, { - joinPoint.proceedSmart().let { rtn -> + joinPoint.proceedCoroutine().let { rtn -> + // A suspend function may still declare Mono as its return type. if (rtn is Mono<*>) { rtn.awaitSingleOrNull()?.let { it as T } } else { @@ -80,14 +94,29 @@ open class ReqShieldAspect( annotation.timeToLiveMillis, ).value } + } @Around("@annotation(com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.annotation.ReqShieldCacheEvict)") - fun aroundReqShieldCacheEvict(joinPoint: ProceedingJoinPoint): Any? = - joinPoint.runCoroutine { + fun aroundReqShieldCacheEvict(joinPoint: ProceedingJoinPoint): Any? { + requireSuspendTarget(joinPoint, "ReqShieldCacheEvict") + + return joinPoint.runCoroutine { val cacheKey = getCacheEvictCacheKey(joinPoint) + // Evict after the method succeeded, like Spring's @CacheEvict default: a failing method + // leaves the cache untouched, and a failing eviction is reported to the caller. + val result = + joinPoint.proceedCoroutine().let { rtn -> + // Spring 6.1+ adapts a suspend target to a cold Mono inside AOP. + if (rtn is Mono<*>) { + rtn.awaitSingleOrNull() + } else { + rtn + } + } asyncCache.evict(cacheKey) - joinPoint.proceedSmart() + result } + } internal open fun getTargetMethod(joinPoint: ProceedingJoinPoint): Method = (joinPoint.signature as MethodSignature).method @@ -99,18 +128,29 @@ open class ReqShieldAspect( AnnotationUtils.getAnnotation(getTargetMethod(joinPoint), ReqShieldCacheEvict::class.java) ?: throw IllegalArgumentException("ReqShieldCacheEvict annotation is required") - internal fun getCacheableCacheKey(joinPoint: ProceedingJoinPoint): String { - val annotation = getCacheableAnnotation(joinPoint) - validateCacheKey(annotation.key, annotation.keyGenerator) + internal fun getCacheableCacheKey(joinPoint: ProceedingJoinPoint): String = + getCacheableAnnotation(joinPoint).let { + namespacedCacheKey(it.cacheName, it.key, it.keyGenerator, joinPoint) + } - return getCacheKeyOrDefault(annotation.key, annotation.keyGenerator, joinPoint) - } + internal fun getCacheEvictCacheKey(joinPoint: ProceedingJoinPoint): String = + getCacheEvictAnnotation(joinPoint).let { + namespacedCacheKey(it.cacheName, it.key, it.keyGenerator, joinPoint) + } - internal fun getCacheEvictCacheKey(joinPoint: ProceedingJoinPoint): String { - val annotation = getCacheEvictAnnotation(joinPoint) - validateCacheKey(annotation.key, annotation.keyGenerator) + /** + * The cache name is part of the key so that two annotations resolving the same key in different + * caches cannot share a cache entry - nor, since ReqShield locks on this key, a lock. + */ + private fun namespacedCacheKey( + cacheName: String, + annotationCacheKey: String, + annotationCacheKeyGenerator: String, + joinPoint: ProceedingJoinPoint, + ): String { + validateCacheKey(annotationCacheKey, annotationCacheKeyGenerator) - return getCacheKeyOrDefault(annotation.key, annotation.keyGenerator, joinPoint) + return "$cacheName::${getCacheKeyOrDefault(annotationCacheKey, annotationCacheKeyGenerator, joinPoint)}" } private fun getCacheKeyOrDefault( @@ -127,11 +167,11 @@ open class ReqShieldAspect( } val context: EvaluationContext = - MethodBasedEvaluationContext(joinPoint.target, method, args, DefaultParameterNameDiscoverer()) + MethodBasedEvaluationContext(joinPoint.target, method, args, parameterNameDiscoverer) val key = if (StringUtils.hasText(annotationCacheKey)) { - val expression: Expression = spelParser.parseExpression(annotationCacheKey) + val expression = expressionMap.computeIfAbsent(annotationCacheKey) { spelParser.parseExpression(it) } expression.getValue(context, String::class.java) } else { val keyGenerator = getOrCreateKeyGenerator(annotationCacheKeyGenerator) @@ -139,7 +179,7 @@ open class ReqShieldAspect( } require(!key.isNullOrBlank()) { - "Null/blank key for @ReqShieldCacheable method=${method.declaringClass.name}.${method.name} " + + "Null/blank key for method=${method.declaringClass.name}.${method.name} " + "args=${args.joinToString(prefix = "[", postfix = "]") { it?.let { arg -> @@ -151,13 +191,34 @@ open class ReqShieldAspect( return key } + /** + * [runCoroutine] hands the block to the join point's own continuation, so a non-suspend target + * would fail later with an opaque cast error. Reject it here instead. + */ + private fun requireSuspendTarget( + joinPoint: ProceedingJoinPoint, + annotationName: String, + ) { + require(joinPoint.args.lastOrNull() is Continuation<*>) { + val method = getTargetMethod(joinPoint) + "@$annotationName in the coroutine module requires a suspend function: " + + "${method.declaringClass.name}.${method.name}" + } + } + private fun getOrCreateReqShield(joinPoint: ProceedingJoinPoint): ReqShield = - reqShieldMap.computeIfAbsent(generateReqShieldKey(joinPoint)) { + reqShieldMap.computeIfAbsent(getTargetMethod(joinPoint)) { createReqShield(joinPoint) } private fun createReqShield(joinPoint: ProceedingJoinPoint): ReqShield { val annotation = getCacheableAnnotation(joinPoint) + val method = getTargetMethod(joinPoint) + + require(annotation.isLocalLock || lockSupport != null) { + "isLocalLock = false on ${method.declaringClass.name}.${method.name} " + + "requires the AsyncCache bean to implement GlobalLockSupport" + } val reqShieldConfiguration = ReqShieldConfiguration( @@ -167,17 +228,14 @@ open class ReqShieldAspect( getCacheFunction = { key -> asyncCache.get(key) }, - globalLockFunction = { key, timeToLiveMillis -> - asyncCache.globalLock(key, timeToLiveMillis) - }, - globalUnLockFunction = { key -> - asyncCache.globalUnLock(key) - }, + globalLockFunction = lockSupport?.let { support -> { key, token, ttl -> support.globalLock(key, token, ttl) } }, + globalUnLockFunction = lockSupport?.let { support -> { key, token -> support.globalUnLock(key, token) } }, isLocalLock = annotation.isLocalLock, lockTimeoutMillis = annotation.lockTimeoutMillis, decisionForUpdate = annotation.decisionForUpdate, maxAttemptGetCache = annotation.maxAttemptGetCache, reqShieldWorkMode = annotation.reqShieldWorkMode, + scope = scope, ) return ReqShield(reqShieldConfiguration) @@ -213,12 +271,6 @@ open class ReqShieldAspect( return major > 6 || (major == 6 && minor >= 1) } - private fun generateReqShieldKey(joinPoint: ProceedingJoinPoint): String { - val method = getTargetMethod(joinPoint) - return "${method.declaringClass.name}.${method.name}-" + - "${getCacheableAnnotation(joinPoint).cacheName}-${getCacheableCacheKey(joinPoint)}" - } - override fun setBeanFactory(beanFactory: BeanFactory) { this.beanFactory = beanFactory } diff --git a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/cache/AsyncCache.kt b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/cache/AsyncCache.kt index 9f8fabc..f561e77 100644 --- a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/cache/AsyncCache.kt +++ b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/cache/AsyncCache.kt @@ -28,32 +28,4 @@ interface AsyncCache { ): Boolean suspend fun evict(key: String): Boolean - - /** - * Attempt a global lock on a specific key. - * - * @param key The key to get the lock. - * @param timeToLiveMillis The validity of the lock in milliseconds. - * @return Whether the lock was successfully obtained. Returns `true` by default. - * - * This method provides a default implementation, but if you need a locking mechanism - * You must implement and use your own locking logic. The default implementation is true, and if the value of ReqShieldConfiguration > isLocalLock is false, you will use the function you implemented. - * actual production environments should override this method appropriately to manage locks. - */ - suspend fun globalLock( - key: String, - timeToLiveMillis: Long, - ): Boolean = true - - /** - * Releases the global lock on a specific key. - * - * @param key The key you want to unlock. - * @return Whether the lock was successfully obtained. Returns `true` by default. - * - * This method provides a default implementation, but if you need a locking mechanism - * You must implement and use your own unlocking logic. The default implementation is true, and if the value of ReqShieldConfiguration > isLocalLock is false, you will use the function you implemented. - * actual production environments should override this method appropriately to manage locks. - */ - suspend fun globalUnLock(key: String): Boolean = true } diff --git a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/cache/GlobalLockSupport.kt b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/cache/GlobalLockSupport.kt new file mode 100644 index 0000000..97b9b59 --- /dev/null +++ b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/cache/GlobalLockSupport.kt @@ -0,0 +1,33 @@ +/* + * 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.spring.webflux.kotlin.coroutine.cache + +/** Implement alongside [AsyncCache] to use `@ReqShieldCacheable(isLocalLock = false)`. */ +interface GlobalLockSupport { + /** Acquire lockKey for the caller identified by token (Redis: `SET lockKey token NX PX timeToLiveMillis`). */ + suspend fun globalLock( + lockKey: String, + token: String, + timeToLiveMillis: Long, + ): Boolean + + /** Release only if the stored value equals token (Redis: compare-and-delete Lua script). */ + suspend fun globalUnLock( + lockKey: String, + token: String, + ): Boolean +} diff --git a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/config/LibAutoConfiguration.kt b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/config/LibAutoConfiguration.kt index 14555fa..7d7a78b 100644 --- a/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/config/LibAutoConfiguration.kt +++ b/core-spring-webflux-kotlin-coroutine/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/config/LibAutoConfiguration.kt @@ -17,11 +17,53 @@ package com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.config import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.aspect.ReqShieldAspect +import kotlinx.coroutines.CoroutineExceptionHandler +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.SupervisorJob +import kotlinx.coroutines.cancel +import org.slf4j.LoggerFactory +import org.springframework.beans.factory.DisposableBean +import org.springframework.context.annotation.Bean import org.springframework.context.annotation.Configuration import org.springframework.context.annotation.EnableAspectJAutoProxy import org.springframework.context.annotation.Import +import kotlin.coroutines.CoroutineContext @Configuration @EnableAspectJAutoProxy(proxyTargetClass = true) @Import(ReqShieldAspect::class) -open class LibAutoConfiguration +open class LibAutoConfiguration { + /** + * Scope shared by every [com.linecorp.cse.reqshield.kotlin.coroutine.ReqShield] the aspect + * creates, used for the fire-and-forget cache writes. + * + * SupervisorJob keeps one failed write from cancelling the others, and the exception handler is + * the last-resort backstop for anything the write path did not already log. + */ + @Bean + open fun reqShieldCoroutineScope(): CoroutineScope = + ReqShieldCoroutineScope( + SupervisorJob() + Dispatchers.IO + + CoroutineExceptionHandler { _, e -> + log.error("[Req-Shield] background task failed", e) + }, + ) + + /** + * [CoroutineScope] has no `cancel` member, so the bean carries its own shutdown hook: closing + * the application context cancels the scope and with it every pending cache write. + */ + private class ReqShieldCoroutineScope( + override val coroutineContext: CoroutineContext, + ) : CoroutineScope, + DisposableBean { + override fun destroy() { + coroutineContext.cancel() + } + } + + companion object { + private val log = LoggerFactory.getLogger(LibAutoConfiguration::class.java) + } +} diff --git a/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/InMemoryAsyncCache.kt b/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/InMemoryAsyncCache.kt index 926c3dc..2abc8a6 100644 --- a/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/InMemoryAsyncCache.kt +++ b/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/InMemoryAsyncCache.kt @@ -17,15 +17,19 @@ package com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.aspect import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.cache.AsyncCache +import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.cache.GlobalLockSupport import com.linecorp.cse.reqshield.support.model.ReqShieldData import java.util.concurrent.ConcurrentHashMap -import java.util.concurrent.Semaphore -class InMemoryAsyncCache : AsyncCache { +class InMemoryAsyncCache : + AsyncCache, + GlobalLockSupport { private data class Entry(val data: ReqShieldData, val expiresAt: Long) + private data class Lock(val token: String, val expiresAt: Long) + private val store = ConcurrentHashMap>() - private val locks = ConcurrentHashMap() + private val locks = ConcurrentHashMap() override suspend fun get(key: String): ReqShieldData? { val now = System.currentTimeMillis() @@ -44,13 +48,40 @@ class InMemoryAsyncCache : AsyncCache { override suspend fun evict(key: String): Boolean = store.remove(key) != null + /** In-memory equivalent of `SET lockKey token NX PX ttl`: the stored token identifies the owner. */ override suspend fun globalLock( - key: String, + lockKey: String, + token: String, timeToLiveMillis: Long, - ): Boolean = locks.computeIfAbsent(key) { Semaphore(1) }.tryAcquire() + ): Boolean { + val now = System.currentTimeMillis() + val owner = + locks.compute(lockKey) { _, current -> + if (current == null || now > current.expiresAt) { + Lock(token, now + timeToLiveMillis) + } else { + current + } + } - override suspend fun globalUnLock(key: String): Boolean { - locks[key]?.release() - return true + return owner?.token == token + } + + /** In-memory equivalent of the compare-and-delete script: a stale owner cannot release the lock. */ + override suspend fun globalUnLock( + lockKey: String, + token: String, + ): Boolean { + var released = false + locks.compute(lockKey) { _, current -> + if (current?.token == token) { + released = true + null + } else { + current + } + } + + return released } } diff --git a/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectIntegrationTest.kt b/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectIntegrationTest.kt index 8a253d5..af14308 100644 --- a/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectIntegrationTest.kt +++ b/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectIntegrationTest.kt @@ -20,22 +20,29 @@ 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.spring.webflux.kotlin.coroutine.config.LibAutoConfiguration +import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.Job import kotlinx.coroutines.async import kotlinx.coroutines.awaitAll import kotlinx.coroutines.delay import kotlinx.coroutines.runBlocking import kotlinx.coroutines.withTimeoutOrNull import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertNotNull +import org.junit.jupiter.api.Assertions.assertNull import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.Test import org.junit.jupiter.api.extension.ExtendWith import org.springframework.beans.factory.annotation.Autowired +import org.springframework.context.annotation.AnnotationConfigApplicationContext import org.springframework.context.annotation.Bean import org.springframework.context.annotation.Configuration import org.springframework.test.context.ContextConfiguration import org.springframework.test.context.junit.jupiter.SpringExtension import java.util.concurrent.atomic.AtomicInteger +import kotlin.test.assertFailsWith @ExtendWith(SpringExtension::class) @ContextConfiguration(classes = [LibAutoConfiguration::class, ReqShieldAspectIntegrationTest.TestConfig::class]) @@ -46,12 +53,18 @@ class ReqShieldAspectIntegrationTest { @Autowired private lateinit var asyncCache: AsyncCache + @Autowired + private lateinit var reqShieldCoroutineScope: CoroutineScope + + /** The aspect namespaces every key with the cache name of the annotation. */ + private fun cacheKeyOf(key: String) = "$CACHE_NAME::$key" + private suspend fun awaitCachePut( key: String, timeoutMillis: Long = 1_000, ): Boolean = withTimeoutOrNull(timeoutMillis) { - while (asyncCache.get(key) == null) { + while (asyncCache.get(cacheKeyOf(key)) == null) { delay(5) } true @@ -84,6 +97,56 @@ class ReqShieldAspectIntegrationTest { assertTrue(v1 != v2) } + @Test + fun shouldEvictOnlyAfterTheAnnotatedMethodSucceeds() = + runBlocking { + val key = "evict-success-${System.nanoTime()}" + service.get(key) + assertTrue(awaitCachePut(key), "Timed out waiting for cache put for key=$key") + + assertTrue(service.evict(key)) + + assertNull(asyncCache.get(cacheKeyOf(key))) + } + + @Test + fun shouldKeepTheCacheWhenTheAnnotatedMethodThrows() = + runBlocking { + val key = "evict-failure-${System.nanoTime()}" + service.get(key) + assertTrue(awaitCachePut(key), "Timed out waiting for cache put for key=$key") + + assertFailsWith { service.evictFailing(key) } + + assertNotNull(asyncCache.get(cacheKeyOf(key))) + } + + @Test + fun shouldExposeTheCacheKeyUnderTheCacheNameNamespace() = + runBlocking { + val key = "namespace-${System.nanoTime()}" + service.get(key) + assertTrue(awaitCachePut(key), "Timed out waiting for cache put for key=$key") + + // The bare key must never be used: only the namespaced one carries the entry. + assertNull(asyncCache.get(key)) + assertNotNull(asyncCache.get(cacheKeyOf(key))) + } + + @Test + fun shouldWireTheCoroutineScopeBeanAndCancelItOnShutdown() { + // The autowired bean proves LibAutoConfiguration alone provides the scope (no component scan). + assertFalse(reqShieldCoroutineScope.coroutineContext[Job]!!.isCancelled) + + val context = AnnotationConfigApplicationContext(LibAutoConfiguration::class.java, TestConfig::class.java) + val scope = context.getBean("reqShieldCoroutineScope", CoroutineScope::class.java) + assertFalse(scope.coroutineContext[Job]!!.isCancelled) + + context.close() + + assertTrue(scope.coroutineContext[Job]?.isCancelled == true) + } + @Configuration open class TestConfig { @Bean @@ -96,10 +159,17 @@ class ReqShieldAspectIntegrationTest { open class TestService { val counter = AtomicInteger(0) - @ReqShieldCacheable(cacheName = "it", key = "#key", timeToLiveMillis = 10_000) + @ReqShieldCacheable(cacheName = CACHE_NAME, key = "#key", timeToLiveMillis = 10_000) open suspend fun get(key: String): String = "value-" + counter.incrementAndGet() - @ReqShieldCacheEvict(cacheName = "it", key = "#key") + @ReqShieldCacheEvict(cacheName = CACHE_NAME, key = "#key") open suspend fun evict(key: String): Boolean = true + + @ReqShieldCacheEvict(cacheName = CACHE_NAME, key = "#key") + open suspend fun evictFailing(key: String): Boolean = throw IllegalStateException("eviction target failed: $key") + } + + companion object { + const val CACHE_NAME = "it" } } diff --git a/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectRedisIntegrationTest.kt b/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectRedisIntegrationTest.kt index 9518e40..2c9dfda 100644 --- a/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectRedisIntegrationTest.kt +++ b/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectRedisIntegrationTest.kt @@ -19,10 +19,13 @@ package com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.aspect 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.spring.webflux.kotlin.coroutine.cache.GlobalLockSupport import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.config.LibAutoConfiguration import com.linecorp.cse.reqshield.support.model.ReqShieldData import com.linecorp.cse.reqshield.support.redis.AbstractRedisTest import io.lettuce.core.RedisClient +import io.lettuce.core.ScriptOutputType +import io.lettuce.core.SetArgs import io.lettuce.core.api.StatefulRedisConnection import io.lettuce.core.api.sync.RedisCommands import kotlinx.coroutines.Dispatchers @@ -31,6 +34,8 @@ import kotlinx.coroutines.awaitAll import kotlinx.coroutines.delay import kotlinx.coroutines.runBlocking import kotlinx.coroutines.withTimeoutOrNull +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertNull import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test @@ -57,12 +62,15 @@ class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { service.resetCounter() } + /** The aspect namespaces every key with the cache name of the annotation. */ + private fun cacheKeyOf(key: String) = "$CACHE_NAME::$key" + private suspend fun awaitCachePut( key: String, timeoutMillis: Long = 2_000, ): Boolean = withTimeoutOrNull(timeoutMillis) { - while (asyncCache.get(key) == null) { + while (asyncCache.get(cacheKeyOf(key)) == null) { delay(10) } true @@ -88,6 +96,21 @@ class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { ) } + @Test + fun shouldCollapseDuplicateRequestsWithRedisGlobalLock() = + runBlocking { + val key = "dup-global-${System.nanoTime()}" + val attempts = 20 + val results = (1..attempts).map { async(Dispatchers.IO) { service.getWithGlobalLock(key) } }.awaitAll() + + assertTrue( + service.getRequestCount() == 1, + "Callable should be invoked only once. actual=${service.getRequestCount()}", + ) + assertTrue(results.all { it != null }, "Expected all results to be valid. results=$results") + assertTrue(awaitCachePut(key), "Timed out waiting for cache put for key=$key") + } + @Test fun shouldEvictAndRecomputeWithRedis() = runBlocking { @@ -99,6 +122,23 @@ class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { val v2 = service.get(key) assertTrue(evicted, "Eviction should return true") assertTrue(v1 != v2, "Values should differ after eviction: v1=$v1, v2=$v2") + assertNull(asyncCache.get(key), "The bare key must never hold an entry") + } + + @Test + fun shouldReleaseTheGlobalLockOnlyForTheOwningToken() = + runBlocking { + val lockSupport = asyncCache as GlobalLockSupport + val lockKey = "lock-token-${System.nanoTime()}" + + assertTrue(lockSupport.globalLock(lockKey, "owner", 10_000)) + // SET NX: a second caller cannot take a held lock. + assertFalse(lockSupport.globalLock(lockKey, "intruder", 10_000)) + // Compare-and-delete: a non-owner cannot release it either. + assertFalse(lockSupport.globalUnLock(lockKey, "intruder")) + assertTrue(lockSupport.globalUnLock(lockKey, "owner")) + // Released, so the next caller can take it. + assertTrue(lockSupport.globalLock(lockKey, "intruder", 10_000)) } @Configuration @@ -121,7 +161,9 @@ class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { // Ensure clean DB state for tests running in CI runCatching { sync.flushdb() } - return object : AsyncCache { + return object : + AsyncCache, + GlobalLockSupport { override suspend fun get(key: String): ReqShieldData? = sync.get(key)?.let { ReqShieldData(value = it, timeToLiveMillis = 10_000) } @@ -137,11 +179,15 @@ class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { override suspend fun evict(key: String): Boolean = sync.del(key) > 0 override suspend fun globalLock( - key: String, + lockKey: String, + token: String, timeToLiveMillis: Long, - ): Boolean = sync.setnx("lock:$key", "1").also { if (it) sync.pexpire("lock:$key", timeToLiveMillis) } + ): Boolean = sync.set(lockKey, token, SetArgs.Builder.nx().px(timeToLiveMillis)) == "OK" - override suspend fun globalUnLock(key: String): Boolean = sync.del("lock:$key") >= 0 + override suspend fun globalUnLock( + lockKey: String, + token: String, + ): Boolean = sync.eval(UNLOCK_SCRIPT, ScriptOutputType.INTEGER, arrayOf(lockKey), token) == 1L } } @@ -159,7 +205,7 @@ class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { open fun getRequestCount(): Int = counter.get() @ReqShieldCacheable( - cacheName = "it", + cacheName = CACHE_NAME, key = "#key", timeToLiveMillis = 10_000, // CI environments can be slow; give enough time for async cache put to be observed by waiters. @@ -168,7 +214,25 @@ class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { ) open suspend fun get(key: String): String = "value-" + counter.incrementAndGet() - @ReqShieldCacheEvict(cacheName = "it", key = "#key") + @ReqShieldCacheable( + cacheName = CACHE_NAME, + key = "#key", + isLocalLock = false, + timeToLiveMillis = 10_000, + maxAttemptGetCache = 200, + lockTimeoutMillis = 10_000, + ) + open suspend fun getWithGlobalLock(key: String): String = "value-" + counter.incrementAndGet() + + @ReqShieldCacheEvict(cacheName = CACHE_NAME, key = "#key") open suspend fun evict(key: String): Boolean = true } + + companion object { + const val CACHE_NAME = "it" + + /** Release the lock only when the caller still owns it, as one atomic step. */ + const val UNLOCK_SCRIPT = + "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end" + } } diff --git a/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectTest.kt b/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectTest.kt index 223007d..a36f58f 100644 --- a/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectTest.kt +++ b/core-spring-webflux-kotlin-coroutine/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/aspect/ReqShieldAspectTest.kt @@ -19,19 +19,29 @@ package com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.aspect 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.spring.webflux.kotlin.coroutine.cache.GlobalLockSupport import com.linecorp.cse.reqshield.support.BaseReqShieldModuleSupportTest +import com.linecorp.cse.reqshield.support.constant.ConfigValues.DEFAULT_LOCK_TIMEOUT_MILLIS +import com.linecorp.cse.reqshield.support.constant.ConfigValues.LOCK_KEY_PREFIX import com.linecorp.cse.reqshield.support.model.Product import com.linecorp.cse.reqshield.support.model.ReqShieldData import io.mockk.coEvery +import io.mockk.coVerify import io.mockk.every import io.mockk.mockk +import io.mockk.slot import io.mockk.spyk +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.ExperimentalCoroutinesApi +import kotlinx.coroutines.Job +import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.async import kotlinx.coroutines.awaitAll import kotlinx.coroutines.test.runTest import org.aspectj.lang.ProceedingJoinPoint import org.junit.jupiter.api.Assertions.assertNotNull +import org.junit.jupiter.api.Assertions.assertNull import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test @@ -45,6 +55,7 @@ import kotlin.coroutines.EmptyCoroutineContext import kotlin.reflect.full.functions import kotlin.reflect.jvm.javaMethod import kotlin.test.assertEquals +import kotlin.test.assertFailsWith private val log = LoggerFactory.getLogger(ReqShieldAspectTest::class.java) @@ -52,7 +63,8 @@ private val log = LoggerFactory.getLogger(ReqShieldAspectTest::class.java) class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { private val asyncCache: AsyncCache = InMemoryAsyncCache() private val joinPoint: ProceedingJoinPoint = mockk() - private val reqShieldAspect: ReqShieldAspect = spyk(ReqShieldAspect(asyncCache)) + private val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default) + private val reqShieldAspect: ReqShieldAspect = spyk(ReqShieldAspect(asyncCache, scope)) private val targetObject = spyk(TestBean()) private val argument = mapOf("x" to "paramX", "y" to "paramY") private val mockContinuation = mockk>() @@ -61,6 +73,7 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { private val cacheKeyGenerator = "customGenerator" private val spelEvaluatedKey = "paramXparamY" private val keyGeneratorKey = "KeyGeneratedByGenerator" + private val namespacedSpelKey = "$cacheName::$spelEvaluatedKey" private val beanFactory = mockk() @@ -75,144 +88,245 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { reqShieldAspect.setBeanFactory(beanFactory) } + /** The single-parameter overload of [name] on [TestBean]; suspend and plain ones both report 2. */ + private fun methodOf(name: String): Method = + TestBean::class + .functions + .find { it.name == name && it.parameters.size == 2 } + ?.javaMethod!! + + /** Cache writes are fire-and-forget, so tests that assert on them must wait for the scope. */ + private suspend fun CoroutineScope.awaitBackgroundWrites() { + coroutineContext[Job]?.children?.toList()?.forEach { it.join() } + } + @Test override fun verifyReqShieldCacheCreation() = runTest { - // Mock the cache data using mockk val reqShieldData = ReqShieldData(methodReturn, 1000) - asyncCache.put(spelEvaluatedKey, reqShieldData, 1000) - coEvery { joinPoint.proceed() } coAnswers { targetObject.cacheableWithCustomKey(argument) } - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - TestBean::class - .functions - .find { - it.name == TestBean::cacheableWithCustomKey.name && it.parameters.size == 2 - }?.javaMethod!! - - // Test the aroundTargetCacheable method + asyncCache.put(namespacedSpelKey, reqShieldData, 1000) + coEvery { joinPoint.proceed(any>()) } coAnswers { targetObject.cacheableWithCustomKey(argument) } + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithCustomKey.name) + val result = reqShieldAspect.aroundReqShieldCacheable(joinPoint) - assertEquals(result, reqShieldData.value) - assertTrue(reqShieldAspect.reqShieldMap.size == 1) - val method = reqShieldAspect.getTargetMethod(joinPoint) - val expectedKey = "${method.declaringClass.name}.${method.name}-$cacheName-$spelEvaluatedKey" - assertNotNull(reqShieldAspect.reqShieldMap[expectedKey]) + assertEquals(reqShieldData.value, result) + assertEquals(1, reqShieldAspect.reqShieldMap.size) + assertNotNull(reqShieldAspect.reqShieldMap[reqShieldAspect.getTargetMethod(joinPoint)]) } @Test override fun reqShieldObjectShouldBeCreatedOnce() = runTest { - // Mock the cache data using mockk val reqShieldData = ReqShieldData(methodReturn, 1000) - asyncCache.put(spelEvaluatedKey, reqShieldData, 1000) - coEvery { joinPoint.proceed() } coAnswers { targetObject.cacheableWithCustomKey(argument) } - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - TestBean::class - .functions - .find { - it.name == TestBean::cacheableWithCustomKey.name && it.parameters.size == 2 - }?.javaMethod!! - - val jobs = - List(20) { - async { - reqShieldAspect.aroundReqShieldCacheable(joinPoint) - } - } + asyncCache.put(namespacedSpelKey, reqShieldData, 1000) + coEvery { joinPoint.proceed(any>()) } coAnswers { targetObject.cacheableWithCustomKey(argument) } + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithCustomKey.name) - jobs.awaitAll() + List(20) { async { reqShieldAspect.aroundReqShieldCacheable(joinPoint) } }.awaitAll() - assertTrue(reqShieldAspect.reqShieldMap.size == 1) - val method = reqShieldAspect.getTargetMethod(joinPoint) - val expectedKey = "${method.declaringClass.name}.${method.name}-$cacheName-$spelEvaluatedKey" - assertNotNull(reqShieldAspect.reqShieldMap[expectedKey]) + assertEquals(1, reqShieldAspect.reqShieldMap.size) + assertNotNull(reqShieldAspect.reqShieldMap[reqShieldAspect.getTargetMethod(joinPoint)]) } @Test override fun verifyReqShieldCacheEviction() = runTest { val reqShieldData = ReqShieldData(methodReturn, 1000) - // Use SpEL-based key to align with eviction method's key - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - TestBean::class - .functions - .find { - it.name == TestBean::cacheableWithCustomKey.name && it.parameters.size == 2 - }?.javaMethod!! + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithCustomKey.name) val generatedKey = reqShieldAspect.getCacheableCacheKey(joinPoint) asyncCache.put(generatedKey, reqShieldData, 1000) - coEvery { joinPoint.proceed() } coAnswers { targetObject.cacheableWithCustomKey(argument) } + coEvery { joinPoint.proceed(any>()) } coAnswers { targetObject.cacheableWithCustomKey(argument) } - // Test the aroundTargetCacheable method using the same SpEL key val result = reqShieldAspect.aroundReqShieldCacheable(joinPoint) assertEquals(reqShieldData.value, result) - // Validate cache eviction using the eviction method (same SpEL key) - // real eviction call - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - TestBean::class - .functions - .find { - it.name == TestBean::evict.name && it.parameters.size == 2 - }?.javaMethod!! - // Mock proceed for eviction - the aspect proceeds to the original method after evicting cache - // proceedSmart() calls proceed(args) with continuation, so we need to mock that as well + // The evict advice resolves its own key from @ReqShieldCacheEvict; it must hit the same entry. + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::evict.name) coEvery { joinPoint.proceed(any>()) } coAnswers { targetObject.evict(argument) } - val removeProductMono = reqShieldAspect.aroundReqShieldCacheEvict(joinPoint) + val evicted = reqShieldAspect.aroundReqShieldCacheEvict(joinPoint) - assertTrue(removeProductMono as Boolean) + assertTrue(evicted as Boolean) + assertNull(asyncCache.get(generatedKey)) + } + + @Test + fun cacheIsEvictedOnlyAfterTheMethodSucceeds() = + runTest { + val reqShieldData = ReqShieldData(methodReturn, 1000) + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::evictFailing.name) + val evictKey = reqShieldAspect.getCacheEvictCacheKey(joinPoint) + asyncCache.put(evictKey, reqShieldData, 1000) + coEvery { joinPoint.proceed(any>()) } coAnswers { targetObject.evictFailing(argument) } + + assertFailsWith { reqShieldAspect.aroundReqShieldCacheEvict(joinPoint) } + + assertNotNull(asyncCache.get(evictKey)) } @Test override fun verifyCacheKeyGenerationWithSpEL() = runTest { - coEvery { reqShieldAspect.getTargetMethod(joinPoint) } returns - TestBean::class - .functions - .find { - it.name == TestBean::cacheableWithCustomKey.name && it.parameters.size == 2 - }?.javaMethod!! - - assertEquals(spelEvaluatedKey, reqShieldAspect.getCacheableCacheKey(joinPoint)) + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithCustomKey.name) + + assertEquals(namespacedSpelKey, reqShieldAspect.getCacheableCacheKey(joinPoint)) } @Test override fun verifyCacheKeyGenerationWithKeyGenerator() = runTest { - coEvery { beanFactory.getBean(cacheKeyGenerator, KeyGenerator::class.java) } returns - CustomGenerator() - coEvery { reqShieldAspect.getTargetMethod(joinPoint) } returns - TestBean::class - .functions - .find { - it.name == TestBean::cacheableWithKeyGenerator.name && it.parameters.size == 2 - }?.javaMethod!! - - assertEquals(keyGeneratorKey, reqShieldAspect.getCacheableCacheKey(joinPoint)) + every { beanFactory.getBean(cacheKeyGenerator, KeyGenerator::class.java) } returns CustomGenerator() + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithKeyGenerator.name) + + assertEquals("$cacheName::$keyGeneratorKey", reqShieldAspect.getCacheableCacheKey(joinPoint)) } @Test override fun verifyCacheKeyGenerationWithDefaultGenerator() = runTest { - coEvery { reqShieldAspect.getTargetMethod(joinPoint) } returns - TestBean::class - .functions - .find { - it.name == TestBean::cacheableWithDefaultKeyGenerator.name && it.parameters.size == 2 - }?.javaMethod!! + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithDefaultKeyGenerator.name) assertEquals( - SimpleKeyGenerator.generateKey(arrayOf(argument)).toString(), + "$cacheName::${SimpleKeyGenerator.generateKey(arrayOf(argument))}", reqShieldAspect.getCacheableCacheKey(joinPoint), ) } + @Test + fun cacheKeysOfTwoCacheNamesNeverCollide() = + runTest { + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithCustomKey.name) + val first = reqShieldAspect.getCacheableCacheKey(joinPoint) + + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithOtherCacheName.name) + val second = reqShieldAspect.getCacheableCacheKey(joinPoint) + + assertEquals(namespacedSpelKey, first) + assertEquals("OtherCacheName::$spelEvaluatedKey", second) + } + + @Test + fun cacheableRejectsNonSuspendTarget() = + runTest { + // A non-suspend join point has no trailing Continuation argument. + every { joinPoint.args } returns arrayOf(argument) + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithKeyGenerator.name) + + val exception = + assertFailsWith { + reqShieldAspect.aroundReqShieldCacheable(joinPoint) + } + + assertTrue(exception.message!!.contains(TestBean::cacheableWithKeyGenerator.name), exception.message) + } + + @Test + fun cacheEvictRejectsNonSuspendTarget() = + runTest { + every { joinPoint.args } returns arrayOf(argument) + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::evictNonSuspend.name) + + val exception = + assertFailsWith { + reqShieldAspect.aroundReqShieldCacheEvict(joinPoint) + } + + assertTrue(exception.message!!.contains(TestBean::evictNonSuspend.name), exception.message) + } + + @Test + fun keyAndKeyGeneratorAreMutuallyExclusive() = + runTest { + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithKeyAndGenerator.name) + + val exception = + assertFailsWith { + reqShieldAspect.getCacheableCacheKey(joinPoint) + } + + assertTrue(exception.message!!.contains("mutually exclusive"), exception.message) + } + + @Test + fun aKeyThatResolvesToNothingIsRejected() = + runTest { + every { reqShieldAspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithUnresolvableKey.name) + + val exception = + assertFailsWith { + reqShieldAspect.getCacheableCacheKey(joinPoint) + } + + assertTrue(exception.message!!.contains(TestBean::cacheableWithUnresolvableKey.name), exception.message) + } + + @Test + fun globalLockRequiresTheCacheToImplementGlobalLockSupport() = + runTest { + val plainCache = mockk>() + val aspect = spyk(ReqShieldAspect(plainCache, scope)) + aspect.setBeanFactory(beanFactory) + every { aspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithGlobalLock.name) + + val exception = + assertFailsWith { + aspect.aroundReqShieldCacheable(joinPoint) + } + + assertTrue(exception.message!!.contains("GlobalLockSupport"), exception.message) + assertTrue(exception.message!!.contains(TestBean::cacheableWithGlobalLock.name), exception.message) + } + + @Test + fun globalLockIsAcquiredAndReleasedWithTheSameToken() = + runTest { + val lockableCache = mockk>() + val tokenSlot = slot() + coEvery { lockableCache.get(any()) } returns null + coEvery { lockableCache.put(any(), any(), any()) } returns true + coEvery { lockableCache.globalLock(any(), capture(tokenSlot), any()) } returns true + coEvery { lockableCache.globalUnLock(any(), any()) } returns true + + val lockScope = CoroutineScope(SupervisorJob() + Dispatchers.Default) + val aspect = spyk(ReqShieldAspect(lockableCache, lockScope)) + aspect.setBeanFactory(beanFactory) + every { aspect.getTargetMethod(joinPoint) } returns methodOf(TestBean::cacheableWithGlobalLock.name) + coEvery { joinPoint.proceed(any>()) } coAnswers { targetObject.cacheableWithGlobalLock(argument) } + + aspect.aroundReqShieldCacheable(joinPoint) + lockScope.awaitBackgroundWrites() + + val expectedLockKey = "$LOCK_KEY_PREFIX${namespacedSpelKey}_CREATE" + assertTrue(tokenSlot.isCaptured, "the aspect never called globalLock") + coVerify(exactly = 1) { + lockableCache.globalLock(expectedLockKey, tokenSlot.captured, DEFAULT_LOCK_TIMEOUT_MILLIS) + } + coVerify(exactly = 1) { lockableCache.globalUnLock(expectedLockKey, tokenSlot.captured) } + } + + /** A cache that opts in to global locking, as production code is expected to do. */ + interface LockableAsyncCache : + AsyncCache, + GlobalLockSupport + class TestBean { @ReqShieldCacheable(cacheName = "TestCacheName", key = "#paramMap['x'] + #paramMap['y']") suspend fun cacheableWithCustomKey(paramMap: Map): String = "ReturnValue: $paramMap" + @ReqShieldCacheable(cacheName = "OtherCacheName", key = "#paramMap['x'] + #paramMap['y']") + suspend fun cacheableWithOtherCacheName(paramMap: Map): String = "ReturnValue: $paramMap" + + @ReqShieldCacheable(cacheName = "TestCacheName", key = "#paramMap['x'] + #paramMap['y']", isLocalLock = false) + suspend fun cacheableWithGlobalLock(paramMap: Map): String = "ReturnValue: $paramMap" + + @ReqShieldCacheable(cacheName = "TestCacheName", key = "#paramMap['missing']") + suspend fun cacheableWithUnresolvableKey(paramMap: Map): String = "ReturnValue: $paramMap" + + @ReqShieldCacheable(cacheName = "TestCacheName", key = "#paramMap['x']", keyGenerator = "customGenerator") + suspend fun cacheableWithKeyAndGenerator(paramMap: Map): String = "ReturnValue: $paramMap" + @ReqShieldCacheable(cacheName = "TestCacheName") suspend fun cacheableWithDefaultKeyGenerator(paramMap: Map): String = "ReturnValue: $paramMap" @@ -224,6 +338,12 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { log.debug("cache eviction") return true } + + @ReqShieldCacheEvict(cacheName = "TestCacheName", key = "#paramMap['x'] + #paramMap['y']") + suspend fun evictFailing(paramMap: Map): Boolean = throw IllegalStateException("eviction target failed: $paramMap") + + @ReqShieldCacheEvict(cacheName = "TestCacheName", key = "#paramMap['x'] + #paramMap['y']") + fun evictNonSuspend(paramMap: Map): Boolean = paramMap.isNotEmpty() } class CustomGenerator : KeyGenerator { diff --git a/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/annotation/ReqShieldCacheEvict.kt b/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/annotation/ReqShieldCacheEvict.kt index 7296434..bcc6aad 100644 --- a/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/annotation/ReqShieldCacheEvict.kt +++ b/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/annotation/ReqShieldCacheEvict.kt @@ -16,13 +16,13 @@ package com.linecorp.cse.reqshield.spring.webflux.annotation +import java.lang.annotation.Inherited + @Target(AnnotationTarget.FUNCTION) @Retention(AnnotationRetention.RUNTIME) +@Inherited annotation class ReqShieldCacheEvict( val cacheName: String, val key: String = "", val keyGenerator: String = "", - val isLocalLock: Boolean = true, - val lockTimeoutMillis: Long = 3000, - val condition: String = "", ) diff --git a/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/annotation/ReqShieldCacheable.kt b/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/annotation/ReqShieldCacheable.kt index 5d7ba67..bdea873 100644 --- a/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/annotation/ReqShieldCacheable.kt +++ b/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/annotation/ReqShieldCacheable.kt @@ -17,6 +17,9 @@ package com.linecorp.cse.reqshield.spring.webflux.annotation import com.linecorp.cse.reqshield.reactor.config.ReqShieldWorkMode +import com.linecorp.cse.reqshield.support.constant.ConfigValues.DEFAULT_DECISION_FOR_UPDATE +import com.linecorp.cse.reqshield.support.constant.ConfigValues.DEFAULT_LOCK_TIMEOUT_MILLIS +import com.linecorp.cse.reqshield.support.constant.ConfigValues.DEFAULT_TIME_TO_LIVE_MILLIS import com.linecorp.cse.reqshield.support.constant.ConfigValues.MAX_ATTEMPT_GET_CACHE import java.lang.annotation.Inherited @@ -28,10 +31,10 @@ annotation class ReqShieldCacheable( val key: String = "", val keyGenerator: String = "", val isLocalLock: Boolean = true, - val lockTimeoutMillis: Long = 3000, - val decisionForUpdate: Int = 90, + val lockTimeoutMillis: Long = DEFAULT_LOCK_TIMEOUT_MILLIS, + val decisionForUpdate: Int = DEFAULT_DECISION_FOR_UPDATE, val maxAttemptGetCache: Int = MAX_ATTEMPT_GET_CACHE, - val timeToLiveMillis: Long = 10 * 60 * 1000, + val timeToLiveMillis: Long = DEFAULT_TIME_TO_LIVE_MILLIS, val reqShieldWorkMode: ReqShieldWorkMode = ReqShieldWorkMode.CREATE_AND_UPDATE_CACHE, val nullHandling: NullHandling = NullHandling.EMIT_EMPTY, ) diff --git a/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspect.kt b/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspect.kt index 429e0b5..53549e6 100644 --- a/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspect.kt +++ b/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspect.kt @@ -18,15 +18,18 @@ package com.linecorp.cse.reqshield.spring.webflux.aspect import com.linecorp.cse.reqshield.reactor.ReqShield import com.linecorp.cse.reqshield.reactor.config.ReqShieldConfiguration +import com.linecorp.cse.reqshield.spring.webflux.annotation.NullHandling 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.spring.webflux.cache.GlobalLockSupport import org.aspectj.lang.ProceedingJoinPoint import org.aspectj.lang.annotation.Around import org.aspectj.lang.annotation.Aspect import org.aspectj.lang.reflect.MethodSignature import org.springframework.beans.factory.BeanFactory import org.springframework.beans.factory.BeanFactoryAware +import org.springframework.beans.factory.annotation.Qualifier import org.springframework.cache.interceptor.KeyGenerator import org.springframework.cache.interceptor.SimpleKeyGenerator import org.springframework.context.expression.MethodBasedEvaluationContext @@ -35,32 +38,40 @@ import org.springframework.core.annotation.AnnotationUtils import org.springframework.expression.EvaluationContext import org.springframework.expression.Expression import org.springframework.expression.spel.standard.SpelExpressionParser -import org.springframework.stereotype.Component import org.springframework.util.StringUtils import org.springframework.util.function.SingletonSupplier import reactor.core.publisher.Mono +import reactor.core.scheduler.Scheduler import java.lang.reflect.Method import java.util.concurrent.ConcurrentHashMap @Aspect -@Component open class ReqShieldAspect( private val asyncCache: AsyncCache, + @Qualifier("reqShieldScheduler") private val scheduler: Scheduler, ) : BeanFactoryAware { private lateinit var beanFactory: BeanFactory private val spelParser = SpelExpressionParser() - private var defaultKeyGenerator = SingletonSupplier.of { SimpleKeyGenerator() } + private val parameterNameDiscoverer = DefaultParameterNameDiscoverer() + private val defaultKeyGenerator = SingletonSupplier.of { SimpleKeyGenerator() } + + /** Global locking is only available when the cache implementation opts in to it. */ + private val lockSupport = asyncCache as? GlobalLockSupport private val keyGeneratorMap = ConcurrentHashMap() - internal val reqShieldMap = ConcurrentHashMap>() + private val expressionMap = ConcurrentHashMap() + + /** A ReqShield is configured by the annotation alone, so one instance per annotated method is enough. */ + internal val reqShieldMap = ConcurrentHashMap>() @Around("@annotation(com.linecorp.cse.reqshield.spring.webflux.annotation.ReqShieldCacheable)") fun aroundTargetCacheable(joinPoint: ProceedingJoinPoint): Mono { val annotation = getCacheableAnnotation(joinPoint) - val reqShield = getOrCreateReqShield(joinPoint) val cacheKey = getCacheableCacheKey(joinPoint) + val reqShield = getOrCreateReqShield(joinPoint) - val resultMono = + // Mono.map rejects a null result, so the null decision is taken on the wrapper instead of on the value. + val reqShieldDataMono = reqShield .getAndSetReqShieldData( cacheKey, @@ -68,13 +79,14 @@ open class ReqShieldAspect( joinPoint.proceed() as Mono }, annotation.timeToLiveMillis, - ).map { it.value } + ) return when (annotation.nullHandling) { - com.linecorp.cse.reqshield.spring.webflux.annotation.NullHandling.EMIT_EMPTY -> - resultMono.flatMap { Mono.justOrEmpty(it) } - com.linecorp.cse.reqshield.spring.webflux.annotation.NullHandling.ERROR -> - resultMono.flatMap { value -> + NullHandling.EMIT_EMPTY -> + reqShieldDataMono.flatMap { Mono.justOrEmpty(it.value) } + NullHandling.ERROR -> + reqShieldDataMono.flatMap { reqShieldData -> + val value = reqShieldData.value if (value == null) { Mono.error(IllegalStateException("ReqShieldCacheable returned null for key=$cacheKey")) } else { @@ -88,37 +100,92 @@ open class ReqShieldAspect( fun aroundReqShieldCacheEvict(joinPoint: ProceedingJoinPoint): Mono { val cacheKey = getCacheEvictCacheKey(joinPoint) + // Same default as Spring's @CacheEvict: a method that fails leaves the cache untouched. + // An empty completion (for example Mono) is a success too, so it evicts as well. return Mono - .defer { - asyncCache.evict(cacheKey) - }.flatMap { - joinPoint.proceed() as Mono - } + .defer { joinPoint.proceed() as Mono } + // A Mono never emits null, so the non-null assertion below can never fail. + .flatMap { result -> asyncCache.evict(cacheKey).thenReturn(result!!) } + .switchIfEmpty(Mono.defer { asyncCache.evict(cacheKey).then(Mono.empty()) }) } - internal fun getCacheableAnnotation(joinPoint: ProceedingJoinPoint): ReqShieldCacheable = - AnnotationUtils.getAnnotation(getTargetMethod(joinPoint), ReqShieldCacheable::class.java) - ?: throw IllegalArgumentException("ReqShieldCacheable annotation is required") + private fun getOrCreateReqShield(joinPoint: ProceedingJoinPoint): ReqShield = + reqShieldMap.computeIfAbsent(getTargetMethod(joinPoint)) { + createReqShield(joinPoint) + } - internal fun getCacheableCacheKey(joinPoint: ProceedingJoinPoint): String { + private fun createReqShield(joinPoint: ProceedingJoinPoint): ReqShield { + val method = getTargetMethod(joinPoint) val annotation = getCacheableAnnotation(joinPoint) - validateCacheKey(annotation.key, annotation.keyGenerator) - return getCacheKeyOrDefault(annotation.key, annotation.keyGenerator, joinPoint) + require(annotation.isLocalLock || lockSupport != null) { + "isLocalLock = false on ${method.declaringClass.name}.${method.name} " + + "requires the AsyncCache bean to implement GlobalLockSupport" + } + + val reqShieldConfiguration = + ReqShieldConfiguration( + setCacheFunction = { key, reqShieldData, timeToLiveMillis -> + asyncCache.put(key, reqShieldData, timeToLiveMillis) + }, + getCacheFunction = { key -> + asyncCache.get(key) + }, + globalLockFunction = + lockSupport?.let { support -> + { lockKey, token, timeToLiveMillis -> support.globalLock(lockKey, token, timeToLiveMillis) } + }, + globalUnLockFunction = + lockSupport?.let { support -> + { lockKey, token -> support.globalUnLock(lockKey, token) } + }, + isLocalLock = annotation.isLocalLock, + lockTimeoutMillis = annotation.lockTimeoutMillis, + scheduler = scheduler, + decisionForUpdate = annotation.decisionForUpdate, + maxAttemptGetCache = annotation.maxAttemptGetCache, + reqShieldWorkMode = annotation.reqShieldWorkMode, + ) + + return ReqShield(reqShieldConfiguration) } + internal open fun getTargetMethod(joinPoint: ProceedingJoinPoint): Method = (joinPoint.signature as MethodSignature).method + + internal fun getCacheableAnnotation(joinPoint: ProceedingJoinPoint): ReqShieldCacheable = + AnnotationUtils.getAnnotation(getTargetMethod(joinPoint), ReqShieldCacheable::class.java) + ?: throw IllegalArgumentException("ReqShieldCacheable annotation is required") + internal fun getCacheEvictAnnotation(joinPoint: ProceedingJoinPoint): ReqShieldCacheEvict = AnnotationUtils.getAnnotation(getTargetMethod(joinPoint), ReqShieldCacheEvict::class.java) ?: throw IllegalArgumentException("ReqShieldCacheEvict annotation is required") + internal fun getCacheableCacheKey(joinPoint: ProceedingJoinPoint): String { + val annotation = getCacheableAnnotation(joinPoint) + + return buildCacheKey(annotation.cacheName, annotation.key, annotation.keyGenerator, joinPoint) + } + internal fun getCacheEvictCacheKey(joinPoint: ProceedingJoinPoint): String { val annotation = getCacheEvictAnnotation(joinPoint) - validateCacheKey(annotation.key, annotation.keyGenerator) - return getCacheKeyOrDefault(annotation.key, annotation.keyGenerator, joinPoint) + return buildCacheKey(annotation.cacheName, annotation.key, annotation.keyGenerator, joinPoint) } - internal open fun getTargetMethod(joinPoint: ProceedingJoinPoint): Method = (joinPoint.signature as MethodSignature).method + /** + * Namespaces the resolved key with the cache name so that entries (and the locks derived from + * them) of different caches cannot collide. Eviction follows the same rule so that it matches. + */ + private fun buildCacheKey( + cacheName: String, + annotationCacheKey: String, + annotationCacheKeyGenerator: String, + joinPoint: ProceedingJoinPoint, + ): String { + validateCacheKey(annotationCacheKey, annotationCacheKeyGenerator) + + return "$cacheName::${getCacheKeyOrDefault(annotationCacheKey, annotationCacheKeyGenerator, joinPoint)}" + } private fun getCacheKeyOrDefault( annotationCacheKey: String, @@ -127,12 +194,11 @@ open class ReqShieldAspect( ): String { val method = getTargetMethod(joinPoint) val context: EvaluationContext = - MethodBasedEvaluationContext(joinPoint.target, method, joinPoint.args, DefaultParameterNameDiscoverer()) + MethodBasedEvaluationContext(joinPoint.target, method, joinPoint.args, parameterNameDiscoverer) val key = if (StringUtils.hasText(annotationCacheKey)) { - val expression: Expression = spelParser.parseExpression(annotationCacheKey) - expression.getValue(context, String::class.java) + getOrParseExpression(annotationCacheKey).getValue(context, String::class.java) } else { val keyGenerator = getOrCreateKeyGenerator(annotationCacheKeyGenerator) keyGenerator.generate(joinPoint.target, method, joinPoint.args).toString() @@ -151,39 +217,6 @@ open class ReqShieldAspect( return key } - private fun getOrCreateReqShield(joinPoint: ProceedingJoinPoint): ReqShield = - reqShieldMap.computeIfAbsent(generateReqShieldKey(joinPoint)) { - createReqShield(joinPoint) - } - - private fun createReqShield(joinPoint: ProceedingJoinPoint): ReqShield { - val annotation = getCacheableAnnotation(joinPoint) - - val reqShieldConfiguration = - ReqShieldConfiguration( - setCacheFunction = { key, reqShieldData, timeToLiveMillis -> - asyncCache.put(key, reqShieldData, timeToLiveMillis) - }, - getCacheFunction = { key -> - asyncCache.get(key) - }, - globalLockFunction = { key, timeToLiveMillis -> - asyncCache.globalLock(key, timeToLiveMillis) - }, - globalUnLockFunction = { key -> - asyncCache.globalUnLock(key) - }, - isLocalLock = annotation.isLocalLock, - lockTimeoutMillis = annotation.lockTimeoutMillis, - decisionForUpdate = annotation.decisionForUpdate, - maxAttemptGetCache = annotation.maxAttemptGetCache, - reqShieldWorkMode = annotation.reqShieldWorkMode, - scheduler = beanFactory.getBean("reqShieldScheduler", reactor.core.scheduler.Scheduler::class.java), - ) - - return ReqShield(reqShieldConfiguration) - } - private fun validateCacheKey( cacheKey: String, cacheKeyGenerator: String, @@ -195,6 +228,11 @@ open class ReqShieldAspect( } } + private fun getOrParseExpression(cacheKeyExpression: String): Expression = + expressionMap.computeIfAbsent(cacheKeyExpression) { + spelParser.parseExpression(it) + } + private fun getOrCreateKeyGenerator(keyGeneratorBeanName: String?): KeyGenerator { if (keyGeneratorBeanName.isNullOrBlank()) { return defaultKeyGenerator.obtain() @@ -205,12 +243,6 @@ open class ReqShieldAspect( } } - private fun generateReqShieldKey(joinPoint: ProceedingJoinPoint): String { - val method = getTargetMethod(joinPoint) - return "${method.declaringClass.name}.${method.name}-" + - "${getCacheableAnnotation(joinPoint).cacheName}-${getCacheableCacheKey(joinPoint)}" - } - override fun setBeanFactory(beanFactory: BeanFactory) { this.beanFactory = beanFactory } diff --git a/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/cache/AsyncCache.kt b/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/cache/AsyncCache.kt index 0690486..a6c7ee7 100644 --- a/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/cache/AsyncCache.kt +++ b/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/cache/AsyncCache.kt @@ -29,32 +29,4 @@ interface AsyncCache { ): Mono fun evict(key: String): Mono - - /** - * Attempt a global lock on a specific key. - * - * @param key The key to get the lock. - * @param timeToLiveMillis The validity of the lock in milliseconds. - * @return Whether the lock was successfully obtained. Returns `true` by default. - * - * This method provides a default implementation, but if you need a locking mechanism - * You must implement and use your own locking logic. The default implementation is true, and if the value of ReqShieldConfiguration > isLocalLock is false, you will use the function you implemented. - * actual production environments should override this method appropriately to manage locks. - */ - fun globalLock( - key: String, - timeToLiveMillis: Long, - ): Mono = Mono.just(true) - - /** - * Releases the global lock on a specific key. - * - * @param key The key you want to unlock. - * @return Whether the lock was successfully obtained. Returns `true` by default. - * - * This method provides a default implementation, but if you need a locking mechanism - * You must implement and use your own unlocking logic. The default implementation is true, and if the value of ReqShieldConfiguration > isLocalLock is false, you will use the function you implemented. - * actual production environments should override this method appropriately to manage locks. - */ - fun globalUnLock(key: String): Mono = Mono.just(true) } diff --git a/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/cache/GlobalLockSupport.kt b/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/cache/GlobalLockSupport.kt new file mode 100644 index 0000000..bae6567 --- /dev/null +++ b/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/cache/GlobalLockSupport.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.spring.webflux.cache + +import reactor.core.publisher.Mono + +/** Implement alongside [AsyncCache] to use `@ReqShieldCacheable(isLocalLock = false)`. */ +interface GlobalLockSupport { + /** Acquire lockKey for the caller identified by token (Redis: `SET lockKey token NX PX timeToLiveMillis`). */ + fun globalLock( + lockKey: String, + token: String, + timeToLiveMillis: Long, + ): Mono + + /** Release only if the stored value equals token (Redis: compare-and-delete Lua script). */ + fun globalUnLock( + lockKey: String, + token: String, + ): Mono +} diff --git a/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/config/LibAutoConfiguration.kt b/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/config/LibAutoConfiguration.kt index 2a46c27..b9af186 100644 --- a/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/config/LibAutoConfiguration.kt +++ b/core-spring-webflux/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/config/LibAutoConfiguration.kt @@ -28,6 +28,12 @@ import reactor.core.scheduler.Schedulers @EnableAspectJAutoProxy(proxyTargetClass = true) @Import(ReqShieldAspect::class) open class LibAutoConfiguration { - @Bean + /** + * Scheduler shared by every [com.linecorp.cse.reqshield.reactor.ReqShield] the aspect creates, used for + * the asynchronous cache writes and for polling the cache while another request holds the lock. + * + * [Schedulers.boundedElastic] is Reactor's process-wide instance, so the container must never dispose it. + */ + @Bean(destroyMethod = "") open fun reqShieldScheduler(): Scheduler = Schedulers.boundedElastic() } diff --git a/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/InMemoryAsyncCache.kt b/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/InMemoryAsyncCache.kt index 15a5631..ffaa7f9 100644 --- a/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/InMemoryAsyncCache.kt +++ b/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/InMemoryAsyncCache.kt @@ -17,16 +17,24 @@ package com.linecorp.cse.reqshield.spring.webflux.aspect import com.linecorp.cse.reqshield.spring.webflux.cache.AsyncCache +import com.linecorp.cse.reqshield.spring.webflux.cache.GlobalLockSupport import com.linecorp.cse.reqshield.support.model.ReqShieldData import reactor.core.publisher.Mono import java.util.concurrent.ConcurrentHashMap -import java.util.concurrent.Semaphore -class InMemoryAsyncCache : AsyncCache { +/** + * In-memory stand-in for a distributed cache, with the same token semantics a Redis + * implementation has: `SET NX PX` on lock and compare-and-delete on unlock. + */ +class InMemoryAsyncCache : + AsyncCache, + GlobalLockSupport { private data class Entry(val data: ReqShieldData, val expiresAt: Long) + private data class Lock(val token: String, val expiresAt: Long) + private val store = ConcurrentHashMap>() - private val locks = ConcurrentHashMap() + private val locks = ConcurrentHashMap() override fun get(key: String): Mono?> = Mono.fromCallable { @@ -48,16 +56,39 @@ class InMemoryAsyncCache : AsyncCache { override fun evict(key: String): Mono = Mono.fromCallable { store.remove(key) != null } override fun globalLock( - key: String, + lockKey: String, + token: String, timeToLiveMillis: Long, ): Mono = Mono.fromCallable { - locks.computeIfAbsent(key) { Semaphore(1) }.tryAcquire() + val now = System.currentTimeMillis() + // SET NX PX: only a key without a live lock can be taken, and the owner is remembered. + val holder = + locks.compute(lockKey) { _, current -> + if (current == null || now > current.expiresAt) { + Lock(token, now + timeToLiveMillis) + } else { + current + } + }!! + holder.token == token } - override fun globalUnLock(key: String): Mono = + override fun globalUnLock( + lockKey: String, + token: String, + ): Mono = Mono.fromCallable { - locks[key]?.release() - true + var released = false + // Compare-and-delete: an owner whose lock already expired must not release the next owner's lock. + locks.compute(lockKey) { _, current -> + if (current != null && current.token == token) { + released = true + null + } else { + current + } + } + released } } diff --git a/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/LocalOnlyAsyncCache.kt b/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/LocalOnlyAsyncCache.kt new file mode 100644 index 0000000..bb86f3f --- /dev/null +++ b/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/LocalOnlyAsyncCache.kt @@ -0,0 +1,28 @@ +/* + * 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.spring.webflux.aspect + +import com.linecorp.cse.reqshield.spring.webflux.cache.AsyncCache + +/** + * A cache that deliberately does not implement + * [com.linecorp.cse.reqshield.spring.webflux.cache.GlobalLockSupport], so `isLocalLock = false` + * must be rejected. + */ +class LocalOnlyAsyncCache( + delegate: AsyncCache = InMemoryAsyncCache(), +) : AsyncCache by delegate diff --git a/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/RedisAsyncCache.kt b/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/RedisAsyncCache.kt new file mode 100644 index 0000000..69d9762 --- /dev/null +++ b/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/RedisAsyncCache.kt @@ -0,0 +1,75 @@ +/* + * 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.spring.webflux.aspect + +import com.linecorp.cse.reqshield.spring.webflux.cache.AsyncCache +import com.linecorp.cse.reqshield.spring.webflux.cache.GlobalLockSupport +import com.linecorp.cse.reqshield.support.model.ReqShieldData +import io.lettuce.core.ScriptOutputType +import io.lettuce.core.SetArgs +import io.lettuce.core.api.reactive.RedisReactiveCommands +import io.lettuce.core.api.sync.RedisCommands +import reactor.core.publisher.Mono + +private const val UNLOCK_SCRIPT = "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end" + +/** + * Reference Redis implementation of the documented locking recipe: `SET NX PX` to acquire and a + * compare-and-delete Lua script to release, so only the token holder can release the lock. + */ +class RedisAsyncCache( + private val sync: RedisCommands, + private val reactive: RedisReactiveCommands, +) : AsyncCache, + GlobalLockSupport { + override fun get(key: String): Mono?> = + Mono.fromCallable { + sync.get(key)?.let { ReqShieldData(value = it, timeToLiveMillis = 10_000) } + } + + override fun put( + key: String, + value: ReqShieldData, + timeToLiveMillis: Long, + ): Mono = + Mono.fromCallable { + sync.psetex(key, timeToLiveMillis, value.value ?: "") + true + } + + override fun evict(key: String): Mono = Mono.fromCallable { sync.del(key) > 0 } + + override fun globalLock( + lockKey: String, + token: String, + timeToLiveMillis: Long, + ): Mono = + reactive + .set(lockKey, token, SetArgs.Builder.nx().px(timeToLiveMillis)) + .map { it == "OK" } + .defaultIfEmpty(false) + + override fun globalUnLock( + lockKey: String, + token: String, + ): Mono = + reactive + .eval(UNLOCK_SCRIPT, ScriptOutputType.INTEGER, arrayOf(lockKey), token) + .next() + .map { it == 1L } + .defaultIfEmpty(false) +} diff --git a/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectIntegrationTest.kt b/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectIntegrationTest.kt index 577ac0c..2d9b5b0 100644 --- a/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectIntegrationTest.kt +++ b/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectIntegrationTest.kt @@ -20,9 +20,13 @@ 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.spring.webflux.config.LibAutoConfiguration +import com.linecorp.cse.reqshield.support.model.ReqShieldData import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNotNull +import org.junit.jupiter.api.Assertions.assertNull import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.Test +import org.junit.jupiter.api.assertThrows import org.junit.jupiter.api.extension.ExtendWith import org.springframework.beans.factory.annotation.Autowired import org.springframework.context.annotation.Bean @@ -32,17 +36,40 @@ import org.springframework.test.context.junit.jupiter.SpringExtension import reactor.core.publisher.Flux import reactor.core.publisher.Mono import reactor.core.scheduler.Schedulers +import reactor.test.StepVerifier import java.util.concurrent.atomic.AtomicInteger +private const val INTEGRATION_CACHE_NAME = "it" + @ExtendWith(SpringExtension::class) @ContextConfiguration(classes = [LibAutoConfiguration::class, ReqShieldAspectIntegrationTest.TestConfig::class]) class ReqShieldAspectIntegrationTest { @Autowired private lateinit var service: TestService + @Autowired + private lateinit var asyncCache: AsyncCache + + private fun cacheKey(key: String) = "$INTEGRATION_CACHE_NAME::$key" + + /** ReqShield writes the cache asynchronously, so a test that needs the entry has to wait for it. */ + private fun awaitCachePut( + key: String, + timeoutMillis: Long = 2_000, + ): Boolean { + val start = System.currentTimeMillis() + while (System.currentTimeMillis() - start < timeoutMillis) { + if (asyncCache.get(key).block() != null) { + return true + } + Thread.sleep(10) + } + return false + } + @Test fun shouldCollapseDuplicateRequests() { - val key = "dup" + val key = "dup-${System.nanoTime()}" val attempts = 20 val result = @@ -59,8 +86,10 @@ class ReqShieldAspectIntegrationTest { @Test fun shouldEvictAndRecompute() { - val key = "evict" + val key = "evict-${System.nanoTime()}" val v1 = service.get(key).block() + assertTrue(awaitCachePut(cacheKey(key)), "Timed out waiting for cache put for key=${cacheKey(key)}") + val evicted = service.evict(key).block() val v2 = service.get(key).block() @@ -70,6 +99,45 @@ class ReqShieldAspectIntegrationTest { assertTrue(v1 != v2) } + @Test + fun shouldEvictAfterAMethodThatCompletesEmpty() { + val key = "evict-void-${System.nanoTime()}" + asyncCache.put(cacheKey(key), ReqShieldData("cached", 10_000), 10_000).block() + + StepVerifier + .create(service.evictWithoutResult(key)) + .verifyComplete() + + assertEquals(1, service.getEmptyEvictCount(), "The annotated method must still run") + assertNull(asyncCache.get(cacheKey(key)).block(), "An empty completion is a success, so it evicts") + } + + @Test + fun shouldNotEvictWhenTheAnnotatedMethodFails() { + val key = "evict-fail-${System.nanoTime()}" + asyncCache.put(cacheKey(key), ReqShieldData("cached", 10_000), 10_000).block() + + assertThrows { service.evictFailing(key).block() } + + assertNotNull(asyncCache.get(cacheKey(key)).block(), "A failed method must leave the cache untouched") + } + + @Test + fun shouldCollapseDuplicateRequestsWithGlobalLock() { + val key = "dup-global-${System.nanoTime()}" + val attempts = 20 + + val result = + Flux + .range(1, attempts) + .flatMap { service.getWithGlobalLock(key).subscribeOn(Schedulers.boundedElastic()) } + .collectList() + .block()!! + + assertEquals(attempts, result.size) + assertEquals(1, service.getGlobalLockCount(), "The backend should be called once") + } + @Configuration open class TestConfig { @Bean @@ -80,12 +148,69 @@ class ReqShieldAspectIntegrationTest { } open class TestService { - val counter = AtomicInteger(0) + private val counter = AtomicInteger(0) + private val globalLockCounter = AtomicInteger(0) + private val emptyEvictCounter = AtomicInteger(0) + + // Read through open methods: a CGLIB proxy cannot delegate the final getter of a Kotlin property. + open fun getGlobalLockCount(): Int = globalLockCounter.get() - @ReqShieldCacheable(cacheName = "it", key = "#key", timeToLiveMillis = 10_000) + open fun getEmptyEvictCount(): Int = emptyEvictCounter.get() + + @ReqShieldCacheable(cacheName = INTEGRATION_CACHE_NAME, key = "#key", timeToLiveMillis = 10_000) open fun get(key: String): Mono = Mono.fromCallable { "value-" + counter.incrementAndGet() } - @ReqShieldCacheEvict(cacheName = "it", key = "#key") + @ReqShieldCacheable(cacheName = INTEGRATION_CACHE_NAME, key = "#key", timeToLiveMillis = 10_000, isLocalLock = false) + open fun getWithGlobalLock(key: String): Mono = Mono.fromCallable { "value-" + globalLockCounter.incrementAndGet() } + + @ReqShieldCacheEvict(cacheName = INTEGRATION_CACHE_NAME, key = "#key") open fun evict(key: String): Mono = Mono.just(true) + + @ReqShieldCacheEvict(cacheName = INTEGRATION_CACHE_NAME, key = "#key") + open fun evictWithoutResult(key: String): Mono = Mono.fromRunnable { emptyEvictCounter.incrementAndGet() } + + @ReqShieldCacheEvict(cacheName = INTEGRATION_CACHE_NAME, key = "#key") + open fun evictFailing(key: String): Mono = Mono.error(IllegalStateException("eviction must not happen")) + } +} + +/** + * `isLocalLock = false` needs its own context: the cache bean here deliberately does not implement + * [com.linecorp.cse.reqshield.spring.webflux.cache.GlobalLockSupport]. + */ +@ExtendWith(SpringExtension::class) +@ContextConfiguration(classes = [LibAutoConfiguration::class, ReqShieldAspectWithoutGlobalLockSupportTest.TestConfig::class]) +class ReqShieldAspectWithoutGlobalLockSupportTest { + @Autowired + private lateinit var service: TestService + + @Test + fun globalLockWithoutGlobalLockSupportFailsOnTheFirstCall() { + val exception = assertThrows { service.getWithGlobalLock("key").block() } + + assertTrue( + exception.message!!.contains("requires the AsyncCache bean to implement GlobalLockSupport"), + "Unexpected message: ${exception.message}", + ) + assertEquals(0, service.getCallCount(), "The backend must not be called") + } + + @Configuration + open class TestConfig { + @Bean + open fun asyncCache(): AsyncCache = LocalOnlyAsyncCache() + + @Bean + open fun service(): TestService = TestService() + } + + open class TestService { + private val counter = AtomicInteger(0) + + // Read through an open method: a CGLIB proxy cannot delegate the final getter of a Kotlin property. + open fun getCallCount(): Int = counter.get() + + @ReqShieldCacheable(cacheName = INTEGRATION_CACHE_NAME, key = "#key", isLocalLock = false) + open fun getWithGlobalLock(key: String): Mono = Mono.fromCallable { "value-" + counter.incrementAndGet() } } } diff --git a/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectRedisIntegrationTest.kt b/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectRedisIntegrationTest.kt index eb69443..58ff1d2 100644 --- a/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectRedisIntegrationTest.kt +++ b/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectRedisIntegrationTest.kt @@ -19,12 +19,13 @@ package com.linecorp.cse.reqshield.spring.webflux.aspect 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.spring.webflux.cache.GlobalLockSupport import com.linecorp.cse.reqshield.spring.webflux.config.LibAutoConfiguration -import com.linecorp.cse.reqshield.support.model.ReqShieldData import com.linecorp.cse.reqshield.support.redis.AbstractRedisTest import io.lettuce.core.RedisClient import io.lettuce.core.api.StatefulRedisConnection import io.lettuce.core.api.sync.RedisCommands +import org.junit.jupiter.api.Assertions.assertFalse import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test @@ -38,8 +39,11 @@ import org.springframework.test.context.junit.jupiter.SpringExtension import reactor.core.publisher.Flux import reactor.core.publisher.Mono import reactor.core.scheduler.Schedulers +import java.util.UUID import java.util.concurrent.atomic.AtomicInteger +private const val CACHE_NAME = "it" + @ExtendWith(SpringExtension::class) @ContextConfiguration(classes = [LibAutoConfiguration::class, ReqShieldAspectRedisIntegrationTest.TestConfig::class]) class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { @@ -49,11 +53,16 @@ class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { @Autowired private lateinit var asyncCache: AsyncCache + @Autowired + private lateinit var lockSupport: GlobalLockSupport + @BeforeEach fun resetCounter() { service.resetCounter() } + private fun cacheKey(key: String) = "$CACHE_NAME::$key" + private fun awaitCachePut( key: String, timeoutMillis: Long = 2_000, @@ -90,12 +99,31 @@ class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { assertTrue(result.size == attempts && result.all { it != null }, "Expected all results to be valid") } + @Test + fun shouldCollapseDuplicateRequestsWithRedisGlobalLock() { + val key = "dup-redis-global-${System.nanoTime()}" // Use unique key for test isolation + val attempts = 20 + + val result = + Flux + .range(1, attempts) + .flatMap { service.getWithGlobalLock(key).subscribeOn(Schedulers.boundedElastic()) } + .collectList() + .block()!! + + assertTrue( + service.getRequestCount() == 1, + "Callable should be invoked only once. actual=${service.getRequestCount()}", + ) + assertTrue(result.size == attempts && result.all { it != null }, "Expected all results to be valid") + } + @Test fun shouldEvictAndRecomputeWithRedis() { val key = "evict-redis-${System.nanoTime()}" // Use unique key for test isolation val v1 = service.get(key).block() // ReqShield stores cache asynchronously; wait until the cache write is observed. - assertTrue(awaitCachePut(key), "Timed out waiting for cache put for key=$key") + assertTrue(awaitCachePut(cacheKey(key)), "Timed out waiting for cache put for key=${cacheKey(key)}") val evicted = service.evict(key).block() val v2 = service.get(key).block() @@ -103,6 +131,20 @@ class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { assertTrue(v1 != null && v2 != null && v1 != v2, "Values should differ after eviction: v1=$v1, v2=$v2") } + @Test + fun globalLockIsOnlyReleasedByTheTokenThatAcquiredIt() { + val lockKey = "lock-redis-${System.nanoTime()}" + val ownerToken = UUID.randomUUID().toString() + val otherToken = UUID.randomUUID().toString() + + assertTrue(lockSupport.globalLock(lockKey, ownerToken, 5_000).block() == true, "The free lock should be acquired") + assertFalse(lockSupport.globalLock(lockKey, otherToken, 5_000).block() == true, "A held lock must not be acquired again") + assertFalse(lockSupport.globalUnLock(lockKey, otherToken).block() == true, "A foreign token must not release the lock") + assertTrue(lockSupport.globalUnLock(lockKey, ownerToken).block() == true, "The owner should release the lock") + assertFalse(lockSupport.globalUnLock(lockKey, ownerToken).block() == true, "Releasing twice should report false") + assertTrue(lockSupport.globalLock(lockKey, otherToken, 5_000).block() == true, "The released lock should be acquirable") + } + @Configuration open class TestConfig { @Value("\${spring.redis.host}") @@ -118,37 +160,12 @@ class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { open fun redisConnection(redisClient: RedisClient): StatefulRedisConnection = redisClient.connect() @Bean - open fun asyncCache(redisConnection: StatefulRedisConnection): AsyncCache { + open fun asyncCache(redisConnection: StatefulRedisConnection): RedisAsyncCache { val sync: RedisCommands = redisConnection.sync() // Ensure clean DB state for tests running in CI runCatching { sync.flushdb() } - return object : AsyncCache { - override fun get(key: String): Mono?> = - Mono.fromCallable { - sync.get(key)?.let { ReqShieldData(value = it, timeToLiveMillis = 10_000) } - } - - override fun put( - key: String, - value: ReqShieldData, - timeToLiveMillis: Long, - ): Mono = - Mono.fromCallable { - sync.psetex(key, timeToLiveMillis, value.value ?: "") - true - } - - override fun evict(key: String): Mono = Mono.fromCallable { sync.del(key) > 0 } - - override fun globalLock( - key: String, - timeToLiveMillis: Long, - ): Mono = - Mono.fromCallable { sync.setnx("lock:$key", "1").also { if (it) sync.pexpire("lock:$key", timeToLiveMillis) } } - - override fun globalUnLock(key: String): Mono = Mono.fromCallable { sync.del("lock:$key") >= 0 } - } + return RedisAsyncCache(sync, redisConnection.reactive()) } @Bean @@ -164,10 +181,13 @@ class ReqShieldAspectRedisIntegrationTest : AbstractRedisTest() { open fun getRequestCount(): Int = counter.get() - @ReqShieldCacheable(cacheName = "it", key = "#key", timeToLiveMillis = 10_000) + @ReqShieldCacheable(cacheName = CACHE_NAME, key = "#key", timeToLiveMillis = 10_000) open fun get(key: String): Mono = Mono.fromCallable { "value-" + counter.incrementAndGet() } - @ReqShieldCacheEvict(cacheName = "it", key = "#key") + @ReqShieldCacheable(cacheName = CACHE_NAME, key = "#key", timeToLiveMillis = 10_000, isLocalLock = false) + open fun getWithGlobalLock(key: String): Mono = Mono.fromCallable { "value-" + counter.incrementAndGet() } + + @ReqShieldCacheEvict(cacheName = CACHE_NAME, key = "#key") open fun evict(key: String): Mono = Mono.just(true) } } diff --git a/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectTest.kt b/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectTest.kt index 8591fd8..d89a3df 100644 --- a/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectTest.kt +++ b/core-spring-webflux/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/aspect/ReqShieldAspectTest.kt @@ -16,6 +16,7 @@ package com.linecorp.cse.reqshield.spring.webflux.aspect +import com.linecorp.cse.reqshield.spring.webflux.annotation.NullHandling 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 @@ -29,6 +30,7 @@ import org.aspectj.lang.ProceedingJoinPoint import org.junit.jupiter.api.Assertions import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test +import org.junit.jupiter.api.assertThrows import org.springframework.beans.factory.BeanFactory import org.springframework.cache.interceptor.KeyGenerator import org.springframework.cache.interceptor.SimpleKeyGenerator @@ -39,19 +41,21 @@ import reactor.core.scheduler.Schedulers import reactor.test.StepVerifier import java.lang.reflect.Method import kotlin.test.assertEquals +import kotlin.test.assertNull import kotlin.test.assertTrue class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { private val asyncCache: AsyncCache = InMemoryAsyncCache() private val joinPoint = mockk() - private val reqShieldAspect = spyk(ReqShieldAspect(asyncCache)) + private val reqShieldAspect = spyk(ReqShieldAspect(asyncCache, Schedulers.boundedElastic())) private val targetObject = spyk(TestBean()) private val argument = mapOf("x" to "paramX", "y" to "paramY") private val cacheName = "TestCacheName" private val cacheKeyGenerator = "customGenerator" - private val spelEvaluatedKey = "paramXparamY" - private val keyGeneratorKey = "KeyGeneratedByGenerator" + private val spelEvaluatedKey = "$cacheName::paramXparamY" + private val keyGeneratorKey = "$cacheName::KeyGeneratedByGenerator" + private val defaultGeneratedKey = "$cacheName::${SimpleKeyGenerator.generateKey(arrayOf(argument))}" private val beanFactory = mockk() @@ -63,24 +67,22 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { every { joinPoint.target } returns targetObject reqShieldAspect.setBeanFactory(beanFactory) - // Provide scheduler bean expected by aspect configuration - every { - beanFactory.getBean("reqShieldScheduler", reactor.core.scheduler.Scheduler::class.java) - } returns Schedulers.boundedElastic() } + private fun stubTargetMethod(methodName: String) { + every { reqShieldAspect.getTargetMethod(joinPoint) } returns findTestBeanMethod(methodName) + } + + private fun findTestBeanMethod(methodName: String): Method = + ReflectionUtils.findMethod(TestBean::class.java, methodName, Map::class.java)!! + @Test override fun verifyReqShieldCacheCreation() { val reqShieldData = ReqShieldData(methodReturn, 1000) // pre-populate cache asyncCache.put(spelEvaluatedKey, reqShieldData, 1000).block() every { joinPoint.proceed() } answers { targetObject.cacheableWithCustomKey(argument) } - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::cacheableWithCustomKey.name, - Map::class.java, - )!! + stubTargetMethod(TestBean::cacheableWithCustomKey.name) // Test the aroundTargetCacheable method val result = reqShieldAspect.aroundTargetCacheable(joinPoint) @@ -91,9 +93,7 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { .assertNext { value -> assertEquals(reqShieldData.value, value) Assertions.assertTrue(reqShieldAspect.reqShieldMap.size == 1) - val method = reqShieldAspect.getTargetMethod(joinPoint) - val expectedKey = "${method.declaringClass.name}.${method.name}-$cacheName-$spelEvaluatedKey" - Assertions.assertNotNull(reqShieldAspect.reqShieldMap[expectedKey]) + Assertions.assertNotNull(reqShieldAspect.reqShieldMap[reqShieldAspect.getTargetMethod(joinPoint)]) }.verifyComplete() } @@ -102,12 +102,7 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { val reqShieldData = ReqShieldData(methodReturn, 1000) asyncCache.put(spelEvaluatedKey, reqShieldData, 1000).block() every { joinPoint.proceed() } answers { targetObject.cacheableWithCustomKey(argument) } - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::cacheableWithCustomKey.name, - Map::class.java, - )!! + stubTargetMethod(TestBean::cacheableWithCustomKey.name) val flux = Flux @@ -121,23 +116,17 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { StepVerifier .create(flux) .assertNext { productList -> + assertEquals(20, productList.size) Assertions.assertTrue(reqShieldAspect.reqShieldMap.size == 1) - val method = reqShieldAspect.getTargetMethod(joinPoint) - val expectedKey = "${method.declaringClass.name}.${method.name}-$cacheName-$spelEvaluatedKey" - Assertions.assertNotNull(reqShieldAspect.reqShieldMap[expectedKey]) + Assertions.assertNotNull(reqShieldAspect.reqShieldMap[reqShieldAspect.getTargetMethod(joinPoint)]) }.verifyComplete() } @Test override fun verifyReqShieldCacheEviction() { val reqShieldData = ReqShieldData(methodReturn, 1000) - asyncCache.put("${SimpleKeyGenerator.generateKey(arrayOf(argument))}", reqShieldData, 1000).block() - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::cacheableWithDefaultKeyGenerator.name, - Map::class.java, - )!! + asyncCache.put(defaultGeneratedKey, reqShieldData, 1000).block() + stubTargetMethod(TestBean::cacheableWithDefaultKeyGenerator.name) every { joinPoint.proceed() } answers { targetObject.cacheableWithCustomKey(argument) } // Test the aroundTargetCacheable method @@ -151,12 +140,7 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { // Validate cache eviction // real eviction call - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::evict.name, - Map::class.java, - )!! + stubTargetMethod(TestBean::evict.name) every { joinPoint.proceed() } answers { targetObject.evict(argument) } val removeProductMono = reqShieldAspect.aroundReqShieldCacheEvict(joinPoint) @@ -166,16 +150,14 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { .assertNext { value -> assertTrue(value as Boolean) }.verifyComplete() + + // The evicted entry is the namespaced one written by @ReqShieldCacheable + assertNull(asyncCache.get(defaultGeneratedKey).block()) } @Test override fun verifyCacheKeyGenerationWithSpEL() { - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::cacheableWithCustomKey.name, - Map::class.java, - )!! + stubTargetMethod(TestBean::cacheableWithCustomKey.name) // when, then Assertions.assertEquals( @@ -184,16 +166,12 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { ) } + @Test override fun verifyCacheKeyGenerationWithKeyGenerator() { // given every { beanFactory.getBean(cacheKeyGenerator, KeyGenerator::class.java) } returns CustomGenerator() - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::cacheableWithKeyGenerator.name, - Map::class.java, - )!! + stubTargetMethod(TestBean::cacheableWithKeyGenerator.name) // when, then Assertions.assertEquals( @@ -205,24 +183,143 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { @Test override fun verifyCacheKeyGenerationWithDefaultGenerator() { // given - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::cacheableWithDefaultKeyGenerator.name, - Map::class.java, - )!! + stubTargetMethod(TestBean::cacheableWithDefaultKeyGenerator.name) // when, then Assertions.assertEquals( - SimpleKeyGenerator.generateKey(arrayOf(argument)).toString(), + defaultGeneratedKey, reqShieldAspect.getCacheableCacheKey(joinPoint), ) } + @Test + fun evictKeyIsNamespacedWithTheCacheNameToo() { + stubTargetMethod(TestBean::evict.name) + + Assertions.assertEquals(defaultGeneratedKey, reqShieldAspect.getCacheEvictCacheKey(joinPoint)) + } + + @Test + fun reqShieldIsCreatedPerAnnotatedMethodEvenWhenTheCacheKeyMatches() { + every { joinPoint.proceed() } answers { targetObject.cacheableWithCustomKey(argument) } + stubTargetMethod(TestBean::cacheableWithCustomKey.name) + reqShieldAspect.aroundTargetCacheable(joinPoint).block() + + // Same cacheName and same resolved key, but another method: it must get its own ReqShield. + stubTargetMethod(TestBean::cacheableWithSameKeyOtherMethod.name) + every { joinPoint.proceed() } answers { targetObject.cacheableWithSameKeyOtherMethod(argument) } + reqShieldAspect.aroundTargetCacheable(joinPoint).block() + + assertEquals(2, reqShieldAspect.reqShieldMap.size) + Assertions.assertNotNull(reqShieldAspect.reqShieldMap[findTestBeanMethod(TestBean::cacheableWithCustomKey.name)]) + Assertions.assertNotNull(reqShieldAspect.reqShieldMap[findTestBeanMethod(TestBean::cacheableWithSameKeyOtherMethod.name)]) + } + + @Test + fun nullValueIsEmittedAsAnEmptyMonoByDefault() { + stubTargetMethod(TestBean::cacheableWithCustomKey.name) + every { joinPoint.proceed() } returns Mono.empty() + + StepVerifier + .create(reqShieldAspect.aroundTargetCacheable(joinPoint)) + .verifyComplete() + } + + @Test + fun nullValueFailsWhenNullHandlingIsError() { + stubTargetMethod(TestBean::cacheableWithNullHandlingError.name) + every { joinPoint.proceed() } returns Mono.empty() + + StepVerifier + .create(reqShieldAspect.aroundTargetCacheable(joinPoint)) + .verifyError(IllegalStateException::class.java) + } + + @Test + fun keyAndKeyGeneratorAreMutuallyExclusive() { + stubTargetMethod(TestBean::cacheableWithKeyAndKeyGenerator.name) + + val exception = assertThrows { reqShieldAspect.getCacheableCacheKey(joinPoint) } + assertTrue(exception.message!!.contains("mutually exclusive")) + } + + @Test + fun aBlankResolvedKeyIsRejected() { + stubTargetMethod(TestBean::cacheableWithUnresolvableKey.name) + + val exception = assertThrows { reqShieldAspect.getCacheableCacheKey(joinPoint) } + assertTrue(exception.message!!.contains("Null/blank key")) + } + + @Test + fun missingAnnotationsAreReported() { + every { reqShieldAspect.getTargetMethod(joinPoint) } returns + ReflectionUtils.findMethod(TestBean::class.java, TestBean::notAnnotated.name, Map::class.java)!! + + assertThrows { reqShieldAspect.getCacheableAnnotation(joinPoint) } + assertThrows { reqShieldAspect.getCacheEvictAnnotation(joinPoint) } + } + + @Test + fun evictionIsSkippedWhenTheAnnotatedMethodFails() { + val reqShieldData = ReqShieldData(methodReturn, 1000) + asyncCache.put(defaultGeneratedKey, reqShieldData, 1000).block() + stubTargetMethod(TestBean::evict.name) + every { joinPoint.proceed() } returns Mono.error(IllegalStateException("boom")) + + StepVerifier + .create(reqShieldAspect.aroundReqShieldCacheEvict(joinPoint)) + .verifyError(IllegalStateException::class.java) + + Assertions.assertNotNull(asyncCache.get(defaultGeneratedKey).block()) + } + + @Test + fun evictionAlsoRunsWhenTheAnnotatedMethodCompletesEmpty() { + val reqShieldData = ReqShieldData(methodReturn, 1000) + asyncCache.put(defaultGeneratedKey, reqShieldData, 1000).block() + stubTargetMethod(TestBean::evict.name) + every { joinPoint.proceed() } returns Mono.empty() + + StepVerifier + .create(reqShieldAspect.aroundReqShieldCacheEvict(joinPoint)) + .verifyComplete() + + assertNull(asyncCache.get(defaultGeneratedKey).block()) + } + + @Test + fun globalLockCollapsesRequestsWhenTheCacheSupportsIt() { + stubTargetMethod(TestBean::cacheableWithGlobalLock.name) + every { joinPoint.proceed() } answers { targetObject.cacheableWithGlobalLock(argument) } + + StepVerifier + .create(reqShieldAspect.aroundTargetCacheable(joinPoint)) + .assertNext { value -> assertEquals(methodReturn, value) } + .verifyComplete() + } + + @Test + fun globalLockRequiresTheCacheToImplementGlobalLockSupport() { + val localOnlyAspect = spyk(ReqShieldAspect(LocalOnlyAsyncCache(), Schedulers.boundedElastic())) + localOnlyAspect.setBeanFactory(beanFactory) + every { localOnlyAspect.getTargetMethod(joinPoint) } returns + findTestBeanMethod(TestBean::cacheableWithGlobalLock.name) + every { joinPoint.proceed() } answers { targetObject.cacheableWithGlobalLock(argument) } + + val exception = assertThrows { localOnlyAspect.aroundTargetCacheable(joinPoint) } + assertTrue(exception.message!!.contains("GlobalLockSupport")) + assertTrue(localOnlyAspect.reqShieldMap.isEmpty()) + } + class TestBean { @ReqShieldCacheable(cacheName = "TestCacheName", key = "#paramMap['x'] + #paramMap['y']") fun cacheableWithCustomKey(paramMap: Map): Mono = Mono.justOrEmpty(Product("testProduct", "testCategory")) + @ReqShieldCacheable(cacheName = "TestCacheName", key = "#paramMap['x'] + #paramMap['y']") + fun cacheableWithSameKeyOtherMethod(paramMap: Map): Mono = + Mono.justOrEmpty(Product("testProduct", "testCategory")) + @ReqShieldCacheable(cacheName = "TestCacheName") fun cacheableWithDefaultKeyGenerator(paramMap: Map): Mono = Mono.justOrEmpty(Product("testProduct", "testCategory")) @@ -231,8 +328,28 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { fun cacheableWithKeyGenerator(paramMap: Map): Mono = Mono.justOrEmpty(Product("testProduct", "testCategory")) + @ReqShieldCacheable(cacheName = "TestCacheName", key = "#paramMap['x']", keyGenerator = "customGenerator") + fun cacheableWithKeyAndKeyGenerator(paramMap: Map): Mono = + Mono.justOrEmpty(Product("testProduct", "testCategory")) + + @ReqShieldCacheable(cacheName = "TestCacheName", key = "#paramMap['absent']") + fun cacheableWithUnresolvableKey(paramMap: Map): Mono = + Mono.justOrEmpty(Product("testProduct", "testCategory")) + + @ReqShieldCacheable( + cacheName = "TestCacheName", + key = "#paramMap['x']", + nullHandling = NullHandling.ERROR, + ) + fun cacheableWithNullHandlingError(paramMap: Map): Mono = Mono.empty() + + @ReqShieldCacheable(cacheName = "TestCacheName", key = "#paramMap['x']", isLocalLock = false) + fun cacheableWithGlobalLock(paramMap: Map): Mono = Mono.justOrEmpty(Product("testProduct", "testCategory")) + @ReqShieldCacheEvict(cacheName = "TestCacheName") fun evict(paramMap: Map): Mono = Mono.just(true) + + fun notAnnotated(paramMap: Map): Mono = Mono.empty() } class CustomGenerator : KeyGenerator { diff --git a/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/annotation/ReqShieldCacheEvict.kt b/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/annotation/ReqShieldCacheEvict.kt index 44b5635..4604d4d 100644 --- a/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/annotation/ReqShieldCacheEvict.kt +++ b/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/annotation/ReqShieldCacheEvict.kt @@ -25,7 +25,4 @@ annotation class ReqShieldCacheEvict( val cacheName: String, val key: String = "", val keyGenerator: String = "", - val isLocalLock: Boolean = true, - val lockTimeoutMillis: Long = 3000, - val condition: String = "", ) diff --git a/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/annotation/ReqShieldCacheable.kt b/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/annotation/ReqShieldCacheable.kt index edaf2c2..e972dfa 100644 --- a/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/annotation/ReqShieldCacheable.kt +++ b/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/annotation/ReqShieldCacheable.kt @@ -17,6 +17,9 @@ package com.linecorp.cse.reqshield.spring.annotation import com.linecorp.cse.reqshield.config.ReqShieldWorkMode +import com.linecorp.cse.reqshield.support.constant.ConfigValues.DEFAULT_DECISION_FOR_UPDATE +import com.linecorp.cse.reqshield.support.constant.ConfigValues.DEFAULT_LOCK_TIMEOUT_MILLIS +import com.linecorp.cse.reqshield.support.constant.ConfigValues.DEFAULT_TIME_TO_LIVE_MILLIS import com.linecorp.cse.reqshield.support.constant.ConfigValues.MAX_ATTEMPT_GET_CACHE import java.lang.annotation.Inherited @@ -28,9 +31,9 @@ annotation class ReqShieldCacheable( val key: String = "", val keyGenerator: String = "", val isLocalLock: Boolean = true, - val lockTimeoutMillis: Long = 30000, - val decisionForUpdate: Int = 90, + val lockTimeoutMillis: Long = DEFAULT_LOCK_TIMEOUT_MILLIS, + val decisionForUpdate: Int = DEFAULT_DECISION_FOR_UPDATE, val maxAttemptGetCache: Int = MAX_ATTEMPT_GET_CACHE, - val timeToLiveMillis: Long = 10 * 60 * 1000, + val timeToLiveMillis: Long = DEFAULT_TIME_TO_LIVE_MILLIS, val reqShieldWorkMode: ReqShieldWorkMode = ReqShieldWorkMode.CREATE_AND_UPDATE_CACHE, ) diff --git a/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/aspect/ReqShieldAspect.kt b/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/aspect/ReqShieldAspect.kt index 9136337..0a7868c 100644 --- a/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/aspect/ReqShieldAspect.kt +++ b/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/aspect/ReqShieldAspect.kt @@ -20,6 +20,7 @@ import com.linecorp.cse.reqshield.ReqShield 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.GlobalLockSupport import com.linecorp.cse.reqshield.spring.cache.ReqShieldCache import org.aspectj.lang.ProceedingJoinPoint import org.aspectj.lang.annotation.Around @@ -27,6 +28,7 @@ import org.aspectj.lang.annotation.Aspect import org.aspectj.lang.reflect.MethodSignature import org.springframework.beans.factory.BeanFactory import org.springframework.beans.factory.BeanFactoryAware +import org.springframework.beans.factory.annotation.Qualifier import org.springframework.cache.interceptor.KeyGenerator import org.springframework.cache.interceptor.SimpleKeyGenerator import org.springframework.context.expression.MethodBasedEvaluationContext @@ -35,29 +37,36 @@ import org.springframework.core.annotation.AnnotationUtils import org.springframework.expression.EvaluationContext import org.springframework.expression.Expression import org.springframework.expression.spel.standard.SpelExpressionParser -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 +import java.util.concurrent.ScheduledExecutorService @Aspect -@Component class ReqShieldAspect( private val reqShieldCache: ReqShieldCache, + @Qualifier("reqShieldExecutor") private val executor: ScheduledExecutorService, ) : BeanFactoryAware { private lateinit var beanFactory: BeanFactory private val spelParser = SpelExpressionParser() + private val parameterNameDiscoverer = DefaultParameterNameDiscoverer() private val defaultKeyGenerator = SingletonSupplier.of { SimpleKeyGenerator() } + /** Global locking is only available when the cache implementation opts in to it. */ + private val lockSupport = reqShieldCache as? GlobalLockSupport + private val keyGeneratorMap = ConcurrentHashMap() - internal val reqShieldMap = ConcurrentHashMap>() + private val expressionMap = ConcurrentHashMap() + + /** A ReqShield is configured by the annotation alone, so one instance per annotated method is enough. */ + internal val reqShieldMap = ConcurrentHashMap>() @Around("@annotation(com.linecorp.cse.reqshield.spring.annotation.ReqShieldCacheable)") fun aroundReqShieldCacheable(joinPoint: ProceedingJoinPoint): Any? { val annotation = getCacheableAnnotation(joinPoint) - val reqShield = getOrCreateReqShield(joinPoint) val cacheKey = getCacheableCacheKey(joinPoint) + val reqShield = getOrCreateReqShield(joinPoint) return reqShield .getAndSetReqShieldData( @@ -71,19 +80,27 @@ class ReqShieldAspect( fun aroundReqShieldCacheEvict(joinPoint: ProceedingJoinPoint): Any? { val cacheKey = getCacheEvictCacheKey(joinPoint) + // Same default as Spring's @CacheEvict: a method that fails leaves the cache untouched + val result = joinPoint.proceed() reqShieldCache.evict(cacheKey) - return joinPoint.proceed() + return result } private fun getOrCreateReqShield(joinPoint: ProceedingJoinPoint): ReqShield = - reqShieldMap.computeIfAbsent(generateReqShieldKey(joinPoint)) { + reqShieldMap.computeIfAbsent(getTargetMethod(joinPoint)) { createReqShield(joinPoint) } private fun createReqShield(joinPoint: ProceedingJoinPoint): ReqShield { + val method = getTargetMethod(joinPoint) val annotation = getCacheableAnnotation(joinPoint) + require(annotation.isLocalLock || lockSupport != null) { + "isLocalLock = false on ${method.declaringClass.name}.${method.name} " + + "requires the ReqShieldCache bean to implement GlobalLockSupport" + } + val reqShieldConfiguration = ReqShieldConfiguration( setCacheFunction = { key, reqShieldData, timeToLiveMillis -> @@ -93,14 +110,17 @@ class ReqShieldAspect( getCacheFunction = { key -> reqShieldCache.get(key) }, - globalLockFunction = { key, timeToLiveMillis -> - reqShieldCache.globalLock(key, timeToLiveMillis) - }, - globalUnLockFunction = { key -> - reqShieldCache.globalUnLock(key) - }, + globalLockFunction = + lockSupport?.let { support -> + { lockKey, token, timeToLiveMillis -> support.globalLock(lockKey, token, timeToLiveMillis) } + }, + globalUnLockFunction = + lockSupport?.let { support -> + { lockKey, token -> support.globalUnLock(lockKey, token) } + }, isLocalLock = annotation.isLocalLock, lockTimeoutMillis = annotation.lockTimeoutMillis, + executor = executor, decisionForUpdate = annotation.decisionForUpdate, maxAttemptGetCache = annotation.maxAttemptGetCache, reqShieldWorkMode = annotation.reqShieldWorkMode, @@ -121,16 +141,29 @@ class ReqShieldAspect( internal fun getCacheableCacheKey(joinPoint: ProceedingJoinPoint): String { val annotation = getCacheableAnnotation(joinPoint) - validateCacheKey(annotation.key, annotation.keyGenerator) - return getCacheKeyOrDefault(annotation.key, annotation.keyGenerator, joinPoint) + return buildCacheKey(annotation.cacheName, annotation.key, annotation.keyGenerator, joinPoint) } internal fun getCacheEvictCacheKey(joinPoint: ProceedingJoinPoint): String { val annotation = getCacheEvictAnnotation(joinPoint) - validateCacheKey(annotation.key, annotation.keyGenerator) - return getCacheKeyOrDefault(annotation.key, annotation.keyGenerator, joinPoint) + return buildCacheKey(annotation.cacheName, annotation.key, annotation.keyGenerator, joinPoint) + } + + /** + * Namespaces the resolved key with the cache name so that entries (and the locks derived from + * them) of different caches cannot collide. Eviction follows the same rule so that it matches. + */ + private fun buildCacheKey( + cacheName: String, + annotationCacheKey: String, + annotationCacheKeyGenerator: String, + joinPoint: ProceedingJoinPoint, + ): String { + validateCacheKey(annotationCacheKey, annotationCacheKeyGenerator) + + return "$cacheName::${getCacheKeyOrDefault(annotationCacheKey, annotationCacheKeyGenerator, joinPoint)}" } private fun getCacheKeyOrDefault( @@ -140,12 +173,11 @@ class ReqShieldAspect( ): String { val method = getTargetMethod(joinPoint) val context: EvaluationContext = - MethodBasedEvaluationContext(joinPoint.target, method, joinPoint.args, DefaultParameterNameDiscoverer()) + MethodBasedEvaluationContext(joinPoint.target, method, joinPoint.args, parameterNameDiscoverer) val key = if (StringUtils.hasText(annotationCacheKey)) { - val expression: Expression = spelParser.parseExpression(annotationCacheKey) - expression.getValue(context, String::class.java) + getOrParseExpression(annotationCacheKey).getValue(context, String::class.java) } else { val keyGenerator = getOrCreateKeyGenerator(annotationCacheKeyGenerator) keyGenerator.generate(joinPoint.target, method, joinPoint.args).toString() @@ -175,6 +207,11 @@ class ReqShieldAspect( } } + private fun getOrParseExpression(cacheKeyExpression: String): Expression = + expressionMap.computeIfAbsent(cacheKeyExpression) { + spelParser.parseExpression(it) + } + private fun getOrCreateKeyGenerator(keyGeneratorBeanName: String?): KeyGenerator { if (keyGeneratorBeanName.isNullOrBlank()) { return defaultKeyGenerator.obtain() @@ -185,12 +222,6 @@ class ReqShieldAspect( } } - private fun generateReqShieldKey(joinPoint: ProceedingJoinPoint): String { - val method = getTargetMethod(joinPoint) - return "${method.declaringClass.name}.${method.name}-" + - "${getCacheableAnnotation(joinPoint).cacheName}-${getCacheableCacheKey(joinPoint)}" - } - override fun setBeanFactory(beanFactory: BeanFactory) { this.beanFactory = beanFactory } diff --git a/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/cache/GlobalLockSupport.kt b/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/cache/GlobalLockSupport.kt new file mode 100644 index 0000000..3741130 --- /dev/null +++ b/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/cache/GlobalLockSupport.kt @@ -0,0 +1,33 @@ +/* + * 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.spring.cache + +/** Implement alongside [ReqShieldCache] to use `@ReqShieldCacheable(isLocalLock = false)`. */ +interface GlobalLockSupport { + /** Acquire lockKey for the caller identified by token (Redis: `SET lockKey token NX PX timeToLiveMillis`). */ + fun globalLock( + lockKey: String, + token: String, + timeToLiveMillis: Long, + ): Boolean + + /** Release only if the stored value equals token (Redis: compare-and-delete Lua script). */ + fun globalUnLock( + lockKey: String, + token: String, + ): Boolean +} diff --git a/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/cache/ReqShieldCache.kt b/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/cache/ReqShieldCache.kt index c83983b..792db5e 100644 --- a/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/cache/ReqShieldCache.kt +++ b/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/cache/ReqShieldCache.kt @@ -28,32 +28,4 @@ interface ReqShieldCache { ) fun evict(key: String): Boolean? - - /** - * Attempt a global lock on a specific key. - * - * @param key The key to get the lock. - * @param timeToLiveMillis The validity of the lock in milliseconds. - * @return Whether the lock was successfully obtained. Returns `true` by default. - * - * This method provides a default implementation, but if you need a locking mechanism - * You must implement and use your own locking logic. The default implementation is true, and if the value of ReqShieldConfiguration > isLocalLock is false, you will use the function you implemented. - * actual production environments should override this method appropriately to manage locks. - */ - fun globalLock( - key: String, - timeToLiveMillis: Long, - ): Boolean = true - - /** - * Releases the global lock on a specific key. - * - * @param key The key you want to unlock. - * @return Whether the lock release was successful. Returns `true` by default. - * - * This method also provides a default implementation, but if you need a locking mechanism - * You must implement and use your own unlocking logic. The default implementation is true, and if the value of ReqShieldConfiguration > isLocalLock is false, you will use the function you implemented. - * actual production environments should override this method appropriately to manage locks. - */ - fun globalUnLock(key: String): Boolean = true } diff --git a/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/config/LibAutoConfiguration.kt b/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/config/LibAutoConfiguration.kt index 0676041..f2e9f99 100644 --- a/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/config/LibAutoConfiguration.kt +++ b/core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/config/LibAutoConfiguration.kt @@ -16,11 +16,36 @@ package com.linecorp.cse.reqshield.spring.config -import org.springframework.context.annotation.ComponentScan +import com.linecorp.cse.reqshield.spring.aspect.ReqShieldAspect +import org.springframework.context.annotation.Bean import org.springframework.context.annotation.Configuration import org.springframework.context.annotation.EnableAspectJAutoProxy +import org.springframework.context.annotation.Import +import java.util.concurrent.Executors +import java.util.concurrent.ScheduledExecutorService +import java.util.concurrent.atomic.AtomicLong @Configuration @EnableAspectJAutoProxy -@ComponentScan(basePackages = ["com.linecorp.cse"]) -open class LibAutoConfiguration +@Import(ReqShieldAspect::class) +open class LibAutoConfiguration { + /** + * Pool shared by every [com.linecorp.cse.reqshield.ReqShield] the aspect creates, used for the + * asynchronous cache writes and for polling the cache while another request holds the lock. + * + * Spring's inferred destroy method calls [ScheduledExecutorService.shutdown] when the context is + * closed; the threads are daemons anyway so a pending task can never block JVM shutdown. + */ + @Bean + open fun reqShieldExecutor(): ScheduledExecutorService { + val threadCounter = AtomicLong(0) + + return Executors.newScheduledThreadPool( + maxOf(2, Runtime.getRuntime().availableProcessors() * 2), + ) { runnable -> + Thread(runnable, "req-shield-executor-${threadCounter.incrementAndGet()}").apply { + isDaemon = true + } + } + } +} diff --git a/core-spring/src/test/kotlin/aspect/ReqShieldAspectIntegrationTest.kt b/core-spring/src/test/kotlin/aspect/ReqShieldAspectIntegrationTest.kt new file mode 100644 index 0000000..3c8d4b9 --- /dev/null +++ b/core-spring/src/test/kotlin/aspect/ReqShieldAspectIntegrationTest.kt @@ -0,0 +1,153 @@ +/* + * 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 aspect + +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.spring.config.LibAutoConfiguration +import com.linecorp.cse.reqshield.support.model.ReqShieldData +import org.awaitility.Awaitility.await +import org.junit.jupiter.api.AfterEach +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNotNull +import org.junit.jupiter.api.Assertions.assertNull +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.BeforeEach +import org.junit.jupiter.api.Test +import org.springframework.context.annotation.AnnotationConfigApplicationContext +import org.springframework.context.annotation.Bean +import org.springframework.context.annotation.Configuration +import java.time.Duration +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.CountDownLatch +import java.util.concurrent.Executors +import java.util.concurrent.ScheduledExecutorService +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicInteger + +/** + * Wires the aspect the way a consumer does - through [LibAutoConfiguration] only - so it also proves + * that the `@Import` of the aspect and the `reqShieldExecutor` bean work without component scanning. + * + * spring-test is not a test dependency of this module, so the context is driven directly. + */ +class ReqShieldAspectIntegrationTest { + private lateinit var context: AnnotationConfigApplicationContext + private lateinit var service: TestService + private lateinit var cache: InMemoryReqShieldCache + + @BeforeEach + fun setUp() { + context = AnnotationConfigApplicationContext(LibAutoConfiguration::class.java, TestConfig::class.java) + service = context.getBean(TestService::class.java) + cache = context.getBean(InMemoryReqShieldCache::class.java) + } + + @AfterEach + fun tearDown() { + context.close() + } + + @Test + fun executorBeanShouldBeProvidedByTheAutoConfiguration() { + assertNotNull(context.getBean("reqShieldExecutor", ScheduledExecutorService::class.java)) + } + + @Test + fun shouldCollapseDuplicateRequests() { + val attempts = 20 + val executorService = Executors.newFixedThreadPool(attempts) + val startLatch = CountDownLatch(1) + val results = ConcurrentHashMap.newKeySet() + + repeat(attempts) { + executorService.submit { + startLatch.await() + results.add(service.get("dup")) + } + } + startLatch.countDown() + executorService.shutdown() + assertTrue(executorService.awaitTermination(10, TimeUnit.SECONDS)) + + // only the request holding the lock calls the backend, the others wait for its result + assertEquals(1, service.callCount()) + assertEquals(setOf("value-1"), results) + } + + @Test + fun shouldEvictAndRecompute() { + val v1 = service.get("evict") + + // the cache is written asynchronously, so wait for it before evicting + await().atMost(Duration.ofSeconds(5)).until { cache.get("integration::evict") != null } + + service.evict("evict") + assertNull(cache.get("integration::evict")) + + val v2 = service.get("evict") + + assertEquals("value-1", v1) + assertEquals("value-2", v2) + assertEquals(2, service.callCount()) + } + + @Configuration + open class TestConfig { + @Bean + open fun reqShieldCache(): InMemoryReqShieldCache = InMemoryReqShieldCache() + + @Bean + open fun testService(): TestService = TestService() + } + + open class TestService { + private val counter = AtomicInteger(0) + + // read through a method: the bean is a CGLIB proxy whose own fields are never initialized + open fun callCount(): Int = counter.get() + + @ReqShieldCacheable(cacheName = "integration", key = "#key", timeToLiveMillis = 10_000) + open fun get(key: String): String { + // slow enough that concurrent callers reach the lock before the first one finishes + Thread.sleep(100) + return "value-" + counter.incrementAndGet() + } + + @ReqShieldCacheEvict(cacheName = "integration", key = "#key") + open fun evict(key: String) { + // nothing to do: the aspect evicts after this method returns + } + } + + class InMemoryReqShieldCache : ReqShieldCache { + private val store = ConcurrentHashMap>() + + override fun get(key: String): ReqShieldData? = store[key] + + override fun put( + key: String, + value: ReqShieldData, + timeToLiveMillis: Long, + ) { + store[key] = value + } + + override fun evict(key: String): Boolean? = store.remove(key) != null + } +} diff --git a/core-spring/src/test/kotlin/aspect/ReqShieldAspectTest.kt b/core-spring/src/test/kotlin/aspect/ReqShieldAspectTest.kt index 38b246c..07cdb05 100644 --- a/core-spring/src/test/kotlin/aspect/ReqShieldAspectTest.kt +++ b/core-spring/src/test/kotlin/aspect/ReqShieldAspectTest.kt @@ -19,9 +19,11 @@ package aspect import com.linecorp.cse.reqshield.spring.annotation.ReqShieldCacheEvict import com.linecorp.cse.reqshield.spring.annotation.ReqShieldCacheable import com.linecorp.cse.reqshield.spring.aspect.ReqShieldAspect +import com.linecorp.cse.reqshield.spring.cache.GlobalLockSupport 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.constant.ConfigValues.DEFAULT_LOCK_TIMEOUT_MILLIS import com.linecorp.cse.reqshield.support.model.Product import com.linecorp.cse.reqshield.support.model.ReqShieldData import io.mockk.every @@ -30,8 +32,10 @@ import io.mockk.spyk import io.mockk.verify import org.aspectj.lang.ProceedingJoinPoint import org.awaitility.Awaitility +import org.junit.jupiter.api.AfterEach import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.Assertions.assertNotNull +import org.junit.jupiter.api.Assertions.assertThrows import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test @@ -42,14 +46,18 @@ import org.springframework.cache.interceptor.SimpleKeyGenerator import org.springframework.util.ReflectionUtils import java.lang.reflect.Method import java.time.Duration +import java.util.Collections +import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit private val log = LoggerFactory.getLogger(ReqShieldAspectTest::class.java) class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { + private val executor = Executors.newScheduledThreadPool(2) private val reqShieldCache: ReqShieldCache = mockk() private val joinPoint = mockk() - private val reqShieldAspect = spyk(ReqShieldAspect(reqShieldCache)) + private val reqShieldAspect = spyk(ReqShieldAspect(reqShieldCache, executor)) private val targetObject = spyk(TestBean()) private val argument = mapOf("x" to "paramX", "y" to "paramY") @@ -70,18 +78,28 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { reqShieldAspect.setBeanFactory(beanFactory) } + @AfterEach + fun tearDown() { + executor.shutdownNow() + } + + private fun stubTargetMethod( + methodName: String, + aspect: ReqShieldAspect = reqShieldAspect, + ): Method { + val method = ReflectionUtils.findMethod(TestBean::class.java, methodName, Map::class.java)!! + every { aspect.getTargetMethod(joinPoint) } returns method + + return method + } + @Test override fun verifyReqShieldCacheCreation() { // given val reqShieldData = ReqShieldData(methodReturn, 1000) every { reqShieldCache.get(any()) } returns reqShieldData every { joinPoint.proceed() } answers { targetObject.cacheableWithCustomKey(argument) } - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::cacheableWithCustomKey.name, - Map::class.java, - )!! + val method = stubTargetMethod(TestBean::cacheableWithCustomKey.name) // when val result = reqShieldAspect.aroundReqShieldCacheable(joinPoint) @@ -90,9 +108,7 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { // then assertEquals(reqShieldData.value, result) assertTrue(reqShieldAspect.reqShieldMap.size == 1) - val method = reqShieldAspect.getTargetMethod(joinPoint) - val expectedKey = "${method.declaringClass.name}.${method.name}-$cacheName-$spelEvaluatedKey" - assertNotNull(reqShieldAspect.reqShieldMap[expectedKey]) + assertNotNull(reqShieldAspect.reqShieldMap[method]) } } @@ -101,12 +117,7 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { // given every { reqShieldCache.get(any()) } returns ReqShieldData(methodReturn, 1000) every { joinPoint.proceed() } answers { targetObject.cacheableWithCustomKey(argument) } - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::cacheableWithCustomKey.name, - Map::class.java, - )!! + val method = stubTargetMethod(TestBean::cacheableWithCustomKey.name) // when val executorService = Executors.newFixedThreadPool(10) @@ -115,27 +126,40 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { reqShieldAspect.aroundReqShieldCacheable(joinPoint) } } + executorService.shutdown() + assertTrue(executorService.awaitTermination(BaseReqShieldTest.AWAIT_TIMEOUT, TimeUnit.MILLISECONDS)) Awaitility.await().atMost(Duration.ofMillis(BaseReqShieldTest.AWAIT_TIMEOUT)).untilAsserted { // then assertTrue(reqShieldAspect.reqShieldMap.size == 1) - val method = reqShieldAspect.getTargetMethod(joinPoint) - val expectedKey = "${method.declaringClass.name}.${method.name}-$cacheName-$spelEvaluatedKey" - assertNotNull(reqShieldAspect.reqShieldMap[expectedKey]) + assertNotNull(reqShieldAspect.reqShieldMap[method]) } } + @Test + fun eachAnnotatedMethodShouldGetItsOwnReqShield() { + // given + every { reqShieldCache.get(any()) } returns ReqShieldData(methodReturn, 1000) + every { joinPoint.proceed() } answers { targetObject.cacheableWithCustomKey(argument) } + + // when + val customKeyMethod = stubTargetMethod(TestBean::cacheableWithCustomKey.name) + reqShieldAspect.aroundReqShieldCacheable(joinPoint) + val defaultKeyMethod = stubTargetMethod(TestBean::cacheableWithDefaultKeyGenerator.name) + reqShieldAspect.aroundReqShieldCacheable(joinPoint) + + // then + assertEquals(2, reqShieldAspect.reqShieldMap.size) + assertNotNull(reqShieldAspect.reqShieldMap[customKeyMethod]) + assertNotNull(reqShieldAspect.reqShieldMap[defaultKeyMethod]) + } + @Test override fun verifyReqShieldCacheEviction() { // given val reqShieldData = ReqShieldData(methodReturn, 10000) every { reqShieldCache.get(any()) } returns reqShieldData - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::cacheableWithCustomKey.name, - Map::class.java, - )!! + stubTargetMethod(TestBean::cacheableWithCustomKey.name) every { joinPoint.proceed() } answers { targetObject.cacheableWithCustomKey(argument) } val cachedResult = reqShieldAspect.aroundReqShieldCacheable(joinPoint) @@ -144,33 +168,106 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { // Validate cache eviction every { reqShieldCache.evict(any()) } returns true - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::evict.name, - Map::class.java, - )!! + stubTargetMethod(TestBean::evict.name) // when reqShieldAspect.aroundReqShieldCacheEvict(joinPoint) + // then the evicted key is the same namespaced key the cacheable advice used + verify(exactly = 1) { reqShieldCache.evict("$cacheName::$spelEvaluatedKey") } + } + + @Test + fun evictionShouldHappenOnlyAfterTheMethodReturned() { + // given + val invocations = Collections.synchronizedList(mutableListOf()) + every { joinPoint.proceed() } answers { + invocations.add("method") + "methodResult" + } + every { reqShieldCache.evict(any()) } answers { + invocations.add("evict") + true + } + stubTargetMethod(TestBean::evict.name) + + // when + val result = reqShieldAspect.aroundReqShieldCacheEvict(joinPoint) + // then - verify(exactly = 1) { reqShieldCache.evict(any()) } + assertEquals("methodResult", result) + assertEquals(listOf("method", "evict"), invocations) + } + + @Test + fun evictionShouldBeSkippedWhenTheMethodThrows() { + // given + every { joinPoint.proceed() } throws IllegalStateException("method failed") + stubTargetMethod(TestBean::evict.name) + + // when + val exception = + assertThrows(IllegalStateException::class.java) { + reqShieldAspect.aroundReqShieldCacheEvict(joinPoint) + } + + // then + assertEquals("method failed", exception.message) + verify(exactly = 0) { reqShieldCache.evict(any()) } + } + + @Test + fun globalLockShouldRequireACacheImplementingGlobalLockSupport() { + // given a cache that does not implement GlobalLockSupport + stubTargetMethod(TestBean::cacheableWithGlobalLock.name) + + // when + val exception = + assertThrows(IllegalArgumentException::class.java) { + reqShieldAspect.aroundReqShieldCacheable(joinPoint) + } + + // then + assertTrue(exception.message!!.contains("requires the ReqShieldCache bean to implement GlobalLockSupport")) + assertTrue(exception.message!!.contains(TestBean::cacheableWithGlobalLock.name)) + } + + @Test + fun globalLockShouldBeAcquiredAndReleasedWithTheSameToken() { + // given + val globalLockCache = GlobalLockReqShieldCache() + val aspect = spyk(ReqShieldAspect(globalLockCache, executor)) + aspect.setBeanFactory(beanFactory) + every { joinPoint.proceed() } returns methodReturn + stubTargetMethod(TestBean::cacheableWithGlobalLock.name, aspect) + + // when + val result = aspect.aroundReqShieldCacheable(joinPoint) + + // then + assertEquals(methodReturn, result) + Awaitility.await().atMost(Duration.ofMillis(BaseReqShieldTest.AWAIT_TIMEOUT)).untilAsserted { + assertEquals(1, globalLockCache.unLockCalls.size) + } + + val lockCall = globalLockCache.lockCalls.single() + val unLockCall = globalLockCache.unLockCalls.single() + + assertTrue(lockCall.lockKey.contains("$cacheName::$spelEvaluatedKey")) + assertEquals(DEFAULT_LOCK_TIMEOUT_MILLIS, lockCall.timeToLiveMillis) + assertEquals(lockCall.lockKey, unLockCall.lockKey) + assertEquals(lockCall.token, unLockCall.token) + assertNotNull(globalLockCache.get("$cacheName::$spelEvaluatedKey")) } @Test override fun verifyCacheKeyGenerationWithSpEL() { // given - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::cacheableWithCustomKey.name, - Map::class.java, - )!! + stubTargetMethod(TestBean::cacheableWithCustomKey.name) // when, then assertEquals( - spelEvaluatedKey, + "$cacheName::$spelEvaluatedKey", reqShieldAspect.getCacheableCacheKey(joinPoint), ) } @@ -180,37 +277,82 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { // given every { beanFactory.getBean(cacheKeyGenerator, KeyGenerator::class.java) } returns CustomGenerator() - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::cacheableWithKeyGenerator.name, - Map::class.java, - )!! + stubTargetMethod(TestBean::cacheableWithKeyGenerator.name) // when, then assertEquals( - keyGeneratorKey, + "$cacheName::$keyGeneratorKey", + reqShieldAspect.getCacheableCacheKey(joinPoint), + ) + // the generator bean is looked up once and then cached + assertEquals( + "$cacheName::$keyGeneratorKey", reqShieldAspect.getCacheableCacheKey(joinPoint), ) + verify(exactly = 1) { beanFactory.getBean(cacheKeyGenerator, KeyGenerator::class.java) } } @Test override fun verifyCacheKeyGenerationWithDefaultGenerator() { // given - every { reqShieldAspect.getTargetMethod(joinPoint) } returns - ReflectionUtils.findMethod( - TestBean::class.java, - TestBean::cacheableWithDefaultKeyGenerator.name, - Map::class.java, - )!! + stubTargetMethod(TestBean::cacheableWithDefaultKeyGenerator.name) // when, then assertEquals( - SimpleKeyGenerator.generateKey(arrayOf(argument)).toString(), + "$cacheName::${SimpleKeyGenerator.generateKey(arrayOf(argument))}", reqShieldAspect.getCacheableCacheKey(joinPoint), ) } + @Test + fun evictionKeyShouldUseTheSameNamespaceAsTheCacheableKey() { + // given + stubTargetMethod(TestBean::evict.name) + + // when, then + assertEquals("$cacheName::$spelEvaluatedKey", reqShieldAspect.getCacheEvictCacheKey(joinPoint)) + } + + @Test + fun keyAndKeyGeneratorShouldBeMutuallyExclusive() { + // given + stubTargetMethod(TestBean::cacheableWithKeyAndKeyGenerator.name) + + // when, then + val exception = + assertThrows(IllegalArgumentException::class.java) { + reqShieldAspect.getCacheableCacheKey(joinPoint) + } + assertTrue(exception.message!!.contains("mutually exclusive")) + } + + @Test + fun blankResolvedKeyShouldBeRejected() { + // given + stubTargetMethod(TestBean::cacheableWithUnresolvableKey.name) + + // when, then + val exception = + assertThrows(IllegalArgumentException::class.java) { + reqShieldAspect.getCacheableCacheKey(joinPoint) + } + assertTrue(exception.message!!.contains("Null/blank key")) + } + + @Test + fun missingAnnotationsShouldBeRejected() { + // given + stubTargetMethod(TestBean::withoutAnnotation.name) + + // when, then + assertThrows(IllegalArgumentException::class.java) { + reqShieldAspect.getCacheableAnnotation(joinPoint) + } + assertThrows(IllegalArgumentException::class.java) { + reqShieldAspect.getCacheEvictAnnotation(joinPoint) + } + } + class TestBean { @ReqShieldCacheable(cacheName = "TestCacheName", key = "#paramMap['x'] + #paramMap['y']") fun cacheableWithCustomKey(paramMap: Map): String { @@ -230,10 +372,29 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { return "ReturnValue: $paramMap" } + @ReqShieldCacheable( + cacheName = "TestCacheName", + key = "#paramMap['x'] + #paramMap['y']", + keyGenerator = "customGenerator", + ) + fun cacheableWithKeyAndKeyGenerator(paramMap: Map): String = "ReturnValue: $paramMap" + + @ReqShieldCacheable(cacheName = "TestCacheName", key = "#paramMap['unknown']") + fun cacheableWithUnresolvableKey(paramMap: Map): String = "ReturnValue: $paramMap" + + @ReqShieldCacheable( + cacheName = "TestCacheName", + key = "#paramMap['x'] + #paramMap['y']", + isLocalLock = false, + ) + fun cacheableWithGlobalLock(paramMap: Map): String = "ReturnValue: $paramMap" + @ReqShieldCacheEvict(cacheName = "TestCacheName", key = "#paramMap['x'] + #paramMap['y']") fun evict(paramMap: Map) { log.debug("cache eviction") } + + fun withoutAnnotation(paramMap: Map): String = "ReturnValue: $paramMap" } class CustomGenerator : KeyGenerator { @@ -243,4 +404,54 @@ class ReqShieldAspectTest : BaseReqShieldModuleSupportTest { vararg params: Any?, ): Any = "KeyGeneratedByGenerator" } + + /** Cache that opts in to global locking; it records every lock call so the token can be compared. */ + class GlobalLockReqShieldCache : + ReqShieldCache, + GlobalLockSupport { + data class LockCall( + val lockKey: String, + val token: String, + val timeToLiveMillis: Long, + ) + + data class UnLockCall( + val lockKey: String, + val token: String, + ) + + val lockCalls: MutableList = Collections.synchronizedList(mutableListOf()) + val unLockCalls: MutableList = Collections.synchronizedList(mutableListOf()) + + private val store = ConcurrentHashMap>() + + override fun get(key: String): ReqShieldData? = store[key] + + override fun put( + key: String, + value: ReqShieldData, + timeToLiveMillis: Long, + ) { + store[key] = value + } + + override fun evict(key: String): Boolean? = store.remove(key) != null + + override fun globalLock( + lockKey: String, + token: String, + timeToLiveMillis: Long, + ): Boolean { + lockCalls.add(LockCall(lockKey, token, timeToLiveMillis)) + return true + } + + override fun globalUnLock( + lockKey: String, + token: String, + ): Boolean { + unLockCalls.add(UnLockCall(lockKey, token)) + return true + } + } } diff --git a/core/src/main/kotlin/com/linecorp/cse/reqshield/KeyGlobalLock.kt b/core/src/main/kotlin/com/linecorp/cse/reqshield/KeyGlobalLock.kt index d991013..0a50d03 100644 --- a/core/src/main/kotlin/com/linecorp/cse/reqshield/KeyGlobalLock.kt +++ b/core/src/main/kotlin/com/linecorp/cse/reqshield/KeyGlobalLock.kt @@ -16,24 +16,39 @@ package com.linecorp.cse.reqshield +import com.linecorp.cse.reqshield.support.constant.ConfigValues.LOCK_KEY_PREFIX +import java.util.UUID + +/** + * Lock backed by a shared store (for example Redis). + * + * @param globalLockFunction (lockKey, token, ttlMillis) -> acquired. Must store the token only when + * the key is absent and must expire on its own, e.g. `SET key token NX PX ttl`. + * @param globalUnLockFunction (lockKey, token) -> released. Must delete the key only when its value + * still equals the token (compare-and-delete). + */ class KeyGlobalLock( - private val globalLockFunction: (String, Long) -> Boolean, - private val globalUnLockFunction: (String) -> Boolean, + private val globalLockFunction: (String, String, Long) -> Boolean, + private val globalUnLockFunction: (String, String) -> Boolean, private val lockTimeoutMillis: Long, ) : KeyLock { override fun tryLock( key: String, lockType: LockType, - ): Boolean { - val completeKey = "${key}_${lockType.name}" - return globalLockFunction(completeKey, lockTimeoutMillis) + ): String? { + // Tokens are compared across processes, so they must be globally unique + val token = UUID.randomUUID().toString() + return if (globalLockFunction(buildLockKey(key, lockType), token, lockTimeoutMillis)) token else null } override fun unLock( key: String, lockType: LockType, - ): Boolean { - val completeKey = "${key}_${lockType.name}" - return globalUnLockFunction(completeKey) - } + token: String, + ): Boolean = globalUnLockFunction(buildLockKey(key, lockType), token) + + private fun buildLockKey( + key: String, + lockType: LockType, + ): String = "$LOCK_KEY_PREFIX${key}_${lockType.name}" } diff --git a/core/src/main/kotlin/com/linecorp/cse/reqshield/KeyLocalLock.kt b/core/src/main/kotlin/com/linecorp/cse/reqshield/KeyLocalLock.kt index 7abf230..a8974ae 100644 --- a/core/src/main/kotlin/com/linecorp/cse/reqshield/KeyLocalLock.kt +++ b/core/src/main/kotlin/com/linecorp/cse/reqshield/KeyLocalLock.kt @@ -16,6 +16,7 @@ package com.linecorp.cse.reqshield +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 import org.slf4j.LoggerFactory @@ -25,6 +26,8 @@ import java.util.concurrent.ScheduledExecutorService import java.util.concurrent.Semaphore import java.util.concurrent.TimeUnit import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.atomic.AtomicLong +import java.util.concurrent.atomic.AtomicReference private val log = LoggerFactory.getLogger(KeyLocalLock::class.java) @@ -47,16 +50,28 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock { * check vs unLock, or monitor cleanup vs unLock). */ val isHeld: AtomicBoolean = AtomicBoolean(false), + /** + * Ownership token of the current holder, null when the lock is not held. + * Only the holder presenting this token may release the lock, which keeps a holder + * whose lock expired from releasing the lock of the holder that took it over. + * @Volatile ensures visibility across threads when updated inside compute() and read by monitor. + */ + @Volatile var token: String? = null, ) companion object { // Global lockMap shared by all instances - CRITICAL FIX for request collapsing private val lockMap = ConcurrentHashMap() + // Ownership tokens are only compared within this JVM, so a counter is enough (and cheaper than UUID) + private val tokenCounter = AtomicLong(0) + // Single scheduler shared by all instances @Volatile private var sharedScheduler: ScheduledExecutorService? = null + private fun nextToken(): String = "local-${tokenCounter.incrementAndGet()}" + // Thread-safe lazy initialization private fun getOrCreateScheduler(): ScheduledExecutorService { return sharedScheduler ?: synchronized(this) { @@ -96,6 +111,7 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock { // This handles the case where unlock() was missed due to exception. // CAS ensures safe release (no-op if already released). if (lockInfo.isHeld.compareAndSet(true, false)) { + lockInfo.token = null lockInfo.semaphore.release() } null // Atomic removal @@ -144,10 +160,10 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock { override fun tryLock( key: String, lockType: LockType, - ): Boolean { - val completeKey = "${key}_${lockType.name}" + ): String? { + val completeKey = buildLockKey(key, lockType) val now = nowToEpochTime() - val result = AtomicBoolean(false) + val acquiredToken = AtomicReference(null) // Use compute() for atomic lock acquisition. // This ensures mutual exclusion with cleanup - they cannot race on the same key. @@ -158,45 +174,68 @@ class KeyLocalLock(private val lockTimeoutMillis: Long) : KeyLock { // Without CAS, if unLock() executes between isHeld.get() and release(), // both threads would call release(), causing over-release (permits > 1). if (now > existing.expiresAt && existing.isHeld.compareAndSet(true, false)) { + // Clear the token so the previous holder can no longer release this lock + existing.token = null existing.semaphore.release() } // Existing entry: try to acquire semaphore if (existing.semaphore.tryAcquire()) { + val token = nextToken() existing.isHeld.set(true) + existing.token = token existing.expiresAt = now + lockTimeoutMillis - result.set(true) + acquiredToken.set(token) } existing } else { // New entry: create and acquire + val token = nextToken() val newLock = LockInfo(Semaphore(1), now + lockTimeoutMillis) newLock.semaphore.tryAcquire() // Always succeeds for new semaphore newLock.isHeld.set(true) - result.set(true) + newLock.token = token + acquiredToken.set(token) newLock } } - return result.get() + return acquiredToken.get() } override fun unLock( key: String, lockType: LockType, + token: String, ): Boolean { - val completeKey = "${key}_${lockType.name}" - val lockInfo = lockMap[completeKey] ?: return false - - // Use CAS to prevent over-release: only release if we actually hold the lock - return if (lockInfo.isHeld.compareAndSet(true, false)) { - lockInfo.semaphore.release() - true - } else { - log.debug("Attempted to unlock key '{}' that is not held", completeKey) - false + val completeKey = buildLockKey(key, lockType) + val released = AtomicBoolean(false) + + // Release inside compute() so that ownership check and release are atomic with respect to + // tryLock() and the monitor cleanup. Reading the token outside compute() would allow + // "read token == mine -> lock expires and is reacquired by someone else -> release" and + // would therefore release the new owner's lock. + lockMap.compute(completeKey) { _, existing -> + if (existing == null) return@compute null + + if (existing.token == token && existing.isHeld.compareAndSet(true, false)) { + existing.token = null + existing.semaphore.release() + released.set(true) + } + existing // Keep the entry, the monitor removes it once expired + } + + if (!released.get()) { + log.debug("Attempted to unlock key '{}' that is not held or is owned by another holder", completeKey) } + return released.get() } + private fun buildLockKey( + key: String, + lockType: LockType, + ): String = "$LOCK_KEY_PREFIX${key}_${lockType.name}" + fun shutdown() { // Shared scheduler is managed globally, no individual shutdown needed log.debug("KeyLocalLock instance shutdown (scheduler managed globally)") diff --git a/core/src/main/kotlin/com/linecorp/cse/reqshield/KeyLock.kt b/core/src/main/kotlin/com/linecorp/cse/reqshield/KeyLock.kt index c9e56cf..1cbcdfe 100644 --- a/core/src/main/kotlin/com/linecorp/cse/reqshield/KeyLock.kt +++ b/core/src/main/kotlin/com/linecorp/cse/reqshield/KeyLock.kt @@ -16,15 +16,33 @@ package com.linecorp.cse.reqshield +/** + * Lock used to collapse concurrent requests for the same cache key. + * + * Ownership is represented by an opaque token: only the holder that acquired the lock can + * release it, so a holder whose lock already expired and was taken over by someone else + * cannot release the new owner's lock. + */ interface KeyLock { + /** + * Acquires the lock for [key] + [lockType]. + * + * @return an opaque ownership token, or null when another holder owns the lock. + */ fun tryLock( key: String, lockType: LockType, - ): Boolean + ): String? + /** + * Releases the lock for [key] + [lockType] only if [token] matches the current owner. + * + * @return false when the lock is not held or the token does not match the current owner. + */ fun unLock( key: String, lockType: LockType, + token: String, ): Boolean } diff --git a/core/src/main/kotlin/com/linecorp/cse/reqshield/ReqShield.kt b/core/src/main/kotlin/com/linecorp/cse/reqshield/ReqShield.kt index 7ae3136..cdab096 100644 --- a/core/src/main/kotlin/com/linecorp/cse/reqshield/ReqShield.kt +++ b/core/src/main/kotlin/com/linecorp/cse/reqshield/ReqShield.kt @@ -19,6 +19,7 @@ package com.linecorp.cse.reqshield import com.linecorp.cse.reqshield.config.ReqShieldConfiguration import com.linecorp.cse.reqshield.config.ReqShieldWorkMode import com.linecorp.cse.reqshield.support.constant.ConfigValues.GET_CACHE_INTERVAL_MILLIS +import com.linecorp.cse.reqshield.support.constant.ConfigValues.MAX_CONSECUTIVE_GET_CACHE_FAILURES import com.linecorp.cse.reqshield.support.exception.ClientException import com.linecorp.cse.reqshield.support.exception.code.ErrorCode import com.linecorp.cse.reqshield.support.model.ReqShieldData @@ -26,9 +27,12 @@ import com.linecorp.cse.reqshield.support.utils.decideToUpdateCache import org.slf4j.LoggerFactory import java.util.concurrent.Callable import java.util.concurrent.CompletableFuture +import java.util.concurrent.CompletionException +import java.util.concurrent.ExecutionException import java.util.concurrent.ScheduledExecutorService import java.util.concurrent.ScheduledFuture import java.util.concurrent.TimeUnit +import java.util.concurrent.TimeoutException import java.util.concurrent.atomic.AtomicInteger private val log = LoggerFactory.getLogger(ReqShield::class.java) @@ -61,27 +65,21 @@ class ReqShield( timeToLiveMillis: Long, ) { val lockType = LockType.UPDATE + val onlyCreateCache = reqShieldConfig.reqShieldWorkMode == ReqShieldWorkMode.ONLY_CREATE_CACHE - fun executeAsyncTask() { + // ONLY_CREATE_CACHE collapses requests on cache creation only, so the update runs without a lock + val token = if (onlyCreateCache) null else reqShieldConfig.keyLock.tryLock(key, lockType) + + if (onlyCreateCache || token != null) { CompletableFuture.runAsync({ val reqShieldData = buildReqShieldData( - executeCallable({ callable.call() }, true, key, lockType), + executeCallable(callable, key, lockType, token), timeToLiveMillis, ) - setReqShieldData( - reqShieldConfig.setCacheFunction, - key, - reqShieldData, - lockType, - ) + executeSetCacheFunction(reqShieldConfig.setCacheFunction, key, reqShieldData, lockType, token) }, reqShieldConfig.executor) - } - - if (reqShieldConfig.reqShieldWorkMode == ReqShieldWorkMode.ONLY_CREATE_CACHE || - reqShieldConfig.keyLock.tryLock(key, lockType) - ) { - return executeAsyncTask() + .whenComplete { _, e -> if (e != null) logAsyncFailure(key, e) } } } @@ -91,11 +89,13 @@ class ReqShield( timeToLiveMillis: Long, ): ReqShieldData { val lockType = LockType.CREATE + val onlyUpdateCache = reqShieldConfig.reqShieldWorkMode == ReqShieldWorkMode.ONLY_UPDATE_CACHE - return if (reqShieldConfig.reqShieldWorkMode == ReqShieldWorkMode.ONLY_UPDATE_CACHE || - reqShieldConfig.keyLock.tryLock(key, lockType) - ) { - createReqShieldData(key, callable, timeToLiveMillis, lockType) + // ONLY_UPDATE_CACHE collapses requests on cache update only, so the creation runs without a lock + val token = if (onlyUpdateCache) null else reqShieldConfig.keyLock.tryLock(key, lockType) + + return if (onlyUpdateCache || token != null) { + createReqShieldData(key, callable, timeToLiveMillis, lockType, token) } else { handleLockFailure(key, callable, timeToLiveMillis) } @@ -106,32 +106,57 @@ class ReqShield( callable: Callable, timeToLiveMillis: Long, lockType: LockType, + token: String?, ): ReqShieldData { val reqShieldData = buildReqShieldData( - executeCallable({ callable.call() }, true, key, lockType), + executeCallable(callable, key, lockType, token), timeToLiveMillis, ) CompletableFuture.runAsync({ - setReqShieldData(reqShieldConfig.setCacheFunction, key, reqShieldData, lockType) + executeSetCacheFunction(reqShieldConfig.setCacheFunction, key, reqShieldData, lockType, token) }, reqShieldConfig.executor) + .whenComplete { _, e -> if (e != null) logAsyncFailure(key, e) } return reqShieldData } + /** + * Another request holds the lock: poll the cache until that request publishes its result. + * + * The supplier is never called from the polling task - it is called on this thread only after + * the wait gave up, so at most one extra supplier call per waiting request happens. + */ private fun handleLockFailure( key: String, callable: Callable, timeToLiveMillis: Long, ): ReqShieldData { - val future = createFuture() - val counter = createCounter() + val future = CompletableFuture?>() + val scheduled = scheduleTask(reqShieldConfig.executor, future, reqShieldConfig.getCacheFunction, key) - scheduleTask(reqShieldConfig.executor, future, counter, reqShieldConfig.getCacheFunction, callable, key) + // The polling task gives up on its own; this timeout only guards against a task that never runs + val waitTimeoutMillis = + reqShieldConfig.maxAttemptGetCache * GET_CACHE_INTERVAL_MILLIS + GET_CACHE_INTERVAL_MILLIS * 10 - val result = future.get() + val cachedData = + try { + future.get(waitTimeoutMillis, TimeUnit.MILLISECONDS) + } catch (e: TimeoutException) { + log.warn("Timed out waiting for the cache to be created for key '{}', falling back to the supplier", key) + null + } catch (e: InterruptedException) { + Thread.currentThread().interrupt() + throw ClientException(ErrorCode.GET_CACHE_ERROR, cause = e) + } catch (e: ExecutionException) { + val cause = e.cause + throw if (cause is ClientException) cause else ClientException(ErrorCode.GET_CACHE_ERROR, cause = cause) + } finally { + scheduled.cancel(false) + } - return buildReqShieldData(result, timeToLiveMillis) + // No lock was acquired by this request, so there is nothing to release on failure + return cachedData ?: buildReqShieldData(executeCallable(callable, key, null, null), timeToLiveMillis) } private fun buildReqShieldData( @@ -143,66 +168,52 @@ class ReqShield( timeToLiveMillis = timeToLiveMillis, ) - private fun setReqShieldData( - cacheSetter: (String, ReqShieldData, Long) -> Boolean, - key: String, - reqShieldData: ReqShieldData, - lockType: LockType, - ) { - executeSetCacheFunction(cacheSetter, key, reqShieldData, lockType) - } - - private fun createFuture(): CompletableFuture = CompletableFuture() - - private fun createCounter(): AtomicInteger = AtomicInteger(0) - + /** + * Polls the cache on a fixed delay and completes [future] with the cached data once it appears. + * + * The future is completed with null to signal "stop waiting, fall back to the supplier", which + * happens when [ReqShieldConfiguration.maxAttemptGetCache] successful-but-empty reads were made + * or when [MAX_CONSECUTIVE_GET_CACHE_FAILURES] reads failed in a row (the cache looks unavailable). + */ private fun scheduleTask( executor: ScheduledExecutorService, - future: CompletableFuture, - counter: AtomicInteger, + future: CompletableFuture?>, cacheGetter: (String) -> ReqShieldData?, - callable: Callable, key: String, - ) { + ): ScheduledFuture<*> { + val attemptCount = AtomicInteger(0) + val consecutiveFailureCount = AtomicInteger(0) + val scheduled: ScheduledFuture<*> = - executor.scheduleAtFixedRate({ - try { - // Early exit if future is already completed to avoid unnecessary work - if (future.isDone) { - return@scheduleAtFixedRate - } + executor.scheduleWithFixedDelay({ + // Early exit if future is already completed to avoid unnecessary work + if (future.isDone) { + return@scheduleWithFixedDelay + } - val funcResult = executeGetCacheFunction(cacheGetter, key) - if (funcResult != null) { - // Use CAS-like complete to handle race condition safely - // If another thread already completed, this is a no-op - future.complete(funcResult.value) - return@scheduleAtFixedRate + try { + val cachedData = cacheGetter.invoke(key) + if (cachedData != null) { + // complete() is a no-op when another thread already completed the future + future.complete(cachedData) + return@scheduleWithFixedDelay } - // Increment first, then check - ensures atomic decision making - val attempts = counter.incrementAndGet() - if (attempts >= reqShieldConfig.maxAttemptGetCache && !future.isDone) { - // Use complete() which handles concurrent completion safely - // If another thread completed between our check and this call, it's ignored - future.complete(executeCallable({ callable.call() }, false)) + consecutiveFailureCount.set(0) + if (attemptCount.incrementAndGet() >= reqShieldConfig.maxAttemptGetCache) { + future.complete(null) } } catch (e: Exception) { - // Handle exception to prevent scheduleAtFixedRate from stopping - // Fallback to callable to ensure service availability - log.error("Error in scheduled cache getter for key '{}', falling back to callable", key, e) - if (!future.isDone) { - try { - future.complete(executeCallable({ callable.call() }, false)) - } catch (fallbackException: Exception) { - log.error("Fallback callable also failed for key '{}'", key, fallbackException) - future.completeExceptionally(fallbackException) - } + log.warn("Cache read failed while waiting for the cache to be created for key '{}'", key, e) + if (consecutiveFailureCount.incrementAndGet() >= MAX_CONSECUTIVE_GET_CACHE_FAILURES) { + future.complete(null) } } }, GET_CACHE_INTERVAL_MILLIS, GET_CACHE_INTERVAL_MILLIS, TimeUnit.MILLISECONDS) future.whenComplete { _, _ -> scheduled.cancel(false) } + + return scheduled } private fun executeGetCacheFunction( @@ -212,7 +223,7 @@ class ReqShield( runCatching { getFunction.invoke(key) }.getOrElse { - throw ClientException(ErrorCode.GET_CACHE_ERROR, originErrorMessage = it.message) + throw ClientException(ErrorCode.GET_CACHE_ERROR, cause = it) } private fun executeSetCacheFunction( @@ -220,15 +231,16 @@ class ReqShield( key: String, value: ReqShieldData, lockType: LockType, + token: String?, ) { try { setFunction.invoke(key, value, value.timeToLiveMillis) } catch (e: Exception) { - throw ClientException(ErrorCode.SET_CACHE_ERROR, originErrorMessage = e.message) + throw ClientException(ErrorCode.SET_CACHE_ERROR, cause = e) } finally { - if (shouldAttemptUnlock(lockType)) { + if (token != null) { // No retry needed: false means lock already released or expired (not an error) - val unlocked = reqShieldConfig.keyLock.unLock(key, lockType) + val unlocked = reqShieldConfig.keyLock.unLock(key, lockType, token) if (!unlocked) { log.debug("Lock already released or expired for key '{}'", key) } @@ -236,22 +248,30 @@ class ReqShield( } } + /** + * Runs the client supplier, releasing the lock identified by [token] when it fails. + * A null [token] means this request holds no lock, so nothing is released. + */ private fun executeCallable( callable: Callable, - isUnlockWhenException: Boolean, - key: String? = null, - lockType: LockType? = null, + key: String, + lockType: LockType?, + token: String?, ): T? = runCatching { callable.call() }.getOrElse { - if (isUnlockWhenException && key != null && lockType != null) { - reqShieldConfig.keyLock.unLock(key, lockType) + if (token != null && lockType != null) { + reqShieldConfig.keyLock.unLock(key, lockType, token) } - throw ClientException(ErrorCode.SUPPLIER_ERROR, originErrorMessage = it.message) + throw ClientException(ErrorCode.SUPPLIER_ERROR, cause = it) } - private fun shouldAttemptUnlock(lockType: LockType): Boolean = - (lockType == LockType.UPDATE && reqShieldConfig.reqShieldWorkMode != ReqShieldWorkMode.ONLY_CREATE_CACHE) || - (lockType == LockType.CREATE && reqShieldConfig.reqShieldWorkMode != ReqShieldWorkMode.ONLY_UPDATE_CACHE) + private fun logAsyncFailure( + key: String, + throwable: Throwable, + ) { + val cause = if (throwable is CompletionException) throwable.cause ?: throwable else throwable + log.error("Asynchronous cache task failed for key '{}'", key, cause) + } } diff --git a/core/src/main/kotlin/com/linecorp/cse/reqshield/config/ReqShieldConfiguration.kt b/core/src/main/kotlin/com/linecorp/cse/reqshield/config/ReqShieldConfiguration.kt index 58f5a58..4198368 100644 --- a/core/src/main/kotlin/com/linecorp/cse/reqshield/config/ReqShieldConfiguration.kt +++ b/core/src/main/kotlin/com/linecorp/cse/reqshield/config/ReqShieldConfiguration.kt @@ -26,18 +26,30 @@ import com.linecorp.cse.reqshield.support.exception.code.ErrorCode import com.linecorp.cse.reqshield.support.model.ReqShieldData import java.util.concurrent.Executors import java.util.concurrent.ScheduledExecutorService +import java.util.concurrent.atomic.AtomicLong data class ReqShieldConfiguration( val setCacheFunction: (String, ReqShieldData, Long) -> Boolean, val getCacheFunction: (String) -> ReqShieldData?, - val globalLockFunction: ((String, Long) -> Boolean)? = null, - val globalUnLockFunction: ((String) -> Boolean)? = null, + /** + * (lockKey, token, ttlMillis) -> acquired. Required when [isLocalLock] is false. + * The implementation must acquire only when the key is absent and must let the lock expire + * on its own, e.g. `SET key token NX PX ttl`. + */ + val globalLockFunction: ((String, String, Long) -> Boolean)? = null, + /** + * (lockKey, token) -> released. Required when [isLocalLock] is false. + * The implementation must be a compare-and-delete: delete the key only while its value still + * equals the token, so an expired holder cannot release the lock of the next holder. + */ + val globalUnLockFunction: ((String, String) -> Boolean)? = null, val isLocalLock: Boolean = true, val lockTimeoutMillis: Long = DEFAULT_LOCK_TIMEOUT_MILLIS, - val executor: ScheduledExecutorService = - Executors.newScheduledThreadPool( - maxOf(2, Runtime.getRuntime().availableProcessors() * 2), - ), + /** + * Executor used for the asynchronous cache writes and for polling the cache while another + * request holds the lock. Defaults to a single pool shared by every configuration instance. + */ + val executor: ScheduledExecutorService = sharedExecutor, val decisionForUpdate: Int = DEFAULT_DECISION_FOR_UPDATE, val keyLock: KeyLock = if (isLocalLock) { @@ -48,6 +60,24 @@ data class ReqShieldConfiguration( val maxAttemptGetCache: Int = MAX_ATTEMPT_GET_CACHE, val reqShieldWorkMode: ReqShieldWorkMode = ReqShieldWorkMode.CREATE_AND_UPDATE_CACHE, ) { + companion object { + private val executorThreadCounter = AtomicLong(0) + + /** + * Shared by every configuration instance: one ReqShield per cache key must not mean one + * thread pool per cache key. Threads are daemons so the pool never blocks JVM shutdown. + */ + private val sharedExecutor: ScheduledExecutorService by lazy { + Executors.newScheduledThreadPool( + maxOf(2, Runtime.getRuntime().availableProcessors() * 2), + ) { runnable -> + Thread(runnable, "req-shield-executor-${executorThreadCounter.incrementAndGet()}").apply { + isDaemon = true + } + } + } + } + init { if (!isLocalLock) { requireNotNull(globalLockFunction) { diff --git a/core/src/test/kotlin/com/linecorp/cse/reqshield/KeyGlobalLockTest.kt b/core/src/test/kotlin/com/linecorp/cse/reqshield/KeyGlobalLockTest.kt index 7c70f48..f861ca1 100644 --- a/core/src/test/kotlin/com/linecorp/cse/reqshield/KeyGlobalLockTest.kt +++ b/core/src/test/kotlin/com/linecorp/cse/reqshield/KeyGlobalLockTest.kt @@ -20,23 +20,27 @@ import com.linecorp.cse.reqshield.support.BaseKeyLockTest import com.linecorp.cse.reqshield.support.BaseReqShieldTest.Companion.AWAIT_TIMEOUT import com.linecorp.cse.reqshield.support.redis.AbstractRedisTest import io.lettuce.core.RedisClient +import io.lettuce.core.ScriptOutputType +import io.lettuce.core.SetArgs import io.lettuce.core.api.sync.RedisCommands import org.awaitility.Awaitility.await import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertFalse import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test import java.time.Duration import java.util.concurrent.Executors import java.util.concurrent.atomic.AtomicInteger -import kotlin.test.Ignore +import kotlin.test.assertNotNull +import kotlin.test.assertNull class KeyGlobalLockTest : AbstractRedisTest(), BaseKeyLockTest { private lateinit var redisCommands: RedisCommands - private lateinit var globalLockFunc: (String, Long) -> Boolean - private lateinit var globalUnLockFunc: (String) -> Boolean + private lateinit var globalLockFunc: (String, String, Long) -> Boolean + private lateinit var globalUnLockFunc: (String, String) -> Boolean @BeforeEach fun init() { @@ -50,13 +54,19 @@ class KeyGlobalLockTest : // Clean up all keys from previous tests for proper test isolation redisCommands.flushdb() - globalLockFunc = { key, timeToLiveMillis -> - redisCommands.setnx(key, key) + // Store the ownership token only when the key is absent, and let Redis expire the lock + globalLockFunc = { key, token, timeToLiveMillis -> + redisCommands.set(key, token, SetArgs.Builder.nx().px(timeToLiveMillis)) == "OK" } - globalUnLockFunc = { key -> - redisCommands.del(key) - true + // Compare-and-delete: never delete a lock that is already owned by someone else + globalUnLockFunc = { key, token -> + redisCommands.eval( + COMPARE_AND_DELETE_SCRIPT, + ScriptOutputType.INTEGER, + arrayOf(key), + token, + ) == 1L } } @@ -71,7 +81,8 @@ class KeyGlobalLockTest : for (i in 0 until 20) { executorService.submit { - if (keyLock.tryLock(key, lockType)) { + val token = keyLock.tryLock(key, lockType) + if (token != null) { try { println("${Thread.currentThread().name} acquired the lock") lockAcquiredCount.incrementAndGet() @@ -79,7 +90,7 @@ class KeyGlobalLockTest : } catch (e: InterruptedException) { e.printStackTrace() } finally { - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, token) println("${Thread.currentThread().name} released the lock") } } else { @@ -94,7 +105,7 @@ class KeyGlobalLockTest : await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { assertEquals(1, lockAcquiredCount.get()) - assertTrue(keyLock.tryLock(key, lockType)) + assertNotNull(keyLock.tryLock(key, lockType)) } } @@ -109,7 +120,8 @@ class KeyGlobalLockTest : for (i in 0 until 20) { val key = if (i % 2 == 0) "myKey1" else "myKey2" executorService.submit { - if (keyLock.tryLock(key, lockType)) { + val token = keyLock.tryLock(key, lockType) + if (token != null) { try { println("${Thread.currentThread().name} acquired the lock") lockAcquiredCount.incrementAndGet() @@ -117,7 +129,7 @@ class KeyGlobalLockTest : } catch (e: InterruptedException) { e.printStackTrace() } finally { - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, token) println("${Thread.currentThread().name} released the lock") } } else { @@ -132,16 +144,35 @@ class KeyGlobalLockTest : await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { assertTrue(lockAcquiredCount.get() <= 4) - assertTrue(keyLock.tryLock("myKey1", lockType)) - assertTrue(keyLock.tryLock("myKey2", lockType)) + assertNotNull(keyLock.tryLock("myKey1", lockType)) + assertNotNull(keyLock.tryLock("myKey2", lockType)) } } @Test - @Ignore override fun testLockExpiration() { - // Global locks do not have an expiration + val keyLock = KeyGlobalLock(globalLockFunc, globalUnLockFunc, lockTimeoutMillis) + val key = "expirationKey" + val lockType = LockType.CREATE + + // Given: the lock is held and carries the TTL handed to the lock function + val tokenOfA = assertNotNull(keyLock.tryLock(key, lockType)) + assertNull(keyLock.tryLock(key, lockType), "Lock must not be acquired twice while held") + assertFalse(keyLock.unLock(key, lockType, "foreign-token"), "A foreign token must not release the lock") + + // When: the TTL passes without an explicit unlock + Thread.sleep(lockTimeoutMillis + 500L) + + // Then: the expired lock can be taken over, and the previous holder cannot release it + val tokenOfB = assertNotNull(keyLock.tryLock(key, lockType), "Expired lock should be acquirable again") + assertFalse(keyLock.unLock(key, lockType, tokenOfA), "A stale token must not release the new holder's lock") + assertTrue(keyLock.unLock(key, lockType, tokenOfB), "The current holder can release the lock") } private fun doWork() = Thread.sleep(1000) + + companion object { + private const val COMPARE_AND_DELETE_SCRIPT = + "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end" + } } diff --git a/core/src/test/kotlin/com/linecorp/cse/reqshield/KeyLocalLockShutdownTest.kt b/core/src/test/kotlin/com/linecorp/cse/reqshield/KeyLocalLockShutdownTest.kt index 0475beb..baa42d8 100644 --- a/core/src/test/kotlin/com/linecorp/cse/reqshield/KeyLocalLockShutdownTest.kt +++ b/core/src/test/kotlin/com/linecorp/cse/reqshield/KeyLocalLockShutdownTest.kt @@ -21,6 +21,8 @@ import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.Test import java.time.Duration +import kotlin.test.assertNotNull +import kotlin.test.assertNull class KeyLocalLockShutdownTest { @Test @@ -31,8 +33,8 @@ class KeyLocalLockShutdownTest { val key1 = "testKey1" val key2 = "testKey2" - keyLock.tryLock(key1, LockType.CREATE) - keyLock.tryLock(key2, LockType.CREATE) + val token1 = assertNotNull(keyLock.tryLock(key1, LockType.CREATE)) + val token2 = assertNotNull(keyLock.tryLock(key2, LockType.CREATE)) // Call shutdown keyLock.shutdown() @@ -40,12 +42,12 @@ class KeyLocalLockShutdownTest { // Test behavioral verification instead of internal state // After shutdown, the instance should still function normally // but should be deregistered from shared monitoring - assertTrue(keyLock.tryLock("newKey", LockType.CREATE), "Lock should still work after shutdown") - keyLock.unLock("newKey", LockType.CREATE) + val newToken = assertNotNull(keyLock.tryLock("newKey", LockType.CREATE), "Lock should still work after shutdown") + keyLock.unLock("newKey", LockType.CREATE, newToken) // Cleanup existing locks - keyLock.unLock(key1, LockType.CREATE) - keyLock.unLock(key2, LockType.CREATE) + keyLock.unLock(key1, LockType.CREATE, token1) + keyLock.unLock(key2, LockType.CREATE, token2) } @Test @@ -57,17 +59,17 @@ class KeyLocalLockShutdownTest { val lockType = LockType.CREATE // Acquire lock - assertTrue(keyLock.tryLock(key, lockType), "Should acquire lock initially") + assertNotNull(keyLock.tryLock(key, lockType), "Should acquire lock initially") // Should fail to acquire same lock again (already acquired) - assertTrue(!keyLock.tryLock(key, lockType), "Should not acquire same lock again") + assertNull(keyLock.tryLock(key, lockType), "Should not acquire same lock again") // Should be automatically cleaned up after lock timeout + cleanup interval // Cleanup interval is 1000ms, so we wait for lock timeout + cleanup interval + buffer await().atMost(Duration.ofMillis(shortTimeout + 1000L + 500L)).untilAsserted { // Should be able to acquire new lock after cleanup - assertTrue(keyLock.tryLock(key, lockType), "Should be able to acquire lock after timeout and cleanup") - keyLock.unLock(key, lockType) // Cleanup + val token = assertNotNull(keyLock.tryLock(key, lockType), "Should be able to acquire lock after timeout and cleanup") + keyLock.unLock(key, lockType, token) // Cleanup } keyLock.shutdown() @@ -89,9 +91,10 @@ class KeyLocalLockShutdownTest { // Should be able to acquire locks again after all keys are cleaned up var availableCount = 0 for (i in 1..5) { - if (keyLock.tryLock("key$i", LockType.CREATE)) { + val token = keyLock.tryLock("key$i", LockType.CREATE) + if (token != null) { availableCount++ - keyLock.unLock("key$i", LockType.CREATE) + keyLock.unLock("key$i", LockType.CREATE, token) } } assertEquals(5, availableCount, "All expired locks should be cleaned up") @@ -126,8 +129,9 @@ class KeyLocalLockShutdownTest { // Test behavioral verification: shutdown should complete normally // and the keyLock should still function correctly after interrupted shutdown - assertTrue(keyLock.tryLock("interruptTestKey", LockType.CREATE), "Lock should work after interrupted shutdown") - keyLock.unLock("interruptTestKey", LockType.CREATE) + val token = + assertNotNull(keyLock.tryLock("interruptTestKey", LockType.CREATE), "Lock should work after interrupted shutdown") + keyLock.unLock("interruptTestKey", LockType.CREATE, token) // Clear interrupt state Thread.interrupted() diff --git a/core/src/test/kotlin/com/linecorp/cse/reqshield/KeyLocalLockTest.kt b/core/src/test/kotlin/com/linecorp/cse/reqshield/KeyLocalLockTest.kt index 9f835c9..c9ece1c 100644 --- a/core/src/test/kotlin/com/linecorp/cse/reqshield/KeyLocalLockTest.kt +++ b/core/src/test/kotlin/com/linecorp/cse/reqshield/KeyLocalLockTest.kt @@ -29,6 +29,8 @@ import java.util.concurrent.CountDownLatch import java.util.concurrent.Executors import java.util.concurrent.TimeUnit import java.util.concurrent.atomic.AtomicInteger +import kotlin.test.assertNotNull +import kotlin.test.assertNull class KeyLocalLockTest : BaseKeyLockTest { @Test @@ -42,7 +44,8 @@ class KeyLocalLockTest : BaseKeyLockTest { for (i in 0 until 20) { executorService.submit { - if (keyLock.tryLock(key, lockType)) { + val token = keyLock.tryLock(key, lockType) + if (token != null) { try { println("${Thread.currentThread().name} acquired the lock") lockAcquiredCount.incrementAndGet() @@ -50,7 +53,7 @@ class KeyLocalLockTest : BaseKeyLockTest { } catch (e: InterruptedException) { e.printStackTrace() } finally { - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, token) println("${Thread.currentThread().name} released the lock") } } else { @@ -65,7 +68,7 @@ class KeyLocalLockTest : BaseKeyLockTest { await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { assertEquals(1, lockAcquiredCount.get()) - assertTrue(keyLock.tryLock(key, lockType)) + assertNotNull(keyLock.tryLock(key, lockType)) } } @@ -80,7 +83,8 @@ class KeyLocalLockTest : BaseKeyLockTest { for (i in 0 until 20) { val key = if (i % 2 == 0) "myKey1" else "myKey2" executorService.submit { - if (keyLock.tryLock(key, lockType)) { + val token = keyLock.tryLock(key, lockType) + if (token != null) { try { println("${Thread.currentThread().name} acquired the lock") lockAcquiredCount.incrementAndGet() @@ -88,7 +92,7 @@ class KeyLocalLockTest : BaseKeyLockTest { } catch (e: InterruptedException) { e.printStackTrace() } finally { - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, token) println("${Thread.currentThread().name} released the lock") } } else { @@ -103,8 +107,8 @@ class KeyLocalLockTest : BaseKeyLockTest { await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { assertTrue(lockAcquiredCount.get() <= 4) - assertTrue(keyLock.tryLock("myKey1", lockType)) - assertTrue(keyLock.tryLock("myKey2", lockType)) + assertNotNull(keyLock.tryLock("myKey1", lockType)) + assertNotNull(keyLock.tryLock("myKey2", lockType)) } } @@ -114,7 +118,7 @@ class KeyLocalLockTest : BaseKeyLockTest { val key = "myKey" val lockType = LockType.CREATE - assertTrue(keyLock.tryLock(key, lockType)) + assertNotNull(keyLock.tryLock(key, lockType)) // Wait for lock timeout + cleanup interval + buffer // lockTimeoutMillis = 3000ms, cleanup interval = 1000ms @@ -122,14 +126,36 @@ class KeyLocalLockTest : BaseKeyLockTest { val executorService = Executors.newSingleThreadExecutor() val future = - executorService.submit { + executorService.submit { keyLock.tryLock(key, lockType) } - await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { - assertTrue(future.get()) - } - assertTrue(keyLock.unLock(key, lockType)) + val token = assertNotNull(future.get()) + assertTrue(keyLock.unLock(key, lockType, token)) + } + + @Test + fun `should not let an expired holder release the lock of the new holder`() { + // Given: a lock that expires quickly, held by holder A + val lockTimeout = 200L + val keyLock = KeyLocalLock(lockTimeout) + val key = "stale-token-key" + val lockType = LockType.CREATE + val tokenOfA = assertNotNull(keyLock.tryLock(key, lockType)) + + // When: A's lock expires and holder B takes it over + Thread.sleep(lockTimeout + 100L) + val tokenOfB = assertNotNull(keyLock.tryLock(key, lockType), "B should take over the expired lock") + assertTrue(tokenOfA != tokenOfB, "Each acquisition must get its own token") + + // Then: A's stale token must not release B's lock + assertFalse(keyLock.unLock(key, lockType, tokenOfA), "A stale token must not release the new holder's lock") + assertNull(keyLock.tryLock(key, lockType), "B must still hold the lock") + + // And: B can release its own lock + assertTrue(keyLock.unLock(key, lockType, tokenOfB)) + + keyLock.shutdown() } @Test @@ -169,7 +195,7 @@ class KeyLocalLockTest : BaseKeyLockTest { val lockType = LockType.CREATE // When: Acquire lock and wait for expiration + cleanup interval - assertTrue(keyLock.tryLock(key, lockType)) + assertNotNull(keyLock.tryLock(key, lockType)) // Then: Cleanup should work efficiently // Wait for: lockTimeout + cleanup interval (1000ms) + buffer @@ -177,8 +203,8 @@ class KeyLocalLockTest : BaseKeyLockTest { // Expired locks should be cleaned up, allowing new lock acquisition await().atMost(Duration.ofSeconds(3)).untilAsserted { - assertTrue(keyLock.tryLock(key, lockType)) - keyLock.unLock(key, lockType) + val token = assertNotNull(keyLock.tryLock(key, lockType)) + keyLock.unLock(key, lockType, token) } keyLock.shutdown() @@ -193,16 +219,16 @@ class KeyLocalLockTest : BaseKeyLockTest { val lockType = LockType.CREATE // When - Instance1 acquires lock - val lock1Result = instance1.tryLock(key, lockType) + val token1 = instance1.tryLock(key, lockType) // Then - Instance2 should not be able to acquire the same lock - val lock2Result = instance2.tryLock(key, lockType) + val token2 = instance2.tryLock(key, lockType) - assertTrue(lock1Result) - assertTrue(!lock2Result, "Instance2 should not acquire lock held by Instance1") + val acquiredToken = assertNotNull(token1) + assertNull(token2, "Instance2 should not acquire lock held by Instance1") // Cleanup - instance1.unLock(key, lockType) + instance1.unLock(key, lockType, acquiredToken) instance1.shutdown() instance2.shutdown() } @@ -230,10 +256,11 @@ class KeyLocalLockTest : BaseKeyLockTest { else -> instance3 } attemptCount.incrementAndGet() - if (instance.tryLock(key, lockType)) { + val token = instance.tryLock(key, lockType) + if (token != null) { successCount.incrementAndGet() Thread.sleep(50) // Hold lock briefly - instance.unLock(key, lockType) + instance.unLock(key, lockType, token) } latch.countDown() } @@ -260,21 +287,39 @@ class KeyLocalLockTest : BaseKeyLockTest { val key = "unlock-shared-key" val lockType = LockType.CREATE - // When - Instance1 acquires lock, Instance2 can also unlock (global lockMap shared) - assertTrue(instance1.tryLock(key, lockType)) - // Instance2 can unlock because isHeld state is global - assertTrue(instance2.unLock(key, lockType), "Global unlock should succeed from any instance") + // When - Instance1 acquires the lock and hands its token to Instance2 + val token = assertNotNull(instance1.tryLock(key, lockType)) + + // Then - Instance2 can release it because the lockMap is global and the token matches + assertTrue(instance2.unLock(key, lockType, token), "Unlock with the owning token should succeed from any instance") - // Then - New lock acquisition should succeed - val newLockResult = instance2.tryLock(key, lockType) - assertTrue(newLockResult, "Should be able to acquire lock after global unlock") + // And - New lock acquisition should succeed + val newToken = assertNotNull(instance2.tryLock(key, lockType), "Should be able to acquire lock after global unlock") // Cleanup - instance2.unLock(key, lockType) + instance2.unLock(key, lockType, newToken) instance1.shutdown() instance2.shutdown() } + @Test + fun `unLock with a foreign token returns false and keeps the lock held`() { + // Given + val keyLock = KeyLocalLock(lockTimeoutMillis) + val key = "foreign-token-key" + val lockType = LockType.CREATE + val token = assertNotNull(keyLock.tryLock(key, lockType)) + + // When - Someone that does not own the lock tries to release it + assertFalse(keyLock.unLock(key, lockType, "someone-else-token"), "A foreign token must not release the lock") + + // Then - The lock is still held + assertNull(keyLock.tryLock(key, lockType), "Lock must still be held after a foreign unlock attempt") + assertTrue(keyLock.unLock(key, lockType, token), "The owner can still release the lock") + + keyLock.shutdown() + } + @Test fun `should not over-release semaphore on multiple unlock calls`() { // Given @@ -283,20 +328,20 @@ class KeyLocalLockTest : BaseKeyLockTest { val lockType = LockType.CREATE // When - Acquire lock - assertTrue(keyLock.tryLock(key, lockType)) + val token = assertNotNull(keyLock.tryLock(key, lockType)) // Then - First unlock should succeed - assertTrue(keyLock.unLock(key, lockType), "First unlock should succeed") + assertTrue(keyLock.unLock(key, lockType, token), "First unlock should succeed") // Second unlock should return false (lock not held) - assertFalse(keyLock.unLock(key, lockType), "Second unlock should fail (over-release prevention)") + assertFalse(keyLock.unLock(key, lockType, token), "Second unlock should fail (over-release prevention)") // Verify semaphore is not over-released: can acquire once, not twice - assertTrue(keyLock.tryLock(key, lockType), "Should acquire lock after proper unlock") - assertFalse(keyLock.tryLock(key, lockType), "Should not acquire lock twice (semaphore intact)") + val newToken = assertNotNull(keyLock.tryLock(key, lockType), "Should acquire lock after proper unlock") + assertNull(keyLock.tryLock(key, lockType), "Should not acquire lock twice (semaphore intact)") // Cleanup - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, newToken) keyLock.shutdown() } @@ -308,19 +353,22 @@ class KeyLocalLockTest : BaseKeyLockTest { val lockType = LockType.CREATE val executor = Executors.newFixedThreadPool(10) val successfulAcquisitions = AtomicInteger(0) + val acquiredTokens = mutableListOf() val latch = CountDownLatch(10) // Simulate over-release attempt - assertTrue(keyLock.tryLock(key, lockType)) - keyLock.unLock(key, lockType) + val token = assertNotNull(keyLock.tryLock(key, lockType)) + keyLock.unLock(key, lockType, token) // Multiple unlock attempts should all return false (not over-release) - repeat(5) { assertFalse(keyLock.unLock(key, lockType)) } + repeat(5) { assertFalse(keyLock.unLock(key, lockType, token)) } // When - Try to acquire lock concurrently repeat(10) { executor.submit { - if (keyLock.tryLock(key, lockType)) { + val acquired = keyLock.tryLock(key, lockType) + if (acquired != null) { successfulAcquisitions.incrementAndGet() + synchronized(acquiredTokens) { acquiredTokens.add(acquired) } } latch.countDown() } @@ -333,7 +381,7 @@ class KeyLocalLockTest : BaseKeyLockTest { assertEquals(1, successfulAcquisitions.get(), "Only one thread should acquire the lock") // Cleanup - keyLock.unLock(key, lockType) + acquiredTokens.forEach { keyLock.unLock(key, lockType, it) } keyLock.shutdown() } @@ -353,10 +401,11 @@ class KeyLocalLockTest : BaseKeyLockTest { keys.forEach { key -> executor.submit { try { - if (instance.tryLock(key, LockType.CREATE)) { + val token = instance.tryLock(key, LockType.CREATE) + if (token != null) { operations.incrementAndGet() Thread.sleep(10) // Brief work simulation - instance.unLock(key, LockType.CREATE) + instance.unLock(key, LockType.CREATE, token) } } catch (e: Exception) { errors.incrementAndGet() @@ -408,25 +457,26 @@ class KeyLocalLockTest : BaseKeyLockTest { // This validates that compute() atomicity prevents race conditions repeat(10) { // Acquire lock - assertTrue(keyLock.tryLock(key, lockType), "Should acquire lock") + val token = assertNotNull(keyLock.tryLock(key, lockType), "Should acquire lock") // Hold until expiration Thread.sleep(lockTimeout + 200) - // Release - keyLock.unLock(key, lockType) + // Release (may already have been force-released by the monitor) + keyLock.unLock(key, lockType, token) // Immediately reacquire - compute() ensures this doesn't race with cleanup - val reacquired = keyLock.tryLock(key, lockType) - if (reacquired) { + val reacquiredToken = keyLock.tryLock(key, lockType) + if (reacquiredToken != null) { // Verify lock exclusivity - second acquire must fail - if (keyLock.tryLock(key, lockType)) { + val secondToken = keyLock.tryLock(key, lockType) + if (secondToken != null) { // This indicates lock was incorrectly removed during acquisition errors.incrementAndGet() - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, secondToken) } successfulCycles.incrementAndGet() - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, reacquiredToken) } } @@ -457,24 +507,26 @@ class KeyLocalLockTest : BaseKeyLockTest { executor.submit { try { // Acquire lock - if (keyLock.tryLock(key, lockType)) { + val token = keyLock.tryLock(key, lockType) + if (token != null) { // Hold past expiration to trigger cleanup consideration Thread.sleep(lockTimeout + 200) // Release and immediately re-acquire - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, token) // With compute(), this operation is atomic with respect to cleanup - val reacquired = keyLock.tryLock(key, lockType) - if (reacquired) { + val reacquiredToken = keyLock.tryLock(key, lockType) + if (reacquiredToken != null) { // Verify lock exclusivity - if (keyLock.tryLock(key, lockType)) { + val secondToken = keyLock.tryLock(key, lockType) + if (secondToken != null) { // This should never happen - compute() ensures atomicity lockRemovedWhileHeld.incrementAndGet() - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, secondToken) } successfulCycles.incrementAndGet() - keyLock.unLock(key, lockType) + keyLock.unLock(key, lockType, reacquiredToken) } } } finally { @@ -507,7 +559,7 @@ class KeyLocalLockTest : BaseKeyLockTest { val lockType = LockType.CREATE // When: Acquire lock but never unlock (simulating exception scenario) - assertTrue(keyLock.tryLock(key, lockType), "Should acquire lock") + assertNotNull(keyLock.tryLock(key, lockType), "Should acquire lock") // DO NOT call unlock - simulating exception scenario // Wait for expiration + cleanup interval + buffer @@ -516,16 +568,18 @@ class KeyLocalLockTest : BaseKeyLockTest { // Then: Cleanup should have force-released and removed the expired lock // A new lock acquisition should succeed + var newToken: String? = null await().atMost(Duration.ofSeconds(3)).untilAsserted { - assertTrue( - keyLock.tryLock(key, lockType), - "Should acquire lock after cleanup removed expired held lock", - ) + newToken = + assertNotNull( + keyLock.tryLock(key, lockType), + "Should acquire lock after cleanup removed expired held lock", + ) } // Verify lock is working normally - assertFalse(keyLock.tryLock(key, lockType), "Second acquire should fail (lock is held)") - assertTrue(keyLock.unLock(key, lockType), "Unlock should succeed") + assertNull(keyLock.tryLock(key, lockType), "Second acquire should fail (lock is held)") + assertTrue(keyLock.unLock(key, lockType, newToken!!), "Unlock should succeed") keyLock.shutdown() } @@ -543,7 +597,7 @@ class KeyLocalLockTest : BaseKeyLockTest { // When: Acquire many locks but never unlock them repeat(keyCount) { i -> - assertTrue(keyLock.tryLock("leak-test-$i", lockType), "Should acquire lock $i") + assertNotNull(keyLock.tryLock("leak-test-$i", lockType), "Should acquire lock $i") } // Wait for all locks to expire and be cleaned up @@ -553,11 +607,12 @@ class KeyLocalLockTest : BaseKeyLockTest { // Then: All expired locks should be cleaned up, allowing reacquisition await().atMost(Duration.ofSeconds(5)).untilAsserted { repeat(keyCount) { i -> - assertTrue( - keyLock.tryLock("leak-test-$i", lockType), - "Should acquire lock $i after cleanup", - ) - keyLock.unLock("leak-test-$i", lockType) + val token = + assertNotNull( + keyLock.tryLock("leak-test-$i", lockType), + "Should acquire lock $i after cleanup", + ) + keyLock.unLock("leak-test-$i", lockType, token) } } diff --git a/core/src/test/kotlin/com/linecorp/cse/reqshield/ReqShieldTest.kt b/core/src/test/kotlin/com/linecorp/cse/reqshield/ReqShieldTest.kt index 54d994b..ed3507b 100644 --- a/core/src/test/kotlin/com/linecorp/cse/reqshield/ReqShieldTest.kt +++ b/core/src/test/kotlin/com/linecorp/cse/reqshield/ReqShieldTest.kt @@ -20,17 +20,16 @@ import com.linecorp.cse.reqshield.config.ReqShieldConfiguration import com.linecorp.cse.reqshield.config.ReqShieldWorkMode import com.linecorp.cse.reqshield.support.BaseReqShieldTest import com.linecorp.cse.reqshield.support.BaseReqShieldTest.Companion.AWAIT_TIMEOUT +import com.linecorp.cse.reqshield.support.constant.ConfigValues.GET_CACHE_INTERVAL_MILLIS +import com.linecorp.cse.reqshield.support.constant.ConfigValues.MAX_CONSECUTIVE_GET_CACHE_FAILURES 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 io.mockk.every import io.mockk.mockk -import io.mockk.mockkStatic -import io.mockk.unmockkStatic import io.mockk.verify import org.awaitility.Awaitility.await -import org.junit.jupiter.api.AfterEach import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.BeforeEach @@ -39,8 +38,11 @@ import org.junit.jupiter.api.assertThrows import java.lang.reflect.InvocationTargetException import java.lang.reflect.Method import java.time.Duration -import java.time.LocalDateTime import java.util.concurrent.Callable +import java.util.concurrent.CountDownLatch +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicInteger import kotlin.test.assertFailsWith import kotlin.test.assertNotNull import kotlin.test.assertNull @@ -59,18 +61,20 @@ class ReqShieldTest : BaseReqShieldTest { private val oldValue = Product("oldTestValue", "oldTestName") private val value = Product("testId", "testName") private val callable: Callable = mockk() + private val createToken = "create-token" + private val updateToken = "update-token" private var timeToLiveMillis: Long = 10000 - private lateinit var globalLockFunc: (String, Long) -> Boolean - private lateinit var globalUnLockFunc: (String) -> Boolean + private lateinit var globalLockFunc: (String, String, Long) -> Boolean + private lateinit var globalUnLockFunc: (String, String) -> Boolean @BeforeEach fun setup() { cacheSetter = mockk<(String, ReqShieldData, Long) -> Boolean>() cacheGetter = mockk<(String) -> ReqShieldData?>() - globalLockFunc = mockk<(String, Long) -> Boolean>() - globalUnLockFunc = mockk<(String) -> Boolean>() + globalLockFunc = mockk<(String, String, Long) -> Boolean>() + globalUnLockFunc = mockk<(String, String) -> Boolean>() keyLock = mockk() keyGlobalLock = KeyGlobalLock(globalLockFunc, globalUnLockFunc, 3000) @@ -117,22 +121,35 @@ class ReqShieldTest : BaseReqShieldTest { keyLock = keyGlobalLock, ), ) - - mockkStatic(LocalDateTime::class) - every { LocalDateTime.now() } returns LocalDateTime.of(2023, 11, 13, 12, 0, 0, 0) } - @AfterEach - fun tearDown() { - unmockkStatic(LocalDateTime::class) - } + /** + * Builds a cache entry that is old enough to be an update target: + * 90% of its TTL has already passed while decisionForUpdate defaults to 80%. + */ + private fun updateTargetData(cachedValue: Product?): ReqShieldData = + ReqShieldData( + value = cachedValue, + status = ReqShieldData.Status.NEW, + createdAt = System.currentTimeMillis() - (timeToLiveMillis * 0.9).toLong(), + timeToLiveMillis = timeToLiveMillis, + ) + + /** Builds a freshly created cache entry, which must not be an update target. */ + private fun freshData(cachedValue: Product?): ReqShieldData = + ReqShieldData( + value = cachedValue, + status = ReqShieldData.Status.NEW, + createdAt = System.currentTimeMillis(), + timeToLiveMillis = timeToLiveMillis, + ) @Test override fun testSetMethodCacheNotExistsAndLocalLockAcquired() { every { cacheGetter.invoke(key) } returns null every { cacheSetter.invoke(key, any(), any()) } returns true - every { keyLock.tryLock(key, LockType.CREATE) } returns true - every { keyLock.unLock(key, LockType.CREATE) } returns true + every { keyLock.tryLock(key, LockType.CREATE) } returns createToken + every { keyLock.unLock(key, LockType.CREATE, createToken) } returns true val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) @@ -141,7 +158,7 @@ class ReqShieldTest : BaseReqShieldTest { verify { cacheGetter.invoke(key) } verify { cacheSetter.invoke(key, result, timeToLiveMillis) } verify { keyLock.tryLock(key, LockType.CREATE) } - verify { keyLock.unLock(key, LockType.CREATE) } + verify { keyLock.unLock(key, LockType.CREATE, createToken) } verify { callable.call() } } } @@ -158,7 +175,7 @@ class ReqShieldTest : BaseReqShieldTest { verify { cacheGetter.invoke(key) } verify { cacheSetter.invoke(key, result, timeToLiveMillis) } verify(inverse = true) { keyLock.tryLock(key, LockType.CREATE) } - verify(inverse = true) { keyLock.unLock(key, LockType.CREATE) } + verify(inverse = true) { keyLock.unLock(key, LockType.CREATE, any()) } verify { callable.call() } } } @@ -168,8 +185,8 @@ class ReqShieldTest : BaseReqShieldTest { every { cacheGetter.invoke(key) } returns null every { cacheSetter.invoke(key, any(), any()) } returns true - every { globalLockFunc(any(), any()) } returns true - every { globalUnLockFunc(any()) } returns true + every { globalLockFunc(any(), any(), any()) } returns true + every { globalUnLockFunc(any(), any()) } returns true val result = reqShieldForGlobalLock.getAndSetReqShieldData(key, callable, timeToLiveMillis) @@ -177,10 +194,8 @@ class ReqShieldTest : BaseReqShieldTest { assertNotNull(result) verify { cacheGetter.invoke(key) } verify { cacheSetter.invoke(key, result, timeToLiveMillis) } - verify { globalLockFunc(any(), any()) } - verify { globalUnLockFunc(any()) } - verify { keyGlobalLock.tryLock(key, LockType.CREATE) } - verify { keyGlobalLock.unLock(key, LockType.CREATE) } + verify { globalLockFunc(any(), any(), any()) } + verify { globalUnLockFunc(any(), any()) } verify { callable.call() } } } @@ -207,8 +222,8 @@ class ReqShieldTest : BaseReqShieldTest { override fun testSetMethodCacheNotExistsAndLocalLockAcquiredAndCallableReturnNull() { every { cacheGetter.invoke(key) } returns null every { cacheSetter.invoke(key, any(), any()) } returns true - every { keyLock.tryLock(key, LockType.CREATE) } returns true - every { keyLock.unLock(key, LockType.CREATE) } returns true + every { keyLock.tryLock(key, LockType.CREATE) } returns createToken + every { keyLock.unLock(key, LockType.CREATE, createToken) } returns true every { callable.call() } returns null val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) @@ -220,7 +235,7 @@ class ReqShieldTest : BaseReqShieldTest { verify { cacheGetter.invoke(key) } verify { cacheSetter.invoke(key, result, timeToLiveMillis) } verify { keyLock.tryLock(key, LockType.CREATE) } - verify { keyLock.unLock(key, LockType.CREATE) } + verify { keyLock.unLock(key, LockType.CREATE, createToken) } verify { callable.call() } } } @@ -230,8 +245,8 @@ class ReqShieldTest : BaseReqShieldTest { every { cacheGetter.invoke(key) } returns null every { cacheSetter.invoke(key, any(), any()) } returns true - every { globalLockFunc(any(), any()) } returns true - every { globalUnLockFunc(any()) } returns true + every { globalLockFunc(any(), any(), any()) } returns true + every { globalUnLockFunc(any(), any()) } returns true every { callable.call() } returns null @@ -243,10 +258,8 @@ class ReqShieldTest : BaseReqShieldTest { verify { cacheGetter.invoke(key) } verify { cacheSetter.invoke(key, result, timeToLiveMillis) } - verify { globalLockFunc(any(), any()) } - verify { globalUnLockFunc(any()) } - verify { keyGlobalLock.tryLock(key, LockType.CREATE) } - verify { keyGlobalLock.unLock(key, LockType.CREATE) } + verify { globalLockFunc(any(), any(), any()) } + verify { globalUnLockFunc(any(), any()) } verify { callable.call() } } } @@ -255,18 +268,19 @@ class ReqShieldTest : BaseReqShieldTest { override fun testSetMethodCacheNotExistsAndLocalLockAcquiredAndThrowCallableClientException() { every { cacheGetter.invoke(key) } returns null every { cacheSetter.invoke(key, any(), any()) } returns true - every { keyLock.tryLock(key, LockType.CREATE) } returns true - every { keyLock.unLock(key, LockType.CREATE) } returns true + every { keyLock.tryLock(key, LockType.CREATE) } returns createToken + every { keyLock.unLock(key, LockType.CREATE, createToken) } returns true every { callable.call() } throws Exception("callable error") - val exceptionCode = - assertThrows { reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) }.errorCode + val exception = + assertThrows { reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) } await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { - assertEquals(ErrorCode.SUPPLIER_ERROR, exceptionCode) + assertEquals(ErrorCode.SUPPLIER_ERROR, exception.errorCode) + assertEquals("callable error", exception.cause?.message) verify { cacheGetter.invoke(key) } verify { keyLock.tryLock(key, LockType.CREATE) } - verify { keyLock.unLock(key, LockType.CREATE) } + verify { keyLock.unLock(key, LockType.CREATE, createToken) } verify { callable.call() } } } @@ -276,8 +290,8 @@ class ReqShieldTest : BaseReqShieldTest { every { cacheGetter.invoke(key) } returns null every { cacheSetter.invoke(key, any(), any()) } returns true - every { globalLockFunc(any(), any()) } returns true - every { globalUnLockFunc(any()) } returns true + every { globalLockFunc(any(), any(), any()) } returns true + every { globalUnLockFunc(any(), any()) } returns true every { callable.call() } throws Exception("callable error") @@ -287,10 +301,8 @@ class ReqShieldTest : BaseReqShieldTest { await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { assertEquals(ErrorCode.SUPPLIER_ERROR, exceptionCode) verify { cacheGetter.invoke(key) } - verify { globalLockFunc(any(), any()) } - verify { globalUnLockFunc(any()) } - verify { keyGlobalLock.tryLock(key, LockType.CREATE) } - verify { keyGlobalLock.unLock(key, LockType.CREATE) } + verify { globalLockFunc(any(), any(), any()) } + verify { globalUnLockFunc(any(), any()) } verify { callable.call() } } } @@ -299,17 +311,16 @@ class ReqShieldTest : BaseReqShieldTest { override fun testSetMethodCacheNotExistsAndLocalLockAcquiredAndThrowGetCacheClientException() { every { cacheGetter.invoke(key) } throws Exception("get cache error") every { cacheSetter.invoke(key, any(), any()) } returns true - every { keyLock.tryLock(key, LockType.CREATE) } returns true - every { keyLock.unLock(key, LockType.CREATE) } returns true - val exceptionCode = - assertThrows { reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) }.errorCode + val exception = + assertThrows { reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) } await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { - assertEquals(ErrorCode.GET_CACHE_ERROR, exceptionCode) + assertEquals(ErrorCode.GET_CACHE_ERROR, exception.errorCode) + assertEquals("get cache error", exception.cause?.message) verify { cacheGetter.invoke(key) } verify(inverse = true) { keyLock.tryLock(key, LockType.CREATE) } - verify(inverse = true) { keyLock.unLock(key, LockType.CREATE) } + verify(inverse = true) { keyLock.unLock(key, LockType.CREATE, any()) } verify(inverse = true) { callable.call() } } } @@ -319,8 +330,8 @@ class ReqShieldTest : BaseReqShieldTest { every { cacheGetter.invoke(key) } throws Exception("get cache error") every { cacheSetter.invoke(key, any(), any()) } returns true - every { globalLockFunc(any(), any()) } returns true - every { globalUnLockFunc(any()) } returns true + every { globalLockFunc(any(), any(), any()) } returns true + every { globalUnLockFunc(any(), any()) } returns true val exceptionCode = assertThrows { reqShieldForGlobalLock.getAndSetReqShieldData(key, callable, timeToLiveMillis) }.errorCode @@ -328,10 +339,8 @@ class ReqShieldTest : BaseReqShieldTest { await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { assertEquals(ErrorCode.GET_CACHE_ERROR, exceptionCode) verify { cacheGetter.invoke(key) } - verify(inverse = true) { globalLockFunc(any(), any()) } - verify(inverse = true) { globalUnLockFunc(any()) } - verify(inverse = true) { keyGlobalLock.tryLock(key, LockType.CREATE) } - verify(inverse = true) { keyGlobalLock.unLock(key, LockType.CREATE) } + verify(inverse = true) { globalLockFunc(any(), any(), any()) } + verify(inverse = true) { globalUnLockFunc(any(), any()) } verify(inverse = true) { callable.call() } } } @@ -339,13 +348,13 @@ class ReqShieldTest : BaseReqShieldTest { @Test override fun testSetMethodCacheNotExistsAndLocalLockNotAcquired() { every { cacheGetter.invoke(key) } returns null - every { keyLock.tryLock(key, LockType.CREATE) } returns false + every { keyLock.tryLock(key, LockType.CREATE) } returns null val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { - result.value != null assertNotNull(result) + assertEquals(value, result.value) verify { cacheGetter.invoke(key) } verify { keyLock.tryLock(key, LockType.CREATE) } verify { callable.call() } @@ -355,32 +364,33 @@ class ReqShieldTest : BaseReqShieldTest { @Test override fun testSetMethodCacheNotExistsAndGlobalLockNotAcquired() { every { cacheGetter.invoke(key) } returns null - every { globalLockFunc(any(), any()) } returns false + every { globalLockFunc(any(), any(), any()) } returns false val result = reqShieldForGlobalLock.getAndSetReqShieldData(key, callable, timeToLiveMillis) await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { - result.value != null assertNotNull(result) + assertEquals(value, result.value) verify { cacheGetter.invoke(key) } - verify { keyGlobalLock.tryLock(key, LockType.CREATE) } + verify { globalLockFunc(any(), any(), any()) } verify { callable.call() } } } @Test override fun testSetMethodCacheExistsButNotTargetedForUpdate() { - val reqShieldData = ReqShieldData(value, timeToLiveMillis) + val reqShieldData = freshData(value) every { cacheGetter.invoke(key) } returns reqShieldData - every { keyLock.tryLock(key, LockType.UPDATE) } returns false val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { assertEquals(reqShieldData, result) verify { cacheGetter.invoke(key) } - verify(inverse = true) { cacheSetter.invoke(key, reqShieldData, timeToLiveMillis) } + // A fresh cache entry is not an update target, so no lock is even attempted + verify(inverse = true) { keyLock.tryLock(key, LockType.UPDATE) } + verify(inverse = true) { cacheSetter.invoke(key, any(), any()) } verify(inverse = true) { callable.call() } } } @@ -388,22 +398,21 @@ class ReqShieldTest : BaseReqShieldTest { @Test override fun testSetMethodCacheExistsAndTheUpdateTarget() { timeToLiveMillis = 1000 - val reqShieldData = ReqShieldData(oldValue, timeToLiveMillis) - val newReqShieldData = ReqShieldData(value, timeToLiveMillis) + val reqShieldData = updateTargetData(oldValue) every { cacheGetter.invoke(key) } returns reqShieldData every { cacheSetter.invoke(key, any(), any()) } answers { true } - every { keyLock.tryLock(key, LockType.UPDATE) } returns true - every { keyLock.unLock(key, LockType.UPDATE) } returns true + every { keyLock.tryLock(key, LockType.UPDATE) } returns updateToken + every { keyLock.unLock(key, LockType.UPDATE, updateToken) } returns true val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { assertEquals(reqShieldData, result) verify { cacheGetter.invoke(key) } - verify { cacheSetter.invoke(key, newReqShieldData, timeToLiveMillis) } + verify { cacheSetter.invoke(key, match { it.value == value }, timeToLiveMillis) } verify { keyLock.tryLock(key, LockType.UPDATE) } - verify { keyLock.unLock(key, LockType.UPDATE) } + verify { keyLock.unLock(key, LockType.UPDATE, updateToken) } verify { callable.call() } } } @@ -411,8 +420,7 @@ class ReqShieldTest : BaseReqShieldTest { @Test override fun testSetMethodCacheExistsAndTheUpdateTargetOnlyCreateCache() { timeToLiveMillis = 1000 - val reqShieldData = ReqShieldData(oldValue, timeToLiveMillis) - val newReqShieldData = ReqShieldData(value, timeToLiveMillis) + val reqShieldData = updateTargetData(oldValue) every { cacheGetter.invoke(key) } returns reqShieldData every { cacheSetter.invoke(key, any(), any()) } answers { true } @@ -422,9 +430,10 @@ class ReqShieldTest : BaseReqShieldTest { await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { assertEquals(reqShieldData, result) verify { cacheGetter.invoke(key) } - verify { cacheSetter.invoke(key, newReqShieldData, timeToLiveMillis) } + verify { cacheSetter.invoke(key, match { it.value == value }, timeToLiveMillis) } + // ONLY_CREATE_CACHE collapses on creation only, so the update takes no lock verify(inverse = true) { keyLock.tryLock(key, LockType.UPDATE) } - verify(inverse = true) { keyLock.unLock(key, LockType.UPDATE) } + verify(inverse = true) { keyLock.unLock(key, LockType.UPDATE, any()) } verify { callable.call() } } } @@ -432,13 +441,12 @@ class ReqShieldTest : BaseReqShieldTest { @Test override fun testSetMethodCacheExistsAndTheUpdateTargetAndCallableReturnNull() { timeToLiveMillis = 1000 - val reqShieldData = ReqShieldData(value, timeToLiveMillis) - val reqShieldDataNull = ReqShieldData(null, timeToLiveMillis) + val reqShieldData = updateTargetData(value) every { cacheGetter.invoke(key) } returns reqShieldData every { cacheSetter.invoke(key, any(), any()) } answers { true } - every { keyLock.tryLock(key, LockType.UPDATE) } returns true - every { keyLock.unLock(key, LockType.UPDATE) } returns true + every { keyLock.tryLock(key, LockType.UPDATE) } returns updateToken + every { keyLock.unLock(key, LockType.UPDATE, updateToken) } returns true every { callable.call() } returns null val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) @@ -446,17 +454,16 @@ class ReqShieldTest : BaseReqShieldTest { await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { assertEquals(reqShieldData, result) verify { cacheGetter.invoke(key) } - verify { cacheSetter.invoke(key, reqShieldDataNull, timeToLiveMillis) } + verify { cacheSetter.invoke(key, match { it.value == null }, timeToLiveMillis) } verify { keyLock.tryLock(key, LockType.UPDATE) } - verify { keyLock.unLock(key, LockType.UPDATE) } + verify { keyLock.unLock(key, LockType.UPDATE, updateToken) } verify { callable.call() } } } @Test override fun executeSetCacheFunctionShouldHandleExceptionFromCacheSetter() { - every { keyLock.tryLock(any(), any()) } returns true - every { keyLock.unLock(any(), any()) } returns true + every { keyLock.unLock(any(), any(), any()) } returns true val key = "key" val reqShieldData = ReqShieldData(value, 1000L) @@ -472,7 +479,7 @@ class ReqShieldTest : BaseReqShieldTest { val exception = assertFailsWith { - method.invoke(reqShield, cacheSetter, key, reqShieldData, lockType) + method.invoke(reqShield, cacheSetter, key, reqShieldData, lockType, createToken) } val cause = exception.cause @@ -480,90 +487,210 @@ class ReqShieldTest : BaseReqShieldTest { assertEquals(ErrorCode.SET_CACHE_ERROR, (cause as ClientException).errorCode) verify { cacheSetter.invoke(key, reqShieldData, 1000L) } - verify { keyLock.unLock(any(), any()) } + verify { keyLock.unLock(key, lockType, createToken) } } @Test - fun `should complete future with callable result when cache getter throws exception in scheduled task`() { - // Given: First cache check returns null (triggers scheduleTask path), - // then subsequent calls in scheduleTask throw exception - var callCount = 0 - every { cacheGetter.invoke(key) } answers { - callCount++ - if (callCount == 1) { - null // First call: cache miss, triggers handleLockForCacheCreation - } else { - throw Exception("cache connection error") // Subsequent calls: exception in scheduleTask + fun `should not unlock when no lock was taken and the supplier fails`() { + every { cacheGetter.invoke(key) } returns null + every { callable.call() } throws Exception("callable error") + + val exception = + assertThrows { + reqShieldOnlyUpdateCache.getAndSetReqShieldData(key, callable, timeToLiveMillis) } - } - every { keyLock.tryLock(key, LockType.CREATE) } returns false - // When: getAndSetReqShieldData is called - // The scheduled task will hit exception, should fallback to callable + assertEquals(ErrorCode.SUPPLIER_ERROR, exception.errorCode) + // ONLY_UPDATE_CACHE takes no CREATE lock, so nothing may be released on failure + verify(inverse = true) { keyLock.unLock(key, LockType.CREATE, any()) } + } + + @Test + fun `should release the lock and keep serving the caller when the asynchronous cache write fails`() { + every { cacheGetter.invoke(key) } returns null + every { cacheSetter.invoke(key, any(), any()) } throws Exception("set cache error") + every { keyLock.tryLock(key, LockType.CREATE) } returns createToken + every { keyLock.unLock(key, LockType.CREATE, createToken) } returns true + + // The cache write is fire-and-forget, so its failure must not reach the caller val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) - // Then: Should return callable result (fallback), not hang + assertEquals(value, result.value) await().atMost(Duration.ofMillis(AWAIT_TIMEOUT)).untilAsserted { - assertNotNull(result) - assertEquals(value, result.value) - verify { callable.call() } + verify { cacheSetter.invoke(key, result, timeToLiveMillis) } + // The lock must still be released even though the write failed + verify { keyLock.unLock(key, LockType.CREATE, createToken) } } } @Test - fun `should not hang when scheduled task encounters repeated cache getter exceptions`() { - // Given: First cache check returns null (triggers scheduleTask path), - // then subsequent calls always fail - var callCount = 0 - every { cacheGetter.invoke(key) } answers { - callCount++ - if (callCount == 1) { - null // First call: cache miss, triggers handleLockForCacheCreation - } else { - throw Exception("persistent cache error $callCount") // Subsequent calls: exception in scheduleTask + fun `should restore the interrupt flag and report a get cache error when the waiting thread is interrupted`() { + every { cacheGetter.invoke(key) } returns null + every { keyLock.tryLock(key, LockType.CREATE) } returns null + + var thrown: Throwable? = null + var interruptFlagRestored = false + val waiter = + Thread { + try { + reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) + } catch (e: Throwable) { + thrown = e + interruptFlagRestored = Thread.currentThread().isInterrupted + } } + + waiter.start() + Thread.sleep(100) // let the waiter enter the cache polling wait + waiter.interrupt() + waiter.join(2000) + + val exception = thrown + assertTrue(exception is ClientException, "Expected a ClientException but was $exception") + assertEquals(ErrorCode.GET_CACHE_ERROR, (exception as ClientException).errorCode) + assertTrue(interruptFlagRestored, "The interrupt flag must be restored on the caller thread") + verify(exactly = 0) { callable.call() } + } + + @Test + fun `should return the data another request wrote into the cache without calling the supplier`() { + val cachedData = freshData(value) + var getCount = 0 + // 1st read: the cache miss that leads into the lock-wait, the 3rd poll finds the data + every { cacheGetter.invoke(key) } answers { + getCount++ + if (getCount >= 4) cachedData else null + } + every { keyLock.tryLock(key, LockType.CREATE) } returns null + + val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) + + assertEquals(cachedData, result) + verify(exactly = 0) { callable.call() } + } + + @Test + fun `should fall back to the supplier once when cache reads keep failing`() { + var getCount = 0 + every { cacheGetter.invoke(key) } answers { + getCount++ + if (getCount == 1) null else throw RuntimeException("cache connection error") } - every { keyLock.tryLock(key, LockType.CREATE) } returns false + every { keyLock.tryLock(key, LockType.CREATE) } returns null - // When: getAndSetReqShieldData is called val startTime = System.currentTimeMillis() val result = reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) val elapsed = System.currentTimeMillis() - startTime - // Then: Should complete within reasonable time (not hang), using callable fallback - assertNotNull(result) assertEquals(value, result.value) - // Should complete within 5 seconds (way less than infinite hang) - assertTrue(elapsed < 5000, "Should not hang - completed in ${elapsed}ms") - verify { callable.call() } + // The supplier is called on the caller thread, exactly once - never from the polling task + verify(exactly = 1) { callable.call() } + // Consecutive failures short-circuit the wait instead of polling maxAttemptGetCache times + assertTrue(elapsed < 1000, "Consecutive cache failures should stop the wait early, took ${elapsed}ms") + assertEquals(1 + MAX_CONSECUTIVE_GET_CACHE_FAILURES, getCount) } @Test - fun `should propagate exception when both cache getter and callable fail`() { - // Given: First cache check returns null (triggers scheduleTask path), - // then cache getter fails and callable also fails - var callCount = 0 + fun `should keep waiting when cache read failures are not consecutive`() { + val maxAttemptGetCache = 5 + val reqShieldSmallAttempt = + ReqShield( + ReqShieldConfiguration( + cacheSetter, + cacheGetter, + keyLock = keyLock, + maxAttemptGetCache = maxAttemptGetCache, + ), + ) + + var getCount = 0 + // After the initial miss the polls alternate: empty, failed, empty, failed, ... + // so the consecutive failure counter is reset before it can reach its limit. every { cacheGetter.invoke(key) } answers { - callCount++ - if (callCount == 1) { - null // First call: cache miss - } else { - throw Exception("cache error") // Subsequent calls: exception in scheduleTask - } + getCount++ + if (getCount > 1 && getCount % 2 == 1) throw RuntimeException("transient cache error") + null } - every { keyLock.tryLock(key, LockType.CREATE) } returns false - every { callable.call() } throws Exception("callable also failed") + every { keyLock.tryLock(key, LockType.CREATE) } returns null - // When/Then: Should propagate the callable exception (via completeExceptionally) - val exception = - assertThrows { - reqShield.getAndSetReqShieldData(key, callable, timeToLiveMillis) - } + val result = reqShieldSmallAttempt.getAndSetReqShieldData(key, callable, timeToLiveMillis) - // The exception should be from the fallback callable failure + assertEquals(value, result.value) + verify(exactly = 1) { callable.call() } + // 1 initial read + 5 empty polls + 4 interleaved failures; bailing out early would read less assertTrue( - exception is ClientException || exception.cause is ClientException, - "Should propagate ClientException from failed callable", + getCount >= 1 + maxAttemptGetCache * 2 - 1, + "Expected at least ${1 + maxAttemptGetCache * 2 - 1} cache reads but was $getCount", ) } + + @Test + fun `should throw ClientException with supplier error when the supplier fails after max attempts`() { + val reqShieldSmallAttempt = + ReqShield( + ReqShieldConfiguration( + cacheSetter, + cacheGetter, + keyLock = keyLock, + maxAttemptGetCache = 3, + ), + ) + + every { cacheGetter.invoke(key) } returns null + every { keyLock.tryLock(key, LockType.CREATE) } returns null + every { callable.call() } throws IllegalStateException("supplier down") + + val exception = + assertThrows { + reqShieldSmallAttempt.getAndSetReqShieldData(key, callable, timeToLiveMillis) + } + + // The failure must not leak as an ExecutionException wrapper and must keep its cause + assertEquals(ErrorCode.SUPPLIER_ERROR, exception.errorCode) + assertEquals("supplier down", exception.cause?.message) + verify(exactly = 1) { callable.call() } + } + + @Test + fun shouldCancelQueuedPollingWhenWaitingTimesOut() { + val executor = Executors.newSingleThreadScheduledExecutor() + val releaseExecutor = CountDownLatch(1) + val executorBlocked = CountDownLatch(1) + val reads = AtomicInteger() + every { keyLock.tryLock(key, LockType.CREATE) } returns null + val shield = + ReqShield( + ReqShieldConfiguration( + setCacheFunction = cacheSetter, + getCacheFunction = { + reads.incrementAndGet() + null + }, + keyLock = keyLock, + executor = executor, + maxAttemptGetCache = 1, + ), + ) + + try { + executor.submit { + executorBlocked.countDown() + releaseExecutor.await() + } + assertTrue(executorBlocked.await(2, TimeUnit.SECONDS)) + + val result = shield.getAndSetReqShieldData(key, callable, timeToLiveMillis) + + assertEquals(value, result.value) + verify(exactly = 1) { callable.call() } + releaseExecutor.countDown() + // The barrier runs after the queued poll would have become eligible. + executor.schedule({}, GET_CACHE_INTERVAL_MILLIS * 2, TimeUnit.MILLISECONDS).get(2, TimeUnit.SECONDS) + assertEquals(1, reads.get(), "Only the initial cache read should run") + } finally { + releaseExecutor.countDown() + executor.shutdownNow() + assertTrue(executor.awaitTermination(2, TimeUnit.SECONDS)) + } + } } diff --git a/libs.versions.toml b/libs.versions.toml index 48a0376..23e714e 100644 --- a/libs.versions.toml +++ b/libs.versions.toml @@ -8,7 +8,7 @@ springBoot3 = "3.3.1" springBoot2 = "2.7.17" springDependency = "1.1.5" junit = "5.9.1" -testcontainers = "1.18.1" +testcontainers = "1.21.4" [libraries] # kotlin diff --git a/req-shield-spring-boot3-example/build.gradle.kts b/req-shield-spring-boot3-example/build.gradle.kts index a75d10c..566592a 100644 --- a/req-shield-spring-boot3-example/build.gradle.kts +++ b/req-shield-spring-boot3-example/build.gradle.kts @@ -5,6 +5,10 @@ plugins { alias(libs.plugins.spring.dependency.management) } +// Spring Boot's BOM pins an older Testcontainers whose Docker client is rejected by current Docker daemons +// ("client version 1.32 is too old"); keep the version from the catalog instead. +extra["testcontainers.version"] = libs.versions.testcontainers.get() + dependencies { implementation(project(":core-spring")) diff --git a/req-shield-spring-boot3-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/cache/ReqShieldCacheImpl.kt b/req-shield-spring-boot3-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/cache/ReqShieldCacheImpl.kt index d6dab9d..a79c5ca 100644 --- a/req-shield-spring-boot3-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/cache/ReqShieldCacheImpl.kt +++ b/req-shield-spring-boot3-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/cache/ReqShieldCacheImpl.kt @@ -1,8 +1,10 @@ package com.linecorp.cse.reqshield.spring3.mvc.example.cache +import com.linecorp.cse.reqshield.spring.cache.GlobalLockSupport import com.linecorp.cse.reqshield.spring.cache.ReqShieldCache import com.linecorp.cse.reqshield.support.model.ReqShieldData import org.springframework.data.redis.core.RedisTemplate +import org.springframework.data.redis.core.script.DefaultRedisScript import org.springframework.stereotype.Service import java.time.Duration @@ -10,7 +12,8 @@ import java.time.Duration class ReqShieldCacheImpl( private val redisTemplate: RedisTemplate>, private val redisTemplateForGlobalLock: RedisTemplate, -) : ReqShieldCache { +) : ReqShieldCache, + GlobalLockSupport { override fun get(key: String): ReqShieldData? = redisTemplate.opsForValue()[key] override fun put( @@ -22,9 +25,22 @@ class ReqShieldCacheImpl( override fun evict(key: String): Boolean? = redisTemplate.delete(key) override fun globalLock( - key: String, + lockKey: String, + token: String, timeToLiveMillis: Long, - ): Boolean = redisTemplateForGlobalLock.opsForValue().setIfAbsent(key, key, Duration.ofMillis(timeToLiveMillis)) ?: false + ): Boolean = redisTemplateForGlobalLock.opsForValue().setIfAbsent(lockKey, token, Duration.ofMillis(timeToLiveMillis)) ?: false + + override fun globalUnLock( + lockKey: String, + token: String, + ): Boolean = redisTemplateForGlobalLock.execute(UN_LOCK_SCRIPT, listOf(lockKey), token) == 1L - override fun globalUnLock(key: String): Boolean = redisTemplateForGlobalLock.delete(key) + companion object { + /** Compare-and-delete, so an expired holder cannot release the lock of the next holder. */ + private val UN_LOCK_SCRIPT = + DefaultRedisScript( + "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end", + Long::class.javaObjectType, + ) + } } diff --git a/req-shield-spring-boot3-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/configuration/RedisConfiguration.kt b/req-shield-spring-boot3-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/configuration/RedisConfiguration.kt index 28e9fc2..166faa1 100644 --- a/req-shield-spring-boot3-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/configuration/RedisConfiguration.kt +++ b/req-shield-spring-boot3-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/configuration/RedisConfiguration.kt @@ -42,15 +42,11 @@ class RedisConfiguration { @Bean fun redisTemplateForGlobalLock(connectionFactory: RedisConnectionFactory): RedisTemplate { - val valueSerializer = - Jackson2JsonRedisSerializer(String::class.java).apply { - setObjectMapper(objectMapper()) - } - val redisTemplate = RedisTemplate() redisTemplate.setConnectionFactory(connectionFactory) redisTemplate.keySerializer = StringRedisSerializer() - redisTemplate.valueSerializer = valueSerializer + // The lock token is stored as a plain string so the compare-and-delete script can compare it as is + redisTemplate.valueSerializer = StringRedisSerializer() return redisTemplate } diff --git a/req-shield-spring-boot3-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/service/CacheAnnotationTest.kt b/req-shield-spring-boot3-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/service/CacheAnnotationTest.kt index d5404c3..076fcb1 100644 --- a/req-shield-spring-boot3-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/service/CacheAnnotationTest.kt +++ b/req-shield-spring-boot3-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/service/CacheAnnotationTest.kt @@ -48,7 +48,7 @@ class CacheAnnotationTest : AbstractRedisTest() { await().atMost(5, TimeUnit.SECONDS).untilAsserted { assertEquals(1, sampleService.getRequestCount()) - assertNotNull(reqShieldCache.get("product-$testProductId")) + assertNotNull(reqShieldCache.get("product::product-$testProductId")) } } @@ -69,7 +69,7 @@ class CacheAnnotationTest : AbstractRedisTest() { await().atMost(5, TimeUnit.SECONDS).untilAsserted { assertEquals(100, sampleService.getRequestCount()) - assertNotNull(reqShieldCache.get("product-$testProductId")) + assertNotNull(reqShieldCache.get("productOnlyUpdateCache::product-$testProductId")) } } @@ -92,7 +92,7 @@ class CacheAnnotationTest : AbstractRedisTest() { await().atMost(5, TimeUnit.SECONDS).untilAsserted { assertEquals(1, sampleService.getRequestCount()) - assertNotNull(reqShieldCache.get("product-$testProductId")) + assertNotNull(reqShieldCache.get("product::product-$testProductId")) } } @@ -103,19 +103,19 @@ class CacheAnnotationTest : AbstractRedisTest() { sampleService.getProduct(testProductId) await().atMost(5, TimeUnit.SECONDS).until { - runCatching { reqShieldCache.get("product-$testProductId") != null }.getOrDefault(false) + runCatching { reqShieldCache.get("product::product-$testProductId") != null }.getOrDefault(false) } - assertNotNull(reqShieldCache.get("product-$testProductId")) + assertNotNull(reqShieldCache.get("product::product-$testProductId")) // when sampleService.removeProduct(testProductId) // then await().atMost(5, TimeUnit.SECONDS).until { - runCatching { reqShieldCache.get("product-$testProductId") == null }.getOrDefault(false) + runCatching { reqShieldCache.get("product::product-$testProductId") == null }.getOrDefault(false) } - assertNull(reqShieldCache.get("product-$testProductId")) + assertNull(reqShieldCache.get("product::product-$testProductId")) } } diff --git a/req-shield-spring-boot3-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/service/GlobalLockTest.kt b/req-shield-spring-boot3-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/service/GlobalLockTest.kt new file mode 100644 index 0000000..22dd3e4 --- /dev/null +++ b/req-shield-spring-boot3-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/mvc/example/service/GlobalLockTest.kt @@ -0,0 +1,52 @@ +/* + * 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.spring3.mvc.example.service + +import com.linecorp.cse.reqshield.spring.cache.GlobalLockSupport +import com.linecorp.cse.reqshield.support.redis.AbstractRedisTest +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.extension.ExtendWith +import org.springframework.beans.factory.annotation.Autowired +import org.springframework.boot.test.context.SpringBootTest +import org.springframework.test.context.junit.jupiter.SpringExtension +import java.util.UUID + +@SpringBootTest +@ExtendWith(SpringExtension::class) +class GlobalLockTest : AbstractRedisTest() { + @Autowired + private lateinit var globalLockSupport: GlobalLockSupport + + @Test + fun `lock is exclusive and can only be released by the token that acquired it`() { + val lockKey = "globalLockTest-${UUID.randomUUID()}" + + assertTrue(globalLockSupport.globalLock(lockKey, "ownerToken", 5000)) + // already held + assertFalse(globalLockSupport.globalLock(lockKey, "otherToken", 5000)) + // compare-and-delete: a foreign token must not release someone else's lock + assertFalse(globalLockSupport.globalUnLock(lockKey, "otherToken")) + assertFalse(globalLockSupport.globalLock(lockKey, "otherToken", 5000)) + + assertTrue(globalLockSupport.globalUnLock(lockKey, "ownerToken")) + // released, so it can be acquired again + assertTrue(globalLockSupport.globalLock(lockKey, "otherToken", 5000)) + assertTrue(globalLockSupport.globalUnLock(lockKey, "otherToken")) + } +} diff --git a/req-shield-spring-boot3-webflux-example/build.gradle.kts b/req-shield-spring-boot3-webflux-example/build.gradle.kts index 94706ad..28e5a39 100644 --- a/req-shield-spring-boot3-webflux-example/build.gradle.kts +++ b/req-shield-spring-boot3-webflux-example/build.gradle.kts @@ -5,6 +5,10 @@ plugins { alias(libs.plugins.spring.dependency.management) } +// Spring Boot's BOM pins an older Testcontainers whose Docker client is rejected by current Docker daemons +// ("client version 1.32 is too old"); keep the version from the catalog instead. +extra["testcontainers.version"] = libs.versions.testcontainers.get() + group = "com.linecorp.cse.reqshield" version = "1.0.0" diff --git a/req-shield-spring-boot3-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/example/cache/AsyncCacheImpl.kt b/req-shield-spring-boot3-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/example/cache/AsyncCacheImpl.kt index 7e0a0d6..341523b 100644 --- a/req-shield-spring-boot3-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/example/cache/AsyncCacheImpl.kt +++ b/req-shield-spring-boot3-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/example/cache/AsyncCacheImpl.kt @@ -1,17 +1,26 @@ package com.linecorp.cse.reqshield.spring3.webflux.example.cache import com.linecorp.cse.reqshield.spring.webflux.cache.AsyncCache +import com.linecorp.cse.reqshield.spring.webflux.cache.GlobalLockSupport import com.linecorp.cse.reqshield.support.model.ReqShieldData import org.springframework.data.redis.core.ReactiveRedisOperations +import org.springframework.data.redis.core.script.RedisScript import org.springframework.stereotype.Service import reactor.core.publisher.Mono import java.time.Duration +/** + * Compare-and-delete: the lock is only released when it still holds the caller's token, so an owner + * whose lock already expired cannot release the lock of the next owner. + */ +private const val UNLOCK_SCRIPT = "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end" + @Service class AsyncCacheImpl( private val redisOperations: ReactiveRedisOperations>, private val redisOperationsForGlobalLock: ReactiveRedisOperations, -) : AsyncCache { +) : AsyncCache, + GlobalLockSupport { override fun get(key: String): Mono?> = redisOperations.opsForValue()[key] override fun put( @@ -23,13 +32,18 @@ class AsyncCacheImpl( override fun evict(key: String): Mono = redisOperations.opsForValue().delete(key) override fun globalLock( - key: String, + lockKey: String, + token: String, timeToLiveMillis: Long, - ): Mono = - redisOperationsForGlobalLock.opsForValue().setIfAbsent(key, key, Duration.ofMillis(timeToLiveMillis)) ?: Mono.just(false) + ): Mono = redisOperationsForGlobalLock.opsForValue().setIfAbsent(lockKey, token, Duration.ofMillis(timeToLiveMillis)) - override fun globalUnLock(key: String): Mono = + override fun globalUnLock( + lockKey: String, + token: String, + ): Mono = redisOperationsForGlobalLock - .delete(key) - .map { count -> count > 0 } + .execute(RedisScript.of(UNLOCK_SCRIPT, Long::class.java), listOf(lockKey), listOf(token)) + .next() + .map { it == 1L } + .defaultIfEmpty(false) } diff --git a/req-shield-spring-boot3-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/example/configuration/RedisConfiguration.kt b/req-shield-spring-boot3-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/example/configuration/RedisConfiguration.kt index b6afcad..c476622 100644 --- a/req-shield-spring-boot3-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/example/configuration/RedisConfiguration.kt +++ b/req-shield-spring-boot3-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/example/configuration/RedisConfiguration.kt @@ -53,10 +53,9 @@ class RedisConfiguration { @Bean("redisOperationsForGlobalLock") fun reactiveRedisOperationsForGlobalLock(factory: ReactiveRedisConnectionFactory): ReactiveRedisOperations { val keySerializer = StringRedisSerializer() - val valueSerializer = - Jackson2JsonRedisSerializer(String::class.java).apply { - setObjectMapper(objectMapper()) - } + // The lock token has to be stored as a plain string: the compare-and-delete Lua script + // compares the stored value with the raw token it receives as a script argument. + val valueSerializer = StringRedisSerializer() val serializationContext = RedisSerializationContext .newSerializationContext(keySerializer) diff --git a/req-shield-spring-boot3-webflux-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/webflux/example/service/CacheAnnotationTest.kt b/req-shield-spring-boot3-webflux-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/webflux/example/service/CacheAnnotationTest.kt index 75ce088..a7b1c4c 100644 --- a/req-shield-spring-boot3-webflux-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/webflux/example/service/CacheAnnotationTest.kt +++ b/req-shield-spring-boot3-webflux-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/webflux/example/service/CacheAnnotationTest.kt @@ -53,7 +53,7 @@ class CacheAnnotationTest : AbstractRedisTest() { .verify() await().atMost(5, TimeUnit.SECONDS).until { - asyncCache.get("product-$testProductId").block() != null + asyncCache.get("product::product-$testProductId").block() != null } } @@ -83,7 +83,7 @@ class CacheAnnotationTest : AbstractRedisTest() { assertEquals(20, sampleService.getRequestCount(), "Request count should be 20") await().atMost(5, TimeUnit.SECONDS).until { - asyncCache.get("product-$testProductId").block() != null + asyncCache.get("productOnlyUpdataCache::product-$testProductId").block() != null } } @@ -108,7 +108,7 @@ class CacheAnnotationTest : AbstractRedisTest() { .verify() await().atMost(5, TimeUnit.SECONDS).until { - asyncCache.get("product-$testProductId").block() != null + asyncCache.get("product::product-$testProductId").block() != null } } @@ -128,9 +128,9 @@ class CacheAnnotationTest : AbstractRedisTest() { // then await().atMost(5, TimeUnit.SECONDS).until { - asyncCache.get("product-$testProductId").block() != null + asyncCache.get("product::product-$testProductId").block() != null } - val cacheMono = asyncCache.get("product-$testProductId").block() + val cacheMono = asyncCache.get("product::product-$testProductId").block() assertNotNull(cacheMono) // when @@ -146,9 +146,9 @@ class CacheAnnotationTest : AbstractRedisTest() { // then await().atMost(5, TimeUnit.SECONDS).until { - asyncCache.get("product-$testProductId").block() == null + asyncCache.get("product::product-$testProductId").block() == null } - val cacheMonoNull = asyncCache.get("product-$testProductId").block() + val cacheMonoNull = asyncCache.get("product::product-$testProductId").block() assertNull(cacheMonoNull) } } diff --git a/req-shield-spring-boot3-webflux-kotlin-coroutine-example/build.gradle.kts b/req-shield-spring-boot3-webflux-kotlin-coroutine-example/build.gradle.kts index c44b648..dd50c0f 100644 --- a/req-shield-spring-boot3-webflux-kotlin-coroutine-example/build.gradle.kts +++ b/req-shield-spring-boot3-webflux-kotlin-coroutine-example/build.gradle.kts @@ -5,6 +5,10 @@ plugins { alias(libs.plugins.spring.dependency.management) } +// Spring Boot's BOM pins an older Testcontainers whose Docker client is rejected by current Docker daemons +// ("client version 1.32 is too old"); keep the version from the catalog instead. +extra["testcontainers.version"] = libs.versions.testcontainers.get() + group = "com.linecorp.cse.reqshield" version = "1.0.0" diff --git a/req-shield-spring-boot3-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/kotlin/coroutine/example/cache/AsyncCacheImpl.kt b/req-shield-spring-boot3-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/kotlin/coroutine/example/cache/AsyncCacheImpl.kt index 20f1f26..1a6b118 100644 --- a/req-shield-spring-boot3-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/kotlin/coroutine/example/cache/AsyncCacheImpl.kt +++ b/req-shield-spring-boot3-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/kotlin/coroutine/example/cache/AsyncCacheImpl.kt @@ -1,8 +1,11 @@ package com.linecorp.cse.reqshield.spring3.webflux.kotlin.coroutine.example.cache import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.cache.AsyncCache +import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.cache.GlobalLockSupport import com.linecorp.cse.reqshield.support.model.ReqShieldData +import kotlinx.coroutines.reactive.awaitFirstOrNull import org.springframework.data.redis.core.* +import org.springframework.data.redis.core.script.RedisScript import org.springframework.stereotype.Service import java.time.Duration @@ -10,7 +13,8 @@ import java.time.Duration class AsyncCacheImpl( private val redisOperations: ReactiveRedisOperations>, private val redisOperationsForGlobalLock: ReactiveRedisOperations, -) : AsyncCache { +) : AsyncCache, + GlobalLockSupport { override suspend fun get(key: String): ReqShieldData? = redisOperations.opsForValue().getAndAwait(key) override suspend fun put( @@ -21,13 +25,34 @@ class AsyncCacheImpl( override suspend fun evict(key: String): Boolean = redisOperations.opsForValue().deleteAndAwait(key) + /** `SET lockKey token NX PX ttl`: the stored token identifies the caller that owns the lock. */ override suspend fun globalLock( - key: String, + lockKey: String, + token: String, timeToLiveMillis: Long, ): Boolean = redisOperationsForGlobalLock .opsForValue() - .setIfAbsentAndAwait(key, key, Duration.ofMillis(timeToLiveMillis)) + .setIfAbsentAndAwait(lockKey, token, Duration.ofMillis(timeToLiveMillis)) + + /** + * Compare-and-delete in a single atomic step, so an owner whose lock already expired can never + * release the lock that the next owner has taken in the meantime. + */ + override suspend fun globalUnLock( + lockKey: String, + token: String, + ): Boolean = + redisOperationsForGlobalLock + .execute(UNLOCK_SCRIPT, listOf(lockKey), listOf(token)) + .awaitFirstOrNull() == 1L - override suspend fun globalUnLock(key: String): Boolean = redisOperationsForGlobalLock.deleteAndAwait(key) > 0 + companion object { + /** Returns the number of keys deleted: 1 when this caller still owned the lock, 0 otherwise. */ + private val UNLOCK_SCRIPT: RedisScript = + RedisScript.of( + "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end", + Long::class.java, + ) + } } diff --git a/req-shield-spring-boot3-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/kotlin/coroutine/example/configuration/RedisConfiguration.kt b/req-shield-spring-boot3-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/kotlin/coroutine/example/configuration/RedisConfiguration.kt index a5d5458..9217a8b 100644 --- a/req-shield-spring-boot3-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/kotlin/coroutine/example/configuration/RedisConfiguration.kt +++ b/req-shield-spring-boot3-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring3/webflux/kotlin/coroutine/example/configuration/RedisConfiguration.kt @@ -53,10 +53,9 @@ class RedisConfiguration { @Bean("redisOperationsForGlobalLock") fun reactiveRedisOperationsForGlobalLock(factory: ReactiveRedisConnectionFactory): ReactiveRedisOperations { val keySerializer = StringRedisSerializer() - val valueSerializer = - Jackson2JsonRedisSerializer(String::class.java).apply { - setObjectMapper(objectMapper()) - } + // Lock values are raw ownership tokens, so they are stored verbatim: a JSON serializer would + // store them quoted, which only stays comparable while every access uses this same template. + val valueSerializer = StringRedisSerializer() val serializationContext = RedisSerializationContext .newSerializationContext(keySerializer) diff --git a/req-shield-spring-boot3-webflux-kotlin-coroutine-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/webflux/kotlin/coroutine/example/CacheAnnotationTest.kt b/req-shield-spring-boot3-webflux-kotlin-coroutine-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/webflux/kotlin/coroutine/example/CacheAnnotationTest.kt index b0a29fe..8c07e67 100644 --- a/req-shield-spring-boot3-webflux-kotlin-coroutine-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/webflux/kotlin/coroutine/example/CacheAnnotationTest.kt +++ b/req-shield-spring-boot3-webflux-kotlin-coroutine-example/src/test/kotlin/com/linecorp/cse/reqshield/spring3/webflux/kotlin/coroutine/example/CacheAnnotationTest.kt @@ -1,8 +1,12 @@ package com.linecorp.cse.reqshield.spring3.webflux.kotlin.coroutine.example +import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.annotation.ReqShieldCacheEvict import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.cache.AsyncCache +import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.cache.GlobalLockSupport import com.linecorp.cse.reqshield.spring3.webflux.kotlin.coroutine.example.dto.Product import com.linecorp.cse.reqshield.spring3.webflux.kotlin.coroutine.example.service.SampleService +import com.linecorp.cse.reqshield.support.constant.ConfigValues.LOCK_KEY_PREFIX +import com.linecorp.cse.reqshield.support.model.ReqShieldData import com.linecorp.cse.reqshield.support.redis.AbstractRedisTest import kotlinx.coroutines.async import kotlinx.coroutines.awaitAll @@ -15,19 +19,41 @@ import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test import org.junit.jupiter.api.extension.ExtendWith import org.springframework.beans.factory.annotation.Autowired +import org.springframework.beans.factory.annotation.Qualifier import org.springframework.boot.test.context.SpringBootTest +import org.springframework.boot.test.context.TestConfiguration +import org.springframework.context.annotation.Bean +import org.springframework.data.redis.core.ReactiveRedisOperations +import org.springframework.data.redis.core.getAndAwait import org.springframework.test.context.junit.jupiter.SpringExtension import java.util.* +import java.util.concurrent.atomic.AtomicBoolean +import kotlin.test.assertFailsWith -@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT) +@SpringBootTest( + classes = [SpringWebfluxCoroutineApplication::class, EvictionTestConfiguration::class], + webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT, +) @ExtendWith(SpringExtension::class) class CacheAnnotationTest : AbstractRedisTest() { @Autowired private lateinit var sampleService: SampleService + @Autowired + private lateinit var evictionTestService: EvictionTestService + @Autowired lateinit var asyncCache: AsyncCache + @Autowired + @Qualifier("redisOperationsForGlobalLock") + lateinit var globalLockOperations: ReactiveRedisOperations + + private val lockSupport: GlobalLockSupport get() = asyncCache as GlobalLockSupport + + /** The aspect namespaces every key with the cache name of the annotation. */ + private fun cacheKeyOf(productId: String) = "product::product-$productId" + @BeforeEach fun `reset request count`() = runTest { @@ -48,7 +74,7 @@ class CacheAnnotationTest : AbstractRedisTest() { delay(500) assertEquals(1, sampleService.getRequestCount()) - assertNotNull(asyncCache.get("product-$testProductId")) + assertNotNull(asyncCache.get(cacheKeyOf(testProductId))) } @Test @@ -81,7 +107,7 @@ class CacheAnnotationTest : AbstractRedisTest() { delay(500) assertEquals(1, sampleService.getRequestCount()) - assertNotNull(asyncCache.get("product-$testProductId")) + assertNotNull(asyncCache.get(cacheKeyOf(testProductId))) } @Test @@ -94,7 +120,7 @@ class CacheAnnotationTest : AbstractRedisTest() { val maxAttempts = 30 var attempts = 0 - while (asyncCache.get("product-$testProductId") == null) { + while (asyncCache.get(cacheKeyOf(testProductId)) == null) { if (attempts >= maxAttempts) { break } @@ -102,13 +128,13 @@ class CacheAnnotationTest : AbstractRedisTest() { delay(100) } - assertNotNull(asyncCache.get("product-$testProductId")) + assertNotNull(asyncCache.get(cacheKeyOf(testProductId))) // when sampleService.removeProduct(testProductId) var attemptsSecond = 0 - while (asyncCache.get("product-$testProductId") != null) { + while (asyncCache.get(cacheKeyOf(testProductId)) != null) { if (attemptsSecond >= maxAttempts) { break } @@ -116,6 +142,117 @@ class CacheAnnotationTest : AbstractRedisTest() { delay(100) } - assertNull(asyncCache.get("product-$testProductId")) + assertNull(asyncCache.get(cacheKeyOf(testProductId))) + } + + @Test + fun cacheEvictionShouldWaitForSuccessfulSuspendMethod() = + runBlocking { + val key = UUID.randomUUID().toString() + putEvictionTestEntry(key) + + val cacheWasPresentDuringMethod = evictionTestService.evictReturningValue(key) + + assertTrue(cacheWasPresentDuringMethod) + assertNull(asyncCache.get(evictionTestCacheKeyOf(key))) + } + + @Test + fun cacheEvictionShouldPreserveTheCacheWhenSuspendMethodFails() = + runBlocking { + val key = UUID.randomUUID().toString() + putEvictionTestEntry(key) + + val exception = assertFailsWith { evictionTestService.evictFailing(key) } + + assertEquals("cacheWasPresent=true", exception.message) + assertNotNull(asyncCache.get(evictionTestCacheKeyOf(key))) + } + + @Test + fun cacheEvictionShouldWaitForSuspendMethodReturningEmptyMono() = + runBlocking { + val key = UUID.randomUUID().toString() + val cacheWasPresentDuringMethod = AtomicBoolean() + putEvictionTestEntry(key) + + assertNull(evictionTestService.evictReturningNull(key, cacheWasPresentDuringMethod)) + + assertTrue(cacheWasPresentDuringMethod.get()) + assertNull(asyncCache.get(evictionTestCacheKeyOf(key))) + } + + @Test + fun globalLockShouldBeReleasedAfterTheOwningRequestCompletes() = + runBlocking { + // given + val testProductId: String = UUID.randomUUID().toString() + val lockKey = "$LOCK_KEY_PREFIX${cacheKeyOf(testProductId)}_CREATE" + + // when + sampleService.getProductForGlobalLock(testProductId) + + // then the background cache write releases the lock it owns + var attempts = 0 + while (globalLockOperations.opsForValue().getAndAwait(lockKey) != null && attempts < 30) { + attempts++ + delay(100) + } + assertNull(globalLockOperations.opsForValue().getAndAwait(lockKey)) + } + + @Test + fun onlyTheOwningTokenShouldReleaseTheGlobalLock() = + runBlocking { + val lockKey = "lock-token-${UUID.randomUUID()}" + + assertTrue(lockSupport.globalLock(lockKey, "owner", 10_000)) + // The token is stored as a plain string, which is what the compare-and-delete script compares. + assertEquals("owner", globalLockOperations.opsForValue().getAndAwait(lockKey)) + // SET NX: a second caller cannot take a held lock. + assertFalse(lockSupport.globalLock(lockKey, "intruder", 10_000)) + // Compare-and-delete: a non-owner cannot release it either. + assertFalse(lockSupport.globalUnLock(lockKey, "intruder")) + assertTrue(lockSupport.globalUnLock(lockKey, "owner")) + assertNull(globalLockOperations.opsForValue().getAndAwait(lockKey)) } + + private suspend fun putEvictionTestEntry(key: String) { + val product = Product(key, "product_$key") + assertTrue(asyncCache.put(evictionTestCacheKeyOf(key), ReqShieldData(product, 10_000), 10_000)) + } + + private fun evictionTestCacheKeyOf(key: String) = "$EVICTION_TEST_CACHE_NAME::$key" +} + +@TestConfiguration(proxyBeanMethods = false) +class EvictionTestConfiguration { + @Bean + fun evictionTestService(asyncCache: AsyncCache): EvictionTestService = EvictionTestService(asyncCache) } + +open class EvictionTestService( + private val asyncCache: AsyncCache, +) { + @ReqShieldCacheEvict(cacheName = EVICTION_TEST_CACHE_NAME, key = "#key") + open suspend fun evictReturningValue(key: String): Boolean = asyncCache.get(cacheKeyOf(key)) != null + + @ReqShieldCacheEvict(cacheName = EVICTION_TEST_CACHE_NAME, key = "#key") + open suspend fun evictFailing(key: String): Boolean { + val cacheWasPresent = asyncCache.get(cacheKeyOf(key)) != null + throw IllegalStateException("cacheWasPresent=$cacheWasPresent") + } + + @ReqShieldCacheEvict(cacheName = EVICTION_TEST_CACHE_NAME, key = "#key") + open suspend fun evictReturningNull( + key: String, + cacheWasPresent: AtomicBoolean, + ): String? { + cacheWasPresent.set(asyncCache.get(cacheKeyOf(key)) != null) + return null + } + + private fun cacheKeyOf(key: String) = "$EVICTION_TEST_CACHE_NAME::$key" +} + +private const val EVICTION_TEST_CACHE_NAME = "eviction-order" diff --git a/req-shield-spring-example/src/main/kotlin/com/linecorp/cse/reqshield/cache/ReqShieldCacheImpl.kt b/req-shield-spring-example/src/main/kotlin/com/linecorp/cse/reqshield/cache/ReqShieldCacheImpl.kt index eae963a..b2ab744 100644 --- a/req-shield-spring-example/src/main/kotlin/com/linecorp/cse/reqshield/cache/ReqShieldCacheImpl.kt +++ b/req-shield-spring-example/src/main/kotlin/com/linecorp/cse/reqshield/cache/ReqShieldCacheImpl.kt @@ -16,9 +16,11 @@ package com.linecorp.cse.reqshield.cache +import com.linecorp.cse.reqshield.spring.cache.GlobalLockSupport import com.linecorp.cse.reqshield.spring.cache.ReqShieldCache import com.linecorp.cse.reqshield.support.model.ReqShieldData import org.springframework.data.redis.core.RedisTemplate +import org.springframework.data.redis.core.script.DefaultRedisScript import org.springframework.stereotype.Service import java.time.Duration @@ -26,7 +28,8 @@ import java.time.Duration class ReqShieldCacheImpl( private val redisTemplate: RedisTemplate>, private val redisTemplateForGlobalLock: RedisTemplate, -) : ReqShieldCache { +) : ReqShieldCache, + GlobalLockSupport { override fun get(key: String): ReqShieldData? = redisTemplate.opsForValue()[key] override fun put( @@ -38,9 +41,22 @@ class ReqShieldCacheImpl( override fun evict(key: String): Boolean? = redisTemplate.delete(key) override fun globalLock( - key: String, + lockKey: String, + token: String, timeToLiveMillis: Long, - ): Boolean = redisTemplateForGlobalLock.opsForValue().setIfAbsent(key, key, Duration.ofMillis(timeToLiveMillis)) ?: false + ): Boolean = redisTemplateForGlobalLock.opsForValue().setIfAbsent(lockKey, token, Duration.ofMillis(timeToLiveMillis)) ?: false + + override fun globalUnLock( + lockKey: String, + token: String, + ): Boolean = redisTemplateForGlobalLock.execute(UN_LOCK_SCRIPT, listOf(lockKey), token) == 1L - override fun globalUnLock(key: String): Boolean = redisTemplateForGlobalLock.delete(key) + companion object { + /** Compare-and-delete, so an expired holder cannot release the lock of the next holder. */ + private val UN_LOCK_SCRIPT = + DefaultRedisScript( + "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end", + Long::class.javaObjectType, + ) + } } diff --git a/req-shield-spring-example/src/main/kotlin/com/linecorp/cse/reqshield/configuration/RedisConfiguration.kt b/req-shield-spring-example/src/main/kotlin/com/linecorp/cse/reqshield/configuration/RedisConfiguration.kt index 8d86718..9018bda 100644 --- a/req-shield-spring-example/src/main/kotlin/com/linecorp/cse/reqshield/configuration/RedisConfiguration.kt +++ b/req-shield-spring-example/src/main/kotlin/com/linecorp/cse/reqshield/configuration/RedisConfiguration.kt @@ -58,15 +58,11 @@ class RedisConfiguration { @Bean fun redisTemplateForGlobalLock(connectionFactory: RedisConnectionFactory): RedisTemplate { - val valueSerializer = - Jackson2JsonRedisSerializer(String::class.java).apply { - setObjectMapper(objectMapper()) - } - val redisTemplate = RedisTemplate() redisTemplate.setConnectionFactory(connectionFactory) redisTemplate.keySerializer = StringRedisSerializer() - redisTemplate.valueSerializer = valueSerializer + // The lock token is stored as a plain string so the compare-and-delete script can compare it as is + redisTemplate.valueSerializer = StringRedisSerializer() return redisTemplate } diff --git a/req-shield-spring-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/service/CacheAnnotationTest.kt b/req-shield-spring-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/service/CacheAnnotationTest.kt index 4509c35..d032e96 100644 --- a/req-shield-spring-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/service/CacheAnnotationTest.kt +++ b/req-shield-spring-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/service/CacheAnnotationTest.kt @@ -65,7 +65,7 @@ class CacheAnnotationTest : AbstractRedisTest() { await().atMost(Duration.ofMillis(BaseReqShieldTest.AWAIT_TIMEOUT)).untilAsserted { assertEquals(1, sampleService.getRequestCount()) - assertNotNull(reqShieldCache.get("product-$testProductId")) + assertNotNull(reqShieldCache.get("product::product-$testProductId")) } } @@ -86,7 +86,7 @@ class CacheAnnotationTest : AbstractRedisTest() { await().atMost(Duration.ofMillis(BaseReqShieldTest.AWAIT_TIMEOUT)).untilAsserted { assertEquals(100, sampleService.getRequestCount()) - assertNotNull(reqShieldCache.get("product-$testProductId")) + assertNotNull(reqShieldCache.get("productOnlyUpdateCache::product-$testProductId")) } } @@ -109,7 +109,7 @@ class CacheAnnotationTest : AbstractRedisTest() { await().atMost(Duration.ofMillis(BaseReqShieldTest.AWAIT_TIMEOUT)).untilAsserted { assertEquals(1, sampleService.getRequestCount()) - assertNotNull(reqShieldCache.get("product-$testProductId")) + assertNotNull(reqShieldCache.get("product::product-$testProductId")) } } @@ -120,19 +120,19 @@ class CacheAnnotationTest : AbstractRedisTest() { sampleService.getProduct(testProductId) await().atMost(5, TimeUnit.SECONDS).until { - reqShieldCache.get("product-$testProductId") != null + reqShieldCache.get("product::product-$testProductId") != null } - assertNotNull(reqShieldCache.get("product-$testProductId")) + assertNotNull(reqShieldCache.get("product::product-$testProductId")) // when sampleService.removeProduct(testProductId) // then await().atMost(5, TimeUnit.SECONDS).until { - reqShieldCache.get("product-$testProductId") == null + reqShieldCache.get("product::product-$testProductId") == null } - assertNull(reqShieldCache.get("product-$testProductId")) + assertNull(reqShieldCache.get("product::product-$testProductId")) } } diff --git a/req-shield-spring-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/service/GlobalLockTest.kt b/req-shield-spring-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/service/GlobalLockTest.kt new file mode 100644 index 0000000..78c9cf4 --- /dev/null +++ b/req-shield-spring-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/service/GlobalLockTest.kt @@ -0,0 +1,52 @@ +/* + * 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.spring.service + +import com.linecorp.cse.reqshield.spring.cache.GlobalLockSupport +import com.linecorp.cse.reqshield.support.redis.AbstractRedisTest +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.extension.ExtendWith +import org.springframework.beans.factory.annotation.Autowired +import org.springframework.boot.test.context.SpringBootTest +import org.springframework.test.context.junit.jupiter.SpringExtension +import java.util.UUID + +@SpringBootTest +@ExtendWith(SpringExtension::class) +class GlobalLockTest : AbstractRedisTest() { + @Autowired + private lateinit var globalLockSupport: GlobalLockSupport + + @Test + fun `lock is exclusive and can only be released by the token that acquired it`() { + val lockKey = "globalLockTest-${UUID.randomUUID()}" + + assertTrue(globalLockSupport.globalLock(lockKey, "ownerToken", 5000)) + // already held + assertFalse(globalLockSupport.globalLock(lockKey, "otherToken", 5000)) + // compare-and-delete: a foreign token must not release someone else's lock + assertFalse(globalLockSupport.globalUnLock(lockKey, "otherToken")) + assertFalse(globalLockSupport.globalLock(lockKey, "otherToken", 5000)) + + assertTrue(globalLockSupport.globalUnLock(lockKey, "ownerToken")) + // released, so it can be acquired again + assertTrue(globalLockSupport.globalLock(lockKey, "otherToken", 5000)) + assertTrue(globalLockSupport.globalUnLock(lockKey, "otherToken")) + } +} diff --git a/req-shield-spring-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/example/cache/AsyncCacheImpl.kt b/req-shield-spring-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/example/cache/AsyncCacheImpl.kt index 5b282c3..6581ac9 100644 --- a/req-shield-spring-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/example/cache/AsyncCacheImpl.kt +++ b/req-shield-spring-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/example/cache/AsyncCacheImpl.kt @@ -17,17 +17,26 @@ package com.linecorp.cse.reqshield.spring.webflux.example.cache import com.linecorp.cse.reqshield.spring.webflux.cache.AsyncCache +import com.linecorp.cse.reqshield.spring.webflux.cache.GlobalLockSupport import com.linecorp.cse.reqshield.support.model.ReqShieldData import org.springframework.data.redis.core.ReactiveRedisOperations +import org.springframework.data.redis.core.script.RedisScript import org.springframework.stereotype.Service import reactor.core.publisher.Mono import java.time.Duration +/** + * Compare-and-delete: the lock is only released when it still holds the caller's token, so an owner + * whose lock already expired cannot release the lock of the next owner. + */ +private const val UNLOCK_SCRIPT = "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end" + @Service class AsyncCacheImpl( private val redisOperations: ReactiveRedisOperations>, private val redisOperationsForGlobalLock: ReactiveRedisOperations, -) : AsyncCache { +) : AsyncCache, + GlobalLockSupport { override fun get(key: String): Mono?> = redisOperations.opsForValue()[key] override fun put( @@ -39,13 +48,18 @@ class AsyncCacheImpl( override fun evict(key: String): Mono = redisOperations.opsForValue().delete(key) override fun globalLock( - key: String, + lockKey: String, + token: String, timeToLiveMillis: Long, - ): Mono = - redisOperationsForGlobalLock.opsForValue().setIfAbsent(key, key, Duration.ofMillis(timeToLiveMillis)) ?: Mono.just(false) + ): Mono = redisOperationsForGlobalLock.opsForValue().setIfAbsent(lockKey, token, Duration.ofMillis(timeToLiveMillis)) - override fun globalUnLock(key: String): Mono = + override fun globalUnLock( + lockKey: String, + token: String, + ): Mono = redisOperationsForGlobalLock - .delete(key) - .map { count -> count > 0 } + .execute(RedisScript.of(UNLOCK_SCRIPT, Long::class.java), listOf(lockKey), listOf(token)) + .next() + .map { it == 1L } + .defaultIfEmpty(false) } diff --git a/req-shield-spring-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/example/configuration/RedisConfiguration.kt b/req-shield-spring-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/example/configuration/RedisConfiguration.kt index 7c3bb5e..bd4bda2 100644 --- a/req-shield-spring-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/example/configuration/RedisConfiguration.kt +++ b/req-shield-spring-webflux-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/example/configuration/RedisConfiguration.kt @@ -69,10 +69,9 @@ class RedisConfiguration { @Bean("redisOperationsForGlobalLock") fun reactiveRedisOperationsForGlobalLock(factory: ReactiveRedisConnectionFactory): ReactiveRedisOperations { val keySerializer = StringRedisSerializer() - val valueSerializer = - Jackson2JsonRedisSerializer(String::class.java).apply { - setObjectMapper(objectMapper()) - } + // The lock token has to be stored as a plain string: the compare-and-delete Lua script + // compares the stored value with the raw token it receives as a script argument. + val valueSerializer = StringRedisSerializer() val serializationContext = RedisSerializationContext .newSerializationContext(keySerializer) diff --git a/req-shield-spring-webflux-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/example/service/CacheAnnotationTest.kt b/req-shield-spring-webflux-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/example/service/CacheAnnotationTest.kt index 0af2d08..52f8b66 100644 --- a/req-shield-spring-webflux-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/example/service/CacheAnnotationTest.kt +++ b/req-shield-spring-webflux-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/example/service/CacheAnnotationTest.kt @@ -70,7 +70,7 @@ class CacheAnnotationTest : AbstractRedisTest() { }.verifyComplete() await().atMost(5, TimeUnit.SECONDS).until { - asyncCache.get("product-$testProductId").block() != null + asyncCache.get("product::product-$testProductId").block() != null } } @@ -100,7 +100,7 @@ class CacheAnnotationTest : AbstractRedisTest() { assertEquals(20, sampleService.getRequestCount(), "Request count should be 20") await().atMost(5, TimeUnit.SECONDS).until { - asyncCache.get("product-$testProductId").block() != null + asyncCache.get("productOnlyUpdataCache::product-$testProductId").block() != null } } @@ -124,7 +124,7 @@ class CacheAnnotationTest : AbstractRedisTest() { }.verifyComplete() await().atMost(5, TimeUnit.SECONDS).until { - asyncCache.get("product-$testProductId").block() != null + asyncCache.get("product::product-$testProductId").block() != null } } @@ -144,9 +144,9 @@ class CacheAnnotationTest : AbstractRedisTest() { // then await().atMost(5, TimeUnit.SECONDS).until { - asyncCache.get("product-$testProductId").block() != null + asyncCache.get("product::product-$testProductId").block() != null } - val cacheMono = asyncCache.get("product-$testProductId").block() + val cacheMono = asyncCache.get("product::product-$testProductId").block() assertNotNull(cacheMono) // when @@ -162,9 +162,9 @@ class CacheAnnotationTest : AbstractRedisTest() { // then await().atMost(5, TimeUnit.SECONDS).until { - asyncCache.get("product-$testProductId").block() == null + asyncCache.get("product::product-$testProductId").block() == null } - val cacheMonoNull = asyncCache.get("product-$testProductId").block() + val cacheMonoNull = asyncCache.get("product::product-$testProductId").block() assertNull(cacheMonoNull) } } diff --git a/req-shield-spring-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/example/cache/AsyncCacheImpl.kt b/req-shield-spring-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/example/cache/AsyncCacheImpl.kt index 5c2587d..8c17db3 100644 --- a/req-shield-spring-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/example/cache/AsyncCacheImpl.kt +++ b/req-shield-spring-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/example/cache/AsyncCacheImpl.kt @@ -17,10 +17,13 @@ package com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.example.cache import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.cache.AsyncCache +import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.cache.GlobalLockSupport import com.linecorp.cse.reqshield.support.model.ReqShieldData +import kotlinx.coroutines.reactive.awaitFirstOrNull import org.springframework.data.redis.core.ReactiveRedisOperations import org.springframework.data.redis.core.deleteAndAwait import org.springframework.data.redis.core.getAndAwait +import org.springframework.data.redis.core.script.RedisScript import org.springframework.data.redis.core.setAndAwait import org.springframework.data.redis.core.setIfAbsentAndAwait import org.springframework.stereotype.Service @@ -30,7 +33,8 @@ import java.time.Duration class AsyncCacheImpl( private val redisOperations: ReactiveRedisOperations>, private val redisOperationsForGlobalLock: ReactiveRedisOperations, -) : AsyncCache { +) : AsyncCache, + GlobalLockSupport { override suspend fun get(key: String): ReqShieldData? = redisOperations.opsForValue().getAndAwait(key) override suspend fun put( @@ -41,10 +45,34 @@ class AsyncCacheImpl( override suspend fun evict(key: String): Boolean = redisOperations.opsForValue().deleteAndAwait(key) + /** `SET lockKey token NX PX ttl`: the stored token identifies the caller that owns the lock. */ override suspend fun globalLock( - key: String, + lockKey: String, + token: String, timeToLiveMillis: Long, - ): Boolean = redisOperationsForGlobalLock.opsForValue().setIfAbsentAndAwait(key, key, Duration.ofMillis(timeToLiveMillis)) + ): Boolean = + redisOperationsForGlobalLock + .opsForValue() + .setIfAbsentAndAwait(lockKey, token, Duration.ofMillis(timeToLiveMillis)) + + /** + * Compare-and-delete in a single atomic step, so an owner whose lock already expired can never + * release the lock that the next owner has taken in the meantime. + */ + override suspend fun globalUnLock( + lockKey: String, + token: String, + ): Boolean = + redisOperationsForGlobalLock + .execute(UNLOCK_SCRIPT, listOf(lockKey), listOf(token)) + .awaitFirstOrNull() == 1L - override suspend fun globalUnLock(key: String): Boolean = redisOperationsForGlobalLock.deleteAndAwait(key) > 0 + companion object { + /** Returns the number of keys deleted: 1 when this caller still owned the lock, 0 otherwise. */ + private val UNLOCK_SCRIPT: RedisScript = + RedisScript.of( + "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end", + Long::class.java, + ) + } } diff --git a/req-shield-spring-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/example/configuration/RedisConfiguration.kt b/req-shield-spring-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/example/configuration/RedisConfiguration.kt index 1efaa32..a900aeb 100644 --- a/req-shield-spring-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/example/configuration/RedisConfiguration.kt +++ b/req-shield-spring-webflux-kotlin-coroutine-example/src/main/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/example/configuration/RedisConfiguration.kt @@ -69,10 +69,9 @@ class RedisConfiguration { @Bean("redisOperationsForGlobalLock") fun reactiveRedisOperationsForGlobalLock(factory: ReactiveRedisConnectionFactory): ReactiveRedisOperations { val keySerializer = StringRedisSerializer() - val valueSerializer = - Jackson2JsonRedisSerializer(String::class.java).apply { - setObjectMapper(objectMapper()) - } + // Lock values are raw ownership tokens, so they are stored verbatim: a JSON serializer would + // store them quoted, which only stays comparable while every access uses this same template. + val valueSerializer = StringRedisSerializer() val serializationContext = RedisSerializationContext .newSerializationContext(keySerializer) diff --git a/req-shield-spring-webflux-kotlin-coroutine-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/example/service/CacheAnnotationTest.kt b/req-shield-spring-webflux-kotlin-coroutine-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/example/service/CacheAnnotationTest.kt index 1425d28..f2910ce 100644 --- a/req-shield-spring-webflux-kotlin-coroutine-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/example/service/CacheAnnotationTest.kt +++ b/req-shield-spring-webflux-kotlin-coroutine-example/src/test/kotlin/com/linecorp/cse/reqshield/spring/webflux/kotlin/coroutine/example/service/CacheAnnotationTest.kt @@ -17,7 +17,9 @@ package com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.example.service import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.cache.AsyncCache +import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.cache.GlobalLockSupport import com.linecorp.cse.reqshield.spring.webflux.kotlin.coroutine.example.dto.Product +import com.linecorp.cse.reqshield.support.constant.ConfigValues.LOCK_KEY_PREFIX import com.linecorp.cse.reqshield.support.redis.AbstractRedisTest import kotlinx.coroutines.async import kotlinx.coroutines.awaitAll @@ -29,7 +31,10 @@ import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test import org.junit.jupiter.api.extension.ExtendWith import org.springframework.beans.factory.annotation.Autowired +import org.springframework.beans.factory.annotation.Qualifier import org.springframework.boot.test.context.SpringBootTest +import org.springframework.data.redis.core.ReactiveRedisOperations +import org.springframework.data.redis.core.getAndAwait import org.springframework.test.context.junit.jupiter.SpringExtension import java.util.UUID @@ -42,6 +47,15 @@ class CacheAnnotationTest : AbstractRedisTest() { @Autowired lateinit var asyncCache: AsyncCache + @Autowired + @Qualifier("redisOperationsForGlobalLock") + lateinit var globalLockOperations: ReactiveRedisOperations + + private val lockSupport: GlobalLockSupport get() = asyncCache as GlobalLockSupport + + /** The aspect namespaces every key with the cache name of the annotation. */ + private fun cacheKeyOf(productId: String) = "product::product-$productId" + @BeforeEach fun `reset request count`() = runTest { @@ -62,7 +76,7 @@ class CacheAnnotationTest : AbstractRedisTest() { delay(500) assertEquals(1, sampleService.getRequestCount()) - assertNotNull(asyncCache.get("product-$testProductId")) + assertNotNull(asyncCache.get(cacheKeyOf(testProductId))) } @Test @@ -107,7 +121,7 @@ class CacheAnnotationTest : AbstractRedisTest() { val maxAttempts = 30 var attempts = 0 - while (asyncCache.get("product-$testProductId") == null) { + while (asyncCache.get(cacheKeyOf(testProductId)) == null) { if (attempts >= maxAttempts) { break } @@ -115,13 +129,13 @@ class CacheAnnotationTest : AbstractRedisTest() { delay(100) } - assertNotNull(asyncCache.get("product-$testProductId")) + assertNotNull(asyncCache.get(cacheKeyOf(testProductId))) // when sampleService.removeProduct(testProductId) var attemptsSecond = 0 - while (asyncCache.get("product-$testProductId") != null) { + while (asyncCache.get(cacheKeyOf(testProductId)) != null) { if (attemptsSecond >= maxAttempts) { break } @@ -129,6 +143,41 @@ class CacheAnnotationTest : AbstractRedisTest() { delay(100) } - assertNull(asyncCache.get("product-$testProductId")) + assertNull(asyncCache.get(cacheKeyOf(testProductId))) + } + + @Test + fun `global lock test - the lock is released after the owning request completes`() = + runBlocking { + // given + val testProductId: String = UUID.randomUUID().toString() + val lockKey = "$LOCK_KEY_PREFIX${cacheKeyOf(testProductId)}_CREATE" + + // when + sampleService.getProductForGlobalLock(testProductId) + + // then the background cache write releases the lock it owns + var attempts = 0 + while (globalLockOperations.opsForValue().getAndAwait(lockKey) != null && attempts < 30) { + attempts++ + delay(100) + } + assertNull(globalLockOperations.opsForValue().getAndAwait(lockKey)) + } + + @Test + fun `global lock test - only the owning token can release the lock`() = + runBlocking { + val lockKey = "lock-token-${UUID.randomUUID()}" + + assertTrue(lockSupport.globalLock(lockKey, "owner", 10_000)) + // The token is stored as a plain string, which is what the compare-and-delete script compares. + assertEquals("owner", globalLockOperations.opsForValue().getAndAwait(lockKey)) + // SET NX: a second caller cannot take a held lock. + assertFalse(lockSupport.globalLock(lockKey, "intruder", 10_000)) + // Compare-and-delete: a non-owner cannot release it either. + assertFalse(lockSupport.globalUnLock(lockKey, "intruder")) + assertTrue(lockSupport.globalUnLock(lockKey, "owner")) + assertNull(globalLockOperations.opsForValue().getAndAwait(lockKey)) } } diff --git a/support/build.gradle.kts b/support/build.gradle.kts index 110db86..aacdc67 100644 --- a/support/build.gradle.kts +++ b/support/build.gradle.kts @@ -16,9 +16,6 @@ plugins { alias(libs.plugins.kotlin.jvm) - alias(libs.plugins.kotlin.spring) - alias(libs.plugins.spring.boot2) - alias(libs.plugins.spring.dependency.management) `java-test-fixtures` } diff --git a/support/src/main/kotlin/com/linecorp/cse/reqshield/support/constant/ConfigValues.kt b/support/src/main/kotlin/com/linecorp/cse/reqshield/support/constant/ConfigValues.kt index bb713ec..140f2c5 100644 --- a/support/src/main/kotlin/com/linecorp/cse/reqshield/support/constant/ConfigValues.kt +++ b/support/src/main/kotlin/com/linecorp/cse/reqshield/support/constant/ConfigValues.kt @@ -19,12 +19,19 @@ package com.linecorp.cse.reqshield.support.constant object ConfigValues { const val DEFAULT_LOCK_TIMEOUT_MILLIS = 3000L const val DEFAULT_DECISION_FOR_UPDATE = 80 + const val DEFAULT_TIME_TO_LIVE_MILLIS = 10 * 60 * 1000L const val LOCK_MONITOR_INTERVAL_MILLIS = 1000L const val MAX_ATTEMPT_GET_CACHE = 60 const val GET_CACHE_INTERVAL_MILLIS = 50L - const val MAX_ATTEMPT_SET_CACHE = 3 - const val SET_CACHE_RETRY_INTERVAL_MILLIS = 100L + /** + * While waiting for another request to fill the cache, this many consecutive cache-read + * failures are treated as a cache outage and the waiter falls back to the supplier at once. + */ + const val MAX_CONSECUTIVE_GET_CACHE_FAILURES = 3 + + /** Prefix applied to every lock key so lock entries can never collide with cache entries. */ + const val LOCK_KEY_PREFIX = "reqshield:lock:" } diff --git a/support/src/main/kotlin/com/linecorp/cse/reqshield/support/exception/ClientException.kt b/support/src/main/kotlin/com/linecorp/cse/reqshield/support/exception/ClientException.kt index 550fdcf..c2d07cd 100644 --- a/support/src/main/kotlin/com/linecorp/cse/reqshield/support/exception/ClientException.kt +++ b/support/src/main/kotlin/com/linecorp/cse/reqshield/support/exception/ClientException.kt @@ -17,16 +17,17 @@ package com.linecorp.cse.reqshield.support.exception import com.linecorp.cse.reqshield.support.exception.code.ErrorCode -import org.slf4j.LoggerFactory - -private val log = LoggerFactory.getLogger(ClientException::class.java) +/** + * Wraps failures raised by client-provided functions (supplier, cache getter/setter, lock functions). + * + * The original exception is chained as [cause] so callers keep the full stack trace. + * This class intentionally does not log: synchronous failures are propagated to the caller, + * and fire-and-forget paths inside ReqShield log at the point where the error is dropped. + */ class ClientException( val errorCode: ErrorCode, override val message: String = errorCode.message, - originErrorMessage: String? = null, -) : RuntimeException(message) { - init { - log.error("[Req-Shield] errorCode : {}, message : {}, originErrorMessage : {}", errorCode.code, message, originErrorMessage) - } -} + cause: Throwable? = null, + val originErrorMessage: String? = cause?.message, +) : RuntimeException(message, cause) diff --git a/support/src/test/kotlin/com/linecorp/cse/reqshield/support/exception/ClientExceptionTest.kt b/support/src/test/kotlin/com/linecorp/cse/reqshield/support/exception/ClientExceptionTest.kt new file mode 100644 index 0000000..4bac987 --- /dev/null +++ b/support/src/test/kotlin/com/linecorp/cse/reqshield/support/exception/ClientExceptionTest.kt @@ -0,0 +1,95 @@ +/* + * 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.exception + +import com.linecorp.cse.reqshield.support.exception.code.ErrorCode +import org.junit.jupiter.api.Test +import kotlin.test.assertEquals +import kotlin.test.assertNull +import kotlin.test.assertSame + +class ClientExceptionTest { + @Test + fun `message defaults to errorCode message and errorCode is exposed`() { + val exception = ClientException(errorCode = ErrorCode.SUPPLIER_ERROR) + + assertEquals(ErrorCode.SUPPLIER_ERROR, exception.errorCode) + assertEquals(ErrorCode.SUPPLIER_ERROR.message, exception.message) + } + + @Test + fun `explicit message overrides errorCode message`() { + val exception = + ClientException( + errorCode = ErrorCode.GET_CACHE_ERROR, + message = "custom message", + ) + + assertEquals("custom message", exception.message) + } + + @Test + fun `cause is chained and originErrorMessage defaults to cause message`() { + val cause = RuntimeException("original failure") + val exception = + ClientException( + errorCode = ErrorCode.SET_CACHE_ERROR, + cause = cause, + ) + + assertSame(cause, exception.cause) + assertEquals("original failure", exception.originErrorMessage) + } + + @Test + fun `explicit originErrorMessage wins over cause message`() { + val cause = RuntimeException("original failure") + val exception = + ClientException( + errorCode = ErrorCode.SET_CACHE_ERROR, + cause = cause, + originErrorMessage = "overridden origin message", + ) + + assertSame(cause, exception.cause) + assertEquals("overridden origin message", exception.originErrorMessage) + } + + @Test + fun `with neither cause nor originErrorMessage both are null`() { + val exception = ClientException(errorCode = ErrorCode.DOES_NOT_EXIST_GLOBAL_LOCK_FUNCTION) + + assertNull(exception.cause) + assertNull(exception.originErrorMessage) + } + + @Test + fun `all ErrorCode entries expose their expected code`() { + val expectedCodes = + mapOf( + ErrorCode.SUPPLIER_ERROR to "1001", + ErrorCode.GET_CACHE_ERROR to "1002", + ErrorCode.SET_CACHE_ERROR to "1003", + ErrorCode.DOES_NOT_EXIST_GLOBAL_LOCK_FUNCTION to "1004", + ErrorCode.DOES_NOT_EXIST_GLOBAL_UNLOCK_FUNCTION to "1005", + ) + + expectedCodes.forEach { (errorCode, expectedCode) -> + assertEquals(expectedCode, errorCode.code) + } + } +} diff --git a/support/src/test/kotlin/com/linecorp/cse/reqshield/support/model/ReqShieldDataTest.kt b/support/src/test/kotlin/com/linecorp/cse/reqshield/support/model/ReqShieldDataTest.kt new file mode 100644 index 0000000..28b4dd1 --- /dev/null +++ b/support/src/test/kotlin/com/linecorp/cse/reqshield/support/model/ReqShieldDataTest.kt @@ -0,0 +1,76 @@ +/* + * 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.model + +import org.junit.jupiter.api.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertNull +import kotlin.test.assertTrue + +class ReqShieldDataTest { + @Test + fun `secondary constructor sets status NEW and ttl with non-null value`() { + val before = System.currentTimeMillis() + val data = ReqShieldData(value = "cached-value", timeToLiveMillis = 5000L) + val after = System.currentTimeMillis() + + assertEquals(ReqShieldData.Status.NEW, data.status) + assertEquals(5000L, data.timeToLiveMillis) + assertEquals("cached-value", data.value) + assertTrue(data.createdAt in (before - 1000)..(after + 1000)) + } + + @Test + fun `secondary constructor allows null value`() { + val data = ReqShieldData(timeToLiveMillis = 1000L) + + assertNull(data.value) + assertEquals(ReqShieldData.Status.NEW, data.status) + } + + @Test + fun `primary constructor keeps explicit status and createdAt`() { + val data = + ReqShieldData( + value = "value", + status = ReqShieldData.Status.NORMAL, + createdAt = 12345L, + timeToLiveMillis = 999L, + ) + + assertEquals(ReqShieldData.Status.NORMAL, data.status) + assertEquals(12345L, data.createdAt) + assertEquals(999L, data.timeToLiveMillis) + } + + @Test + fun `data class equality holds for identical fields`() { + val first = ReqShieldData(value = "same", status = ReqShieldData.Status.CREATING, createdAt = 1L, timeToLiveMillis = 2L) + val second = ReqShieldData(value = "same", status = ReqShieldData.Status.CREATING, createdAt = 1L, timeToLiveMillis = 2L) + + assertEquals(first, second) + } + + @Test + fun `data class equality fails when value differs`() { + val first = ReqShieldData(value = "one", status = ReqShieldData.Status.CREATING, createdAt = 1L, timeToLiveMillis = 2L) + val second = ReqShieldData(value = "two", status = ReqShieldData.Status.CREATING, createdAt = 1L, timeToLiveMillis = 2L) + + assertFalse(first == second) + } +} diff --git a/support/src/test/kotlin/com/linecorp/cse/reqshield/support/utils/TimeUtilsTest.kt b/support/src/test/kotlin/com/linecorp/cse/reqshield/support/utils/TimeUtilsTest.kt new file mode 100644 index 0000000..764627f --- /dev/null +++ b/support/src/test/kotlin/com/linecorp/cse/reqshield/support/utils/TimeUtilsTest.kt @@ -0,0 +1,40 @@ +/* + * 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 + +import org.junit.jupiter.api.Test +import kotlin.test.assertTrue + +class TimeUtilsTest { + @Test + fun `nowToEpochTime is close to System currentTimeMillis`() { + val before = System.currentTimeMillis() + val epochTime = nowToEpochTime() + val after = System.currentTimeMillis() + + assertTrue(epochTime in (before - 1000)..(after + 1000)) + } + + @Test + fun `nowToEpochTime is monotonic non-decreasing across calls`() { + val first = nowToEpochTime() + Thread.sleep(10) + val second = nowToEpochTime() + + assertTrue(second >= first) + } +} diff --git a/support/src/test/kotlin/com/linecorp/cse/reqshield/utils/support/CommonUtilsTest.kt b/support/src/test/kotlin/com/linecorp/cse/reqshield/utils/support/CommonUtilsTest.kt index 8c5c1c8..ae0b33f 100644 --- a/support/src/test/kotlin/com/linecorp/cse/reqshield/utils/support/CommonUtilsTest.kt +++ b/support/src/test/kotlin/com/linecorp/cse/reqshield/utils/support/CommonUtilsTest.kt @@ -42,4 +42,40 @@ class CommonUtilsTest { val decide = decideToUpdateCache(createdAt.toInstant().toEpochMilli(), expireTime, 80) assertFalse(decide) } + + @Test + fun `elapsed well past the threshold returns true`() { + val ttl = Duration.ofMinutes(1).toMillis() + // decisionForUpdate 50% of a 1-minute ttl is 30s, use a generous margin past that boundary. + val createdAt = System.currentTimeMillis() - Duration.ofSeconds(45).toMillis() + + val decide = decideToUpdateCache(createdAt, ttl, 50) + assertTrue(decide) + } + + @Test + fun `decisionForUpdate of zero returns true immediately`() { + val ttl = Duration.ofMinutes(10).toMillis() + val createdAt = System.currentTimeMillis() + + val decide = decideToUpdateCache(createdAt, ttl, 0) + assertTrue(decide) + } + + @Test + fun `decisionForUpdate of 100 with fresh createdAt returns false`() { + val ttl = Duration.ofMinutes(10).toMillis() + val createdAt = System.currentTimeMillis() + + val decide = decideToUpdateCache(createdAt, ttl, 100) + assertFalse(decide) + } + + @Test + fun `ttl of zero returns true`() { + val createdAt = System.currentTimeMillis() + + val decide = decideToUpdateCache(createdAt, 0L, 80) + assertTrue(decide) + } }