issue-52: Refactor ReqShield to use Executor instead of ScheduledExec…

.gitignore

@@ -40,3 +40,5 @@ bin/
 ### Mac OS ###
 .DS_Store
 /.idea/
+
+.cross-memory.toml

README.md

@@ -116,8 +116,9 @@ refresh dependencies with `./gradlew build --refresh-dependencies`.
 
 ### 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` accepts any `java.util.concurrent.Executor` for its background cache writes; only `execute` is called and the
+  library never shuts a caller-supplied pool down. The default is a shared daemon pool; the Spring adapter exposes it as
+  the `reqShieldExecutor` bean (an `ExecutorService` the context shuts down), 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

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

@@ -42,12 +42,12 @@ 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
+import java.util.concurrent.Executor
 
 @Aspect
 class ReqShieldAspect<T>(
     private val reqShieldCache: ReqShieldCache<T>,
-    @Qualifier("reqShieldExecutor") private val executor: ScheduledExecutorService,
+    @Qualifier("reqShieldExecutor") private val executor: Executor,
 ) : BeanFactoryAware {
     private lateinit var beanFactory: BeanFactory
     private val spelParser = SpelExpressionParser()

core-spring/src/main/kotlin/com/linecorp/cse/reqshield/spring/config/LibAutoConfiguration.kt

@@ -21,8 +21,8 @@ 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.ExecutorService
 import java.util.concurrent.Executors
-import java.util.concurrent.ScheduledExecutorService
 import java.util.concurrent.atomic.AtomicLong
 
 @Configuration
@@ -31,13 +31,14 @@ import java.util.concurrent.atomic.AtomicLong
 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.
+     * asynchronous cache writes.
      *
-     * 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.
+     * Declared as an [ExecutorService] because this pool is owned by the context: Spring's inferred
+     * destroy method calls [ExecutorService.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 {
+    open fun reqShieldExecutor(): ExecutorService {
         val threadCounter = AtomicLong(0)
 
         return Executors.newScheduledThreadPool(

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

@@ -35,8 +35,8 @@ import org.springframework.context.annotation.Configuration
 import java.time.Duration
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.CountDownLatch
+import java.util.concurrent.ExecutorService
 import java.util.concurrent.Executors
-import java.util.concurrent.ScheduledExecutorService
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicInteger
 
@@ -65,7 +65,7 @@ class ReqShieldAspectIntegrationTest {
 
     @Test
     fun executorBeanShouldBeProvidedByTheAutoConfiguration() {
-        assertNotNull(context.getBean("reqShieldExecutor", ScheduledExecutorService::class.java))
+        assertNotNull(context.getBean("reqShieldExecutor", ExecutorService::class.java))
     }
 
     @Test

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

@@ -28,12 +28,6 @@ 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)
 
@@ -139,39 +133,51 @@ class ReqShield<T>(
     /**
      * 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.
+     * Polling runs on the caller's thread, which is blocked for the duration of the wait either way.
+     * Keeping it here gives the wait a single termination condition and means a saturated executor
+     * cannot stall it. The supplier is never called from the polling loop - it is called on this
+     * thread only after the wait gave up, so at most one extra supplier call per waiting request
+     * happens.
+     *
+     * The wait ends after [ReqShieldConfiguration.maxAttemptGetCache] polls, or earlier when
+     * [MAX_CONSECUTIVE_GET_CACHE_FAILURES] reads failed in a row (the cache looks unavailable).
+     * A poll whose cache read fails still counts as an attempt, so the wait stays bounded.
      */
     private fun handleLockFailure(
         key: String,
         callable: Callable<T?>,
         timeToLiveMillis: Long,
     ): ReqShieldData<T> {
-        val future = CompletableFuture<ReqShieldData<T>?>()
-        val scheduled = scheduleTask(reqShieldConfig.executor, future, reqShieldConfig.getCacheFunction, key)
+        var attempts = 0
+        var consecutiveFailures = 0
 
-        // 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
+        while (attempts < reqShieldConfig.maxAttemptGetCache) {
+            attempts++
+            sleepBetweenPolls(key)
 
-        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)
+                val cachedData = reqShieldConfig.getCacheFunction.invoke(key)
+                if (cachedData != null) return cachedData
+
+                consecutiveFailures = 0
+            } catch (e: Exception) {
+                log.warn("Cache read failed while waiting for the cache to be created for key '{}'", key, e)
+                if (++consecutiveFailures >= MAX_CONSECUTIVE_GET_CACHE_FAILURES) break
             }
+        }
 
         // No lock was acquired by this request, so there is nothing to release on failure
-        return cachedData ?: buildReqShieldData(executeCallable(callable, key, null, null), timeToLiveMillis)
+        return buildReqShieldData(executeCallable(callable, key, null, null), timeToLiveMillis)
+    }
+
+    private fun sleepBetweenPolls(key: String) {
+        try {
+            Thread.sleep(GET_CACHE_INTERVAL_MILLIS)
+        } catch (e: InterruptedException) {
+            Thread.currentThread().interrupt()
+            log.warn("Interrupted while waiting for the cache to be created for key '{}'", key)
+            throw ClientException(ErrorCode.GET_CACHE_ERROR, cause = e)
+        }
     }
 
     private fun buildReqShieldData(
@@ -183,54 +189,6 @@ class ReqShield<T>(
             timeToLiveMillis = timeToLiveMillis,
         )
 
-    /**
-     * 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<ReqShieldData<T>?>,
-        cacheGetter: (String) -> ReqShieldData<T>?,
-        key: String,
-    ): ScheduledFuture<*> {
-        val attemptCount = AtomicInteger(0)
-        val consecutiveFailureCount = AtomicInteger(0)
-
-        val scheduled: ScheduledFuture<*> =
-            executor.scheduleWithFixedDelay({
-                // Early exit if future is already completed to avoid unnecessary work
-                if (future.isDone) {
-                    return@scheduleWithFixedDelay
-                }
-
-                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
-                    }
-
-                    consecutiveFailureCount.set(0)
-                    if (attemptCount.incrementAndGet() >= reqShieldConfig.maxAttemptGetCache) {
-                        future.complete(null)
-                    }
-                } catch (e: Exception) {
-                    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(
         getFunction: (String) -> ReqShieldData<T>?,
         key: String,

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

@@ -24,6 +24,7 @@ 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 java.util.concurrent.Executor
 import java.util.concurrent.Executors
 import java.util.concurrent.ScheduledExecutorService
 import java.util.concurrent.atomic.AtomicLong
@@ -46,10 +47,11 @@ data class ReqShieldConfiguration<T>(
     val isLocalLock: Boolean = true,
     val lockTimeoutMillis: Long = DEFAULT_LOCK_TIMEOUT_MILLIS,
     /**
-     * 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.
+     * Executor used for the asynchronous cache writes. Only [Executor.execute] is called, so any
+     * pool works, and the library never shuts the pool down - a caller-supplied one stays the
+     * caller's to manage. Defaults to a single pool shared by every configuration instance.
      */
-    val executor: ScheduledExecutorService = sharedExecutor,
+    val executor: Executor = sharedExecutor,
     val decisionForUpdate: Int = DEFAULT_DECISION_FOR_UPDATE,
     val keyLock: KeyLock = defaultKeyLock(isLocalLock, globalLockFunction, globalUnLockFunction, lockTimeoutMillis),
     val maxAttemptGetCache: Int = MAX_ATTEMPT_GET_CACHE,

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

@@ -20,7 +20,6 @@ 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.LOCK_KEY_PREFIX
 import com.linecorp.cse.reqshield.support.constant.ConfigValues.MAX_CONSECUTIVE_GET_CACHE_FAILURES
 import com.linecorp.cse.reqshield.support.exception.ClientException
@@ -759,7 +758,9 @@ class ReqShieldTest : BaseReqShieldTest {
 
     @Test
     fun `should keep waiting when cache read failures are not consecutive`() {
-        val maxAttemptGetCache = 5
+        // 9 polls alternate empty/failed, so without the reset the counter would reach the
+        // 3-failure limit on the 6th poll and the wait would stop early.
+        val maxAttemptGetCache = 9
         val reqShieldSmallAttempt =
             ReqShield(
                 ReqShieldConfiguration(
@@ -784,11 +785,8 @@ class ReqShieldTest : BaseReqShieldTest {
 
         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(
-            getCount >= 1 + maxAttemptGetCache * 2 - 1,
-            "Expected at least ${1 + maxAttemptGetCache * 2 - 1} cache reads but was $getCount",
-        )
+        // A failed read counts as an attempt, so the wait is bounded by maxAttemptGetCache polls
+        assertEquals(1 + maxAttemptGetCache, getCount)
     }
 
     @Test
@@ -819,23 +817,20 @@ class ReqShieldTest : BaseReqShieldTest {
     }
 
     @Test
-    fun shouldCancelQueuedPollingWhenWaitingTimesOut() {
+    fun `should keep polling on the caller thread even when the executor is saturated`() {
         val executor = Executors.newSingleThreadScheduledExecutor()
         val releaseExecutor = CountDownLatch(1)
         val executorBlocked = CountDownLatch(1)
+        val cachedData = freshData(value)
         val reads = AtomicInteger()
         every { keyLock.tryLock(key, LockType.CREATE) } returns null
         val shield =
             ReqShield(
                 ReqShieldConfiguration(
                     setCacheFunction = cacheSetter,
-                    getCacheFunction = {
-                        reads.incrementAndGet()
-                        null
-                    },
+                    getCacheFunction = { if (reads.incrementAndGet() >= 3) cachedData else null },
                     keyLock = keyLock,
                     executor = executor,
-                    maxAttemptGetCache = 1,
                 ),
             )
 
@@ -848,12 +843,9 @@ class ReqShieldTest : BaseReqShieldTest {
 
             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")
+            // The wait runs on the caller thread, so a busy executor cannot stop it from seeing the write
+            assertEquals(cachedData, result)
+            verify(exactly = 0) { callable.call() }
         } finally {
             releaseExecutor.countDown()
             executor.shutdownNow()