diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 24d2478..c137833 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -3,7 +3,7 @@ name: build on: push: branches: [ "master" ] - pull_request: + pull_request_target: branches: [ "master" ] permissions: diff --git a/README.md b/README.md index 2572110..fadeac5 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # Krate Krate is a rate limiter library for Kotlin. -[![Kotlin](https://img.shields.io/badge/kotlin-1.8.0-blue.svg?logo=kotlin)](http://kotlinlang.org) +[![Kotlin](https://img.shields.io/badge/kotlin-1.9.21-blue.svg?logo=kotlin)](http://kotlinlang.org) ![GitHub](https://img.shields.io/github/license/lpicanco/krate) ![Build Status](https://img.shields.io/github/actions/workflow/status/lpicanco/krate/jdk11.yml?branch=master) [![Quality Gate Status](https://sonarcloud.io/api/project_badges/measure?project=lpicanco_krate&metric=alert_status)](https://sonarcloud.io/summary/new_code?id=lpicanco_krate) diff --git a/build.gradle.kts b/build.gradle.kts index a4d26bd..fd90c18 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -2,7 +2,7 @@ import org.jetbrains.kotlin.gradle.tasks.KotlinCompile plugins { kotlin("jvm") version "1.9.21" - id("org.jlleitschuh.gradle.ktlint") version "10.3.0" + id("org.jlleitschuh.gradle.ktlint") version "12.0.3" id("org.sonarqube") version "4.4.1.3373" jacoco id("maven-publish") @@ -52,7 +52,7 @@ subprojects { dependencies { implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.7.3") testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-test:1.7.3") - testImplementation("io.mockk:mockk:1.13.8") + testImplementation("io.mockk:mockk:1.13.9") testImplementation("org.junit.jupiter:junit-jupiter-api:5.10.1") testRuntimeOnly("org.junit.jupiter:junit-jupiter-engine:5.10.1") } diff --git a/krate-core/build.gradle.kts b/krate-core/build.gradle.kts index e69de29..8337712 100644 --- a/krate-core/build.gradle.kts +++ b/krate-core/build.gradle.kts @@ -0,0 +1 @@ +// diff --git a/krate-core/src/main/kotlin/com/neutrine/krate/RateLimiterBuilder.kt b/krate-core/src/main/kotlin/com/neutrine/krate/RateLimiterBuilder.kt index 644f9e6..6c1b033 100644 --- a/krate-core/src/main/kotlin/com/neutrine/krate/RateLimiterBuilder.kt +++ b/krate-core/src/main/kotlin/com/neutrine/krate/RateLimiterBuilder.kt @@ -35,7 +35,6 @@ import kotlin.math.roundToLong * @param maxRate the maximum rate at which tokens can be consumed */ class RateLimiterBuilder(private val maxRate: Long) { - /** * The maximum number of tokens that can be stored in the bucket. Defaults to [maxRate]. */ @@ -63,7 +62,7 @@ class RateLimiterBuilder(private val maxRate: Long) { capacity = maxBurst, refillTokenInterval = Duration.ofMillis(refillTokenIntervalInMillis.roundToLong()), stateStorage = stateStorage, - clock = clock + clock = clock, ) } } @@ -86,6 +85,9 @@ class RateLimiterBuilder(private val maxRate: Long) { * } * } */ -fun rateLimiter(maxRate: Long, init: RateLimiterBuilder.() -> Unit = {}): RateLimiter { +fun rateLimiter( + maxRate: Long, + init: RateLimiterBuilder.() -> Unit = {}, +): RateLimiter { return RateLimiterBuilder(maxRate).apply(init).build() } diff --git a/krate-core/src/main/kotlin/com/neutrine/krate/algorithms/TokenBucketLimiter.kt b/krate-core/src/main/kotlin/com/neutrine/krate/algorithms/TokenBucketLimiter.kt index 912feb4..2c02ab1 100644 --- a/krate-core/src/main/kotlin/com/neutrine/krate/algorithms/TokenBucketLimiter.kt +++ b/krate-core/src/main/kotlin/com/neutrine/krate/algorithms/TokenBucketLimiter.kt @@ -44,7 +44,7 @@ class TokenBucketLimiter( val capacity: Long, val refillTokenInterval: Duration, private val clock: Clock, - private val stateStorage: StateStorage = MemoryStateStorage() + private val stateStorage: StateStorage = MemoryStateStorage(), ) : RateLimiter { override suspend fun tryTake(): Boolean { return tryTakeFromState(null) @@ -80,7 +80,7 @@ class TokenBucketLimiter( current.copy( remainingTokens = max(0, totalTokens - 1), - lastUpdated = lastUpdated + lastUpdated = lastUpdated, ) } } @@ -100,5 +100,5 @@ class TokenBucketLimiter( */ data class BucketState( val remainingTokens: Long, - val lastUpdated: Instant + val lastUpdated: Instant, ) diff --git a/krate-core/src/main/kotlin/com/neutrine/krate/storage/StateStorage.kt b/krate-core/src/main/kotlin/com/neutrine/krate/storage/StateStorage.kt index 32e0ca1..0102c79 100644 --- a/krate-core/src/main/kotlin/com/neutrine/krate/storage/StateStorage.kt +++ b/krate-core/src/main/kotlin/com/neutrine/krate/storage/StateStorage.kt @@ -30,7 +30,6 @@ import kotlin.time.toKotlinDuration * A state storage is a component that allows to store the state of a rate limiter. */ interface StateStorage { - /** * Get the state of the bucket for the given [key]. * @param key the key to use to get the state @@ -43,10 +42,12 @@ interface StateStorage { * @param key the key to use to set the state * @param compareAndSetFunction a Compare-And-Set function that takes the current state of the bucket and returns the new state */ - suspend fun compareAndSet(key: String, compareAndSetFunction: (current: BucketState?) -> BucketState) + suspend fun compareAndSet( + key: String, + compareAndSetFunction: (current: BucketState?) -> BucketState, + ) companion object { - // Default expiration check interval val DEFAULT_RETRY_DELAY = Duration.ofMillis(100L).toKotlinDuration() } diff --git a/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/BucketStateMap.kt b/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/BucketStateMap.kt index fba5618..6b5fc17 100644 --- a/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/BucketStateMap.kt +++ b/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/BucketStateMap.kt @@ -28,7 +28,6 @@ import java.util.concurrent.atomic.AtomicReference * Interface the provides a map implementation to store the state of the rate limiter. */ interface BucketStateMap { - /** * Returns the [AtomicReference] of the [BucketState] for the given [key]. * If the [key] is not present in the map, it returns `null`. @@ -40,5 +39,8 @@ interface BucketStateMap { * If the [key] is already present, it returns the [AtomicReference] of the [BucketState] for the given [key]. * If the [key] is not present in the map, it returns `null`. */ - fun putIfAbsent(key: String, value: AtomicReference): AtomicReference? + fun putIfAbsent( + key: String, + value: AtomicReference, + ): AtomicReference? } diff --git a/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorage.kt b/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorage.kt index 671ac4f..6d27f9c 100644 --- a/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorage.kt +++ b/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorage.kt @@ -32,13 +32,16 @@ import java.util.concurrent.atomic.AtomicReference * @param bucketStateMap The [BucketStateMap] to use to store the state. */ class MemoryStateStorage( - private val bucketStateMap: BucketStateMap = SimpleBucketStateMap() + private val bucketStateMap: BucketStateMap = SimpleBucketStateMap(), ) : StateStorage { override fun getBucketState(key: String): BucketState? { return bucketStateMap.getBucketStateReference(key)?.get() } - override suspend fun compareAndSet(key: String, compareAndSetFunction: (current: BucketState?) -> BucketState) { + override suspend fun compareAndSet( + key: String, + compareAndSetFunction: (current: BucketState?) -> BucketState, + ) { val currentState = bucketStateMap.getBucketStateReference(key) val currentStateValue = currentState?.get() val newStateValue = compareAndSetFunction(currentStateValue) diff --git a/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorageBuilder.kt b/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorageBuilder.kt index 2f9b865..84d54fa 100644 --- a/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorageBuilder.kt +++ b/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorageBuilder.kt @@ -40,7 +40,6 @@ class MemoryStateStorageBuilder { * A builder for [MemoryStateStorage] instances with eviction. */ class MemoryStateStorageWithEvictionBuilder { - /** * The time-to-live after the last access. */ @@ -66,15 +65,16 @@ class MemoryStateStorageWithEvictionBuilder { */ var coroutineScope: CoroutineScope = CoroutineScope(Dispatchers.Default) - fun build() = MemoryStateStorage( - SimpleBucketStateMapWithEviction( - clock, - bucketStateMap, - ttlAfterLastAccess, - expirationCheckInterval, - coroutineScope + fun build() = + MemoryStateStorage( + SimpleBucketStateMapWithEviction( + clock, + bucketStateMap, + ttlAfterLastAccess, + expirationCheckInterval, + coroutineScope, + ), ) - ) } /** diff --git a/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/SimpleBucketStateMap.kt b/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/SimpleBucketStateMap.kt index 0228c38..f6d58eb 100644 --- a/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/SimpleBucketStateMap.kt +++ b/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/SimpleBucketStateMap.kt @@ -30,14 +30,16 @@ import java.util.concurrent.atomic.AtomicReference * This implementation is thread-safe and uses [ConcurrentHashMap] to store the state. */ class SimpleBucketStateMap : BucketStateMap { - internal val state: ConcurrentHashMap> = ConcurrentHashMap() override fun getBucketStateReference(key: String): AtomicReference? { return state[key] } - override fun putIfAbsent(key: String, value: AtomicReference): AtomicReference? { + override fun putIfAbsent( + key: String, + value: AtomicReference, + ): AtomicReference? { return state.putIfAbsent(key, value) } } diff --git a/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/SimpleBucketStateMapWithEviction.kt b/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/SimpleBucketStateMapWithEviction.kt index 90fb0d3..84ba4ac 100644 --- a/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/SimpleBucketStateMapWithEviction.kt +++ b/krate-core/src/main/kotlin/com/neutrine/krate/storage/memory/SimpleBucketStateMapWithEviction.kt @@ -43,9 +43,8 @@ class SimpleBucketStateMapWithEviction( private val stateMap: SimpleBucketStateMap = SimpleBucketStateMap(), private val ttlAfterLastAccess: Duration = 2.hours, expirationCheckInterval: Duration = 10.minutes, - coroutineScope: CoroutineScope = CoroutineScope(Dispatchers.Default) + coroutineScope: CoroutineScope = CoroutineScope(Dispatchers.Default), ) : BucketStateMap by stateMap { - init { coroutineScope.launch { while (true) { diff --git a/krate-core/src/test/kotlin/com/neutrine/krate/RateLimiterBuilderTest.kt b/krate-core/src/test/kotlin/com/neutrine/krate/RateLimiterBuilderTest.kt index e76a47c..d232c14 100644 --- a/krate-core/src/test/kotlin/com/neutrine/krate/RateLimiterBuilderTest.kt +++ b/krate-core/src/test/kotlin/com/neutrine/krate/RateLimiterBuilderTest.kt @@ -12,13 +12,13 @@ import java.time.Duration import java.time.temporal.ChronoUnit internal class RateLimiterBuilderTest { - @Test fun `should return an instance of TokenBucketLimiter with a rate of 60 per minute`() { - val rateLimiter = rateLimiter(maxRate = 60) { - maxBurst = 70 - maxRateTimeUnit = ChronoUnit.MINUTES - } + val rateLimiter = + rateLimiter(maxRate = 60) { + maxBurst = 70 + maxRateTimeUnit = ChronoUnit.MINUTES + } assertTrue(rateLimiter is TokenBucketLimiter) val tokenBucketLimiter = rateLimiter as TokenBucketLimiter @@ -29,10 +29,11 @@ internal class RateLimiterBuilderTest { @Test fun `should return an instance of TokenBucketLimiter with a rate of 5 per minute`() { - val rateLimiter = rateLimiter(maxRate = 5) { - maxBurst = 70 - maxRateTimeUnit = ChronoUnit.MINUTES - } + val rateLimiter = + rateLimiter(maxRate = 5) { + maxBurst = 70 + maxRateTimeUnit = ChronoUnit.MINUTES + } assertTrue(rateLimiter is TokenBucketLimiter) val tokenBucketLimiter = rateLimiter as TokenBucketLimiter @@ -43,10 +44,11 @@ internal class RateLimiterBuilderTest { @Test fun `should return an instance of TokenBucketLimiter with a rate of 5 per second`() { - val rateLimiter = rateLimiter(maxRate = 5) { - maxBurst = 10 - maxRateTimeUnit = ChronoUnit.SECONDS - } + val rateLimiter = + rateLimiter(maxRate = 5) { + maxBurst = 10 + maxRateTimeUnit = ChronoUnit.SECONDS + } assertTrue(rateLimiter is TokenBucketLimiter) val tokenBucketLimiter = rateLimiter as TokenBucketLimiter @@ -56,11 +58,12 @@ internal class RateLimiterBuilderTest { } @Test - fun `should return an instance of TokenBucketLimiter a rate of 30 per second`() { - val rateLimiter = rateLimiter(maxRate = 30) { - maxBurst = 50 - maxRateTimeUnit = ChronoUnit.SECONDS - } + fun `should return an instance of TokenBucketLimiter with a rate of 30 per second`() { + val rateLimiter = + rateLimiter(maxRate = 30) { + maxBurst = 50 + maxRateTimeUnit = ChronoUnit.SECONDS + } assertTrue(rateLimiter is TokenBucketLimiter) val tokenBucketLimiter = rateLimiter as TokenBucketLimiter @@ -70,11 +73,12 @@ internal class RateLimiterBuilderTest { } @Test - fun `should return an instance of TokenBucketLimiter a rate of 5 per hour`() { - val rateLimiter = rateLimiter(maxRate = 5) { - maxBurst = 5 - maxRateTimeUnit = ChronoUnit.HOURS - } + fun `should return an instance of TokenBucketLimiter with a rate of 5 per hour`() { + val rateLimiter = + rateLimiter(maxRate = 5) { + maxBurst = 5 + maxRateTimeUnit = ChronoUnit.HOURS + } assertTrue(rateLimiter is TokenBucketLimiter) val tokenBucketLimiter = rateLimiter as TokenBucketLimiter @@ -84,16 +88,18 @@ internal class RateLimiterBuilderTest { } @Test - fun `should return an instance of TokenBucketLimiter with custom state storage`() = runTest { - val customStateStorage = mockk(relaxed = true) - val rateLimiter = rateLimiter(maxRate = 5) { - stateStorage = customStateStorage + fun `should return an instance of TokenBucketLimiter with custom state storage`() = + runTest { + val customStateStorage = mockk(relaxed = true) + val rateLimiter = + rateLimiter(maxRate = 5) { + stateStorage = customStateStorage + } + + assertTrue(rateLimiter is TokenBucketLimiter) + val tokenBucketLimiter = rateLimiter as TokenBucketLimiter + + tokenBucketLimiter.tryTake("42") + coVerify { customStateStorage.compareAndSet("42", any()) } } - - assertTrue(rateLimiter is TokenBucketLimiter) - val tokenBucketLimiter = rateLimiter as TokenBucketLimiter - - tokenBucketLimiter.tryTake("42") - coVerify { customStateStorage.compareAndSet("42", any()) } - } } diff --git a/krate-core/src/test/kotlin/com/neutrine/krate/algorithms/TokenBucketLimiterTest.kt b/krate-core/src/test/kotlin/com/neutrine/krate/algorithms/TokenBucketLimiterTest.kt index fd71d17..0cb637d 100644 --- a/krate-core/src/test/kotlin/com/neutrine/krate/algorithms/TokenBucketLimiterTest.kt +++ b/krate-core/src/test/kotlin/com/neutrine/krate/algorithms/TokenBucketLimiterTest.kt @@ -27,133 +27,143 @@ internal class TokenBucketLimiterTest { @Nested inner class TryTakeTest { @Test - fun `should return true if there are remaining tokens`() = runTest { - val tokenBucket = TokenBucketLimiter(1, Duration.ofMinutes(1), clock) - assertTrue(tokenBucket.tryTake()) - } + fun `should return true if there are remaining tokens`() = + runTest { + val tokenBucket = TokenBucketLimiter(1, Duration.ofMinutes(1), clock) + assertTrue(tokenBucket.tryTake()) + } @Test - fun `should return false if there are no remaining tokens`() = runTest { - val tokenBucket = TokenBucketLimiter(1, Duration.ofMinutes(1), clock) - assertTrue(tokenBucket.tryTake()) - currentTime = currentTime.plusSeconds(10) - assertFalse(tokenBucket.tryTake()) - } + fun `should return false if there are no remaining tokens`() = + runTest { + val tokenBucket = TokenBucketLimiter(1, Duration.ofMinutes(1), clock) + assertTrue(tokenBucket.tryTake()) + currentTime = currentTime.plusSeconds(10) + assertFalse(tokenBucket.tryTake()) + } @Test - fun `should restore all tokens after the full refill interval`() = runTest { - val tokenBucket = TokenBucketLimiter(2L, Duration.ofMinutes(1), clock = clock) + fun `should restore all tokens after the full refill interval`() = + runTest { + val tokenBucket = TokenBucketLimiter(2L, Duration.ofMinutes(1), clock = clock) - assertTrue(tokenBucket.tryTake()) - assertTrue(tokenBucket.tryTake()) - assertFalse(tokenBucket.tryTake()) + assertTrue(tokenBucket.tryTake()) + assertTrue(tokenBucket.tryTake()) + assertFalse(tokenBucket.tryTake()) - currentTime = currentTime.plus(Duration.ofMinutes(2)) + currentTime = currentTime.plus(Duration.ofMinutes(2)) - assertTrue(tokenBucket.tryTake()) - assertTrue(tokenBucket.tryTake()) - assertFalse(tokenBucket.tryTake()) - } + assertTrue(tokenBucket.tryTake()) + assertTrue(tokenBucket.tryTake()) + assertFalse(tokenBucket.tryTake()) + } @Test - fun `should restore one token after refill interval`() = runTest { - val tokenBucket = TokenBucketLimiter(2L, Duration.ofMinutes(1), clock = clock) + fun `should restore one token after refill interval`() = + runTest { + val tokenBucket = TokenBucketLimiter(2L, Duration.ofMinutes(1), clock = clock) - assertTrue(tokenBucket.tryTake()) - assertTrue(tokenBucket.tryTake()) - assertFalse(tokenBucket.tryTake()) + assertTrue(tokenBucket.tryTake()) + assertTrue(tokenBucket.tryTake()) + assertFalse(tokenBucket.tryTake()) - currentTime = currentTime.plusSeconds(60) + currentTime = currentTime.plusSeconds(60) - assertTrue(tokenBucket.tryTake()) - assertFalse(tokenBucket.tryTake()) - } + assertTrue(tokenBucket.tryTake()) + assertFalse(tokenBucket.tryTake()) + } @Test - fun `should return true if taking between refill intervals`() = runTest { - val tokenBucket = TokenBucketLimiter(3L, Duration.ofMinutes(1), clock = clock) - - assertTrue(tokenBucket.tryTake()) - currentTime = currentTime.plusSeconds(59) - assertTrue(tokenBucket.tryTake()) - assertTrue(tokenBucket.tryTake()) - assertFalse(tokenBucket.tryTake()) - - currentTime = currentTime.plusSeconds(1) - - assertTrue(tokenBucket.tryTake()) - assertFalse(tokenBucket.tryTake()) - - currentTime = currentTime.plus(Duration.ofMinutes(1)) - assertTrue(tokenBucket.tryTake()) - assertFalse(tokenBucket.tryTake()) - - currentTime = currentTime.plus(Duration.ofMinutes(2)) - assertTrue(tokenBucket.tryTake()) - assertTrue(tokenBucket.tryTake()) - assertFalse(tokenBucket.tryTake()) - - currentTime = currentTime.plus(Duration.ofMinutes(3)) - assertTrue(tokenBucket.tryTake()) - assertTrue(tokenBucket.tryTake()) - assertTrue(tokenBucket.tryTake()) - assertFalse(tokenBucket.tryTake()) - } + fun `should return true if taking between refill intervals`() = + runTest { + val tokenBucket = TokenBucketLimiter(3L, Duration.ofMinutes(1), clock = clock) + + assertTrue(tokenBucket.tryTake()) + currentTime = currentTime.plusSeconds(59) + assertTrue(tokenBucket.tryTake()) + assertTrue(tokenBucket.tryTake()) + assertFalse(tokenBucket.tryTake()) + + currentTime = currentTime.plusSeconds(1) + + assertTrue(tokenBucket.tryTake()) + assertFalse(tokenBucket.tryTake()) + + currentTime = currentTime.plus(Duration.ofMinutes(1)) + assertTrue(tokenBucket.tryTake()) + assertFalse(tokenBucket.tryTake()) + + currentTime = currentTime.plus(Duration.ofMinutes(2)) + assertTrue(tokenBucket.tryTake()) + assertTrue(tokenBucket.tryTake()) + assertFalse(tokenBucket.tryTake()) + + currentTime = currentTime.plus(Duration.ofMinutes(3)) + assertTrue(tokenBucket.tryTake()) + assertTrue(tokenBucket.tryTake()) + assertTrue(tokenBucket.tryTake()) + assertFalse(tokenBucket.tryTake()) + } @Test - fun `should return true if there are remaining tokens by key`() = runTest { - val tokenBucket = TokenBucketLimiter(1, Duration.ofMinutes(1), clock = clock) - assertTrue(tokenBucket.tryTake("42")) - assertTrue(tokenBucket.tryTake("43")) - } + fun `should return true if there are remaining tokens by key`() = + runTest { + val tokenBucket = TokenBucketLimiter(1, Duration.ofMinutes(1), clock = clock) + assertTrue(tokenBucket.tryTake("42")) + assertTrue(tokenBucket.tryTake("43")) + } @Test - fun `should return false if there are no remaining tokens by key`() = runTest { - val tokenBucket = TokenBucketLimiter(1, Duration.ofMinutes(1), clock = clock) - assertTrue(tokenBucket.tryTake("42")) - currentTime = currentTime.plusSeconds(10) - assertFalse(tokenBucket.tryTake("42")) - } + fun `should return false if there are no remaining tokens by key`() = + runTest { + val tokenBucket = TokenBucketLimiter(1, Duration.ofMinutes(1), clock = clock) + assertTrue(tokenBucket.tryTake("42")) + currentTime = currentTime.plusSeconds(10) + assertFalse(tokenBucket.tryTake("42")) + } } @Nested inner class AwaitUntilTakeTest { @Test - fun `should not await if there are remaining tokens`() = runTest { - val tokenBucket = TokenBucketLimiter(1, Duration.ofMinutes(1), clock = clock) - withTimeout(100) { - tokenBucket.awaitUntilTake() + fun `should not await if there are remaining tokens`() = + runTest { + val tokenBucket = TokenBucketLimiter(1, Duration.ofMinutes(1), clock = clock) + withTimeout(100) { + tokenBucket.awaitUntilTake() + } } - } @Test - fun `should await until one token is available`() = runTest { - val tokenBucket = TokenBucketLimiter(1L, Duration.ofMinutes(1), clock = clock) - var completed = false - - launch { - tokenBucket.awaitUntilTake() - tokenBucket.awaitUntilTake() - completed = true - } - - withTimeout(100) { - advanceTimeBy(1000) - assertFalse(completed) - currentTime = currentTime.plusSeconds(60) - advanceTimeBy(1000) - - assertTrue(completed) + fun `should await until one token is available`() = + runTest { + val tokenBucket = TokenBucketLimiter(1L, Duration.ofMinutes(1), clock = clock) + var completed = false + + launch { + tokenBucket.awaitUntilTake() + tokenBucket.awaitUntilTake() + completed = true + } + + withTimeout(100) { + advanceTimeBy(1000) + assertFalse(completed) + currentTime = currentTime.plusSeconds(60) + advanceTimeBy(1000) + + assertTrue(completed) + } } - } @Test - fun `should not await if there are remaining tokens by key`() = runTest { - val tokenBucket = TokenBucketLimiter(1, Duration.ofMinutes(1), clock = clock) - withTimeout(100) { - tokenBucket.awaitUntilTake("42") - tokenBucket.awaitUntilTake("43") + fun `should not await if there are remaining tokens by key`() = + runTest { + val tokenBucket = TokenBucketLimiter(1, Duration.ofMinutes(1), clock = clock) + withTimeout(100) { + tokenBucket.awaitUntilTake("42") + tokenBucket.awaitUntilTake("43") + } } - } } } diff --git a/krate-core/src/test/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorageBuilderTest.kt b/krate-core/src/test/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorageBuilderTest.kt index 658a9ff..3e8ec60 100644 --- a/krate-core/src/test/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorageBuilderTest.kt +++ b/krate-core/src/test/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorageBuilderTest.kt @@ -38,7 +38,6 @@ import kotlin.time.Duration.Companion.hours import kotlin.time.Duration.Companion.minutes class MemoryStateStorageBuilderTest { - @Test fun `should return an instance of MemoryStateStorage`() { val memoryStateStorage = memoryStateStorage() @@ -46,32 +45,34 @@ class MemoryStateStorageBuilderTest { } @Test - fun `should return an instance of MemoryStateStorage with a custom clock, stateStorage, ttl of 2h and checkInterval of 30m`() = runTest { - val ttl = 2.hours - val checkInterval = 30.minutes - val fixedClock = Clock.fixed(Instant.parse("2023-04-07T11:00:00Z"), ZoneOffset.UTC) - val customStateMap = spyk() - val testScope = TestScope() + fun `should return an instance of MemoryStateStorage with a custom clock, stateStorage, ttl of 2h and checkInterval of 30m`() = + runTest { + val ttl = 2.hours + val checkInterval = 30.minutes + val fixedClock = Clock.fixed(Instant.parse("2023-04-07T11:00:00Z"), ZoneOffset.UTC) + val customStateMap = spyk() + val testScope = TestScope() - val stateStorage = memoryStateStorageWithEviction { - clock = fixedClock - ttlAfterLastAccess = ttl - expirationCheckInterval = checkInterval - bucketStateMap = customStateMap - coroutineScope = testScope - } + val stateStorage = + memoryStateStorageWithEviction { + clock = fixedClock + ttlAfterLastAccess = ttl + expirationCheckInterval = checkInterval + bucketStateMap = customStateMap + coroutineScope = testScope + } - assertTrue(stateStorage is MemoryStateStorage) + assertTrue(stateStorage is MemoryStateStorage) - val state42 = BucketState(10, fixedClock.instant().minusMillis(ttl.inWholeMilliseconds - 1)) - val state410 = BucketState(10, fixedClock.instant().minusMillis(ttl.inWholeMilliseconds)) - stateStorage.compareAndSet("42") { state42 } - coVerify { customStateMap.putIfAbsent("42", any()) } - stateStorage.compareAndSet("410") { state410 } + val state42 = BucketState(10, fixedClock.instant().minusMillis(ttl.inWholeMilliseconds - 1)) + val state410 = BucketState(10, fixedClock.instant().minusMillis(ttl.inWholeMilliseconds)) + stateStorage.compareAndSet("42") { state42 } + coVerify { customStateMap.putIfAbsent("42", any()) } + stateStorage.compareAndSet("410") { state410 } - testScope.advanceTimeBy(checkInterval.inWholeMilliseconds + 1) + testScope.advanceTimeBy(checkInterval.inWholeMilliseconds + 1) - assertEquals(state42, stateStorage.getBucketState("42")) - assertNull(stateStorage.getBucketState("410")) - } + assertEquals(state42, stateStorage.getBucketState("42")) + assertNull(stateStorage.getBucketState("410")) + } } diff --git a/krate-core/src/test/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorageTest.kt b/krate-core/src/test/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorageTest.kt index a17e5ed..89d649a 100644 --- a/krate-core/src/test/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorageTest.kt +++ b/krate-core/src/test/kotlin/com/neutrine/krate/storage/memory/MemoryStateStorageTest.kt @@ -33,7 +33,6 @@ import java.time.Instant import java.time.ZoneOffset.UTC internal class MemoryStateStorageTest { - private val clock = Clock.fixed(Instant.parse("2022-08-14T00:44:00Z"), UTC) private lateinit var storage: MemoryStateStorage @@ -43,71 +42,76 @@ internal class MemoryStateStorageTest { } @Test - fun `should compare and set a new state when not exists`() = runTest { - val newState = BucketState(10, clock.instant()) - storage.compareAndSet("42") { current -> - assertNull(current) - newState - } + fun `should compare and set a new state when not exists`() = + runTest { + val newState = BucketState(10, clock.instant()) + storage.compareAndSet("42") { current -> + assertNull(current) + newState + } - assertEquals(newState, storage.getBucketState("42")) - } + assertEquals(newState, storage.getBucketState("42")) + } @Test - fun `should compare and set a new state when not exists concurrently`() = runTest { - val key = "42" - val newState = BucketState(10, clock.instant()) - val concurrentState = BucketState(8, clock.instant().plusSeconds(10)) - - storage.compareAndSet(key) { current -> - runBlocking { - storage.compareAndSet(key) { concurrentState } + fun `should compare and set a new state when not exists concurrently`() = + runTest { + val key = "42" + val newState = BucketState(10, clock.instant()) + val concurrentState = BucketState(8, clock.instant().plusSeconds(10)) + + storage.compareAndSet(key) { current -> + runBlocking { + storage.compareAndSet(key) { concurrentState } + } + current?.copy(current.remainingTokens - 1) ?: newState } - current?.copy(current.remainingTokens - 1) ?: newState - } - assertEquals(7, storage.getBucketState(key)?.remainingTokens) - } + assertEquals(7, storage.getBucketState(key)?.remainingTokens) + } @Test - fun `should compare and set a new state when exists`() = runTest { - val currentState = BucketState(10, clock.instant()) - val newState = BucketState(9, clock.instant().plusSeconds(10)) - storage.compareAndSet("42") { currentState } - - storage.compareAndSet("42") { current -> - assertEquals(currentState, current) - newState - } + fun `should compare and set a new state when exists`() = + runTest { + val currentState = BucketState(10, clock.instant()) + val newState = BucketState(9, clock.instant().plusSeconds(10)) + storage.compareAndSet("42") { currentState } + + storage.compareAndSet("42") { current -> + assertEquals(currentState, current) + newState + } - assertEquals(newState, storage.getBucketState("42")) - } + assertEquals(newState, storage.getBucketState("42")) + } @Test - fun `should compare and set a new state concurrently`() = runTest { - val key = "42" - val currentState = BucketState(10, clock.instant()) - val concurrentState = BucketState(8, clock.instant().plusSeconds(10)) - storage.compareAndSet(key) { currentState } - - storage.compareAndSet(key) { current -> - runBlocking { - storage.compareAndSet(key) { concurrentState } + fun `should compare and set a new state concurrently`() = + runTest { + val key = "42" + val currentState = BucketState(10, clock.instant()) + val concurrentState = BucketState(8, clock.instant().plusSeconds(10)) + storage.compareAndSet(key) { currentState } + + storage.compareAndSet(key) { current -> + runBlocking { + storage.compareAndSet(key) { concurrentState } + } + + current!!.copy(current.remainingTokens - 1) } - current!!.copy(current.remainingTokens - 1) + assertEquals(7, storage.getBucketState(key)?.remainingTokens) } - assertEquals(7, storage.getBucketState(key)?.remainingTokens) - } - @Test - fun `should return the current state`() = runTest { - val currentState = BucketState(10, clock.instant()) - storage.compareAndSet("42") { currentState } + fun `should return the current state`() = + runTest { + val currentState = BucketState(10, clock.instant()) + storage.compareAndSet("42") { currentState } - assertEquals(currentState, storage.getBucketState("42")) - } + assertEquals(currentState, storage.getBucketState("42")) + } @Test fun `should return null when the state not exists`() { diff --git a/krate-core/src/test/kotlin/com/neutrine/krate/storage/memory/SimpleBucketStateMapWithEvictionTest.kt b/krate-core/src/test/kotlin/com/neutrine/krate/storage/memory/SimpleBucketStateMapWithEvictionTest.kt index 70e6fd1..bb675de 100644 --- a/krate-core/src/test/kotlin/com/neutrine/krate/storage/memory/SimpleBucketStateMapWithEvictionTest.kt +++ b/krate-core/src/test/kotlin/com/neutrine/krate/storage/memory/SimpleBucketStateMapWithEvictionTest.kt @@ -63,30 +63,32 @@ class SimpleBucketStateMapWithEvictionTest { } @Test - fun `should expire keys not accessed for more than ttlAfterLastAccess`() = runTest { - val state42 = AtomicReference(BucketState(10, clock.instant().minusMillis(ttlAfterLastAccess.inWholeMilliseconds - 1))) - val state410 = AtomicReference(BucketState(10, clock.instant().minusMillis(ttlAfterLastAccess.inWholeMilliseconds))) + fun `should expire keys not accessed for more than ttlAfterLastAccess`() = + runTest { + val state42 = AtomicReference(BucketState(10, clock.instant().minusMillis(ttlAfterLastAccess.inWholeMilliseconds - 1))) + val state410 = AtomicReference(BucketState(10, clock.instant().minusMillis(ttlAfterLastAccess.inWholeMilliseconds))) - stateMap.putIfAbsent("42", state42) - coVerify { baseStateMap.putIfAbsent("42", state42) } - stateMap.putIfAbsent("410", state410) + stateMap.putIfAbsent("42", state42) + coVerify { baseStateMap.putIfAbsent("42", state42) } + stateMap.putIfAbsent("410", state410) - testScope.advanceTimeBy(expirationCheckInterval.inWholeMilliseconds + 1) + testScope.advanceTimeBy(expirationCheckInterval.inWholeMilliseconds + 1) - assertEquals(state42, stateMap.getBucketStateReference("42")) - verify { baseStateMap.getBucketStateReference("42") } - assertNull(stateMap.getBucketStateReference("410")) - } + assertEquals(state42, stateMap.getBucketStateReference("42")) + verify { baseStateMap.getBucketStateReference("42") } + assertNull(stateMap.getBucketStateReference("410")) + } @Test - fun `should not expire keys before expirationCheckInterval`() = runTest { - val state410 = AtomicReference(BucketState(10, clock.instant().minusMillis(ttlAfterLastAccess.inWholeMilliseconds))) - stateMap.putIfAbsent("410", state410) + fun `should not expire keys before expirationCheckInterval`() = + runTest { + val state410 = AtomicReference(BucketState(10, clock.instant().minusMillis(ttlAfterLastAccess.inWholeMilliseconds))) + stateMap.putIfAbsent("410", state410) - testScope.advanceTimeBy(expirationCheckInterval.inWholeMilliseconds) - assertEquals(state410, stateMap.getBucketStateReference("410")) + testScope.advanceTimeBy(expirationCheckInterval.inWholeMilliseconds) + assertEquals(state410, stateMap.getBucketStateReference("410")) - testScope.advanceTimeBy(1) - assertNull(stateMap.getBucketStateReference("410")) - } + testScope.advanceTimeBy(1) + assertNull(stateMap.getBucketStateReference("410")) + } } diff --git a/krate-storage-memory-caffeine/src/main/kotlin/com/neutrine/krate/storage/memory/caffeine/CaffeineBucketStateMap.kt b/krate-storage-memory-caffeine/src/main/kotlin/com/neutrine/krate/storage/memory/caffeine/CaffeineBucketStateMap.kt index 3bbb40f..de941ba 100644 --- a/krate-storage-memory-caffeine/src/main/kotlin/com/neutrine/krate/storage/memory/caffeine/CaffeineBucketStateMap.kt +++ b/krate-storage-memory-caffeine/src/main/kotlin/com/neutrine/krate/storage/memory/caffeine/CaffeineBucketStateMap.kt @@ -34,18 +34,22 @@ import kotlin.time.toJavaDuration */ class CaffeineBucketStateMap( maximumSize: Long, - expireAfterAccess: Duration + expireAfterAccess: Duration, ) : BucketStateMap { - private val cache: Cache> = Caffeine.newBuilder() - .maximumSize(maximumSize) - .expireAfterAccess(expireAfterAccess.toJavaDuration()) - .build() + private val cache: Cache> = + Caffeine.newBuilder() + .maximumSize(maximumSize) + .expireAfterAccess(expireAfterAccess.toJavaDuration()) + .build() override fun getBucketStateReference(key: String): AtomicReference? { return cache.getIfPresent(key) } - override fun putIfAbsent(key: String, value: AtomicReference): AtomicReference? { + override fun putIfAbsent( + key: String, + value: AtomicReference, + ): AtomicReference? { return cache.asMap().putIfAbsent(key, value) } } diff --git a/krate-storage-memory-caffeine/src/test/kotlin/com/neutrine/krate/storage/memory/caffeine/MemoryCaffeineStateStorageBuilderTest.kt b/krate-storage-memory-caffeine/src/test/kotlin/com/neutrine/krate/storage/memory/caffeine/MemoryCaffeineStateStorageBuilderTest.kt index ee7c78b..2812032 100644 --- a/krate-storage-memory-caffeine/src/test/kotlin/com/neutrine/krate/storage/memory/caffeine/MemoryCaffeineStateStorageBuilderTest.kt +++ b/krate-storage-memory-caffeine/src/test/kotlin/com/neutrine/krate/storage/memory/caffeine/MemoryCaffeineStateStorageBuilderTest.kt @@ -31,17 +31,18 @@ import java.time.Instant class MemoryCaffeineStateStorageBuilderTest { @Test - fun `should return an instance of MemoryStateStorage with CaffeineBucketStateMap`() = runTest { - val stateStorage = memoryCaffeineStateStorage() - assertTrue(stateStorage is MemoryStateStorage) + fun `should return an instance of MemoryStateStorage with CaffeineBucketStateMap`() = + runTest { + val stateStorage = memoryCaffeineStateStorage() + assertTrue(stateStorage is MemoryStateStorage) - val state = BucketState(10, Instant.now()) - stateStorage.compareAndSet("A") { state } - stateStorage.compareAndSet("B") { state } - stateStorage.compareAndSet("C") { state } + val state = BucketState(10, Instant.now()) + stateStorage.compareAndSet("A") { state } + stateStorage.compareAndSet("B") { state } + stateStorage.compareAndSet("C") { state } - assertEquals(state, stateStorage.getBucketState("A")) - assertEquals(state, stateStorage.getBucketState("B")) - assertEquals(state, stateStorage.getBucketState("C")) - } + assertEquals(state, stateStorage.getBucketState("A")) + assertEquals(state, stateStorage.getBucketState("B")) + assertEquals(state, stateStorage.getBucketState("C")) + } } diff --git a/krate-storage-redis/src/main/kotlin/com/neutrine/krate/storage/redis/RedisStateStorage.kt b/krate-storage-redis/src/main/kotlin/com/neutrine/krate/storage/redis/RedisStateStorage.kt index 78fd9a7..67a0df2 100644 --- a/krate-storage-redis/src/main/kotlin/com/neutrine/krate/storage/redis/RedisStateStorage.kt +++ b/krate-storage-redis/src/main/kotlin/com/neutrine/krate/storage/redis/RedisStateStorage.kt @@ -35,7 +35,7 @@ import java.time.Instant */ class RedisStateStorage( val host: String, - val port: Int + val port: Int, ) : StateStorage { private val redisPool: JedisPool = JedisPool(host, port) @@ -45,7 +45,10 @@ class RedisStateStorage( } } - override suspend fun compareAndSet(key: String, compareAndSetFunction: (current: BucketState?) -> BucketState): Unit = + override suspend fun compareAndSet( + key: String, + compareAndSetFunction: (current: BucketState?) -> BucketState, + ): Unit = redisPool.resource.use { jedis -> val computedKey = computeKey(key) jedis.watch(computedKey) @@ -63,16 +66,18 @@ class RedisStateStorage( private fun computeKey(key: String) = "krate:instance:$key" - private fun BucketState.toMap(): Map = mapOf( - KEY_REMAINING_TOKENS to remainingTokens.toString(), - KEY_LAST_UPDATED to lastUpdated.toEpochMilli().toString() - ) + private fun BucketState.toMap(): Map = + mapOf( + KEY_REMAINING_TOKENS to remainingTokens.toString(), + KEY_LAST_UPDATED to lastUpdated.toEpochMilli().toString(), + ) private fun Map.toBucketState(): BucketState? { val remainingTokens: Long? = this["remainingTokens"]?.toLongOrNull() - val lastUpdated: Instant? = this["lastUpdated"]?.toLongOrNull()?.let { - Instant.ofEpochMilli(it) - } + val lastUpdated: Instant? = + this["lastUpdated"]?.toLongOrNull()?.let { + Instant.ofEpochMilli(it) + } if (remainingTokens == null || lastUpdated == null) { return null diff --git a/krate-storage-redis/src/test/kotlin/com/neutrine/krate/storage/redis/RedisStateStorageBuilderTest.kt b/krate-storage-redis/src/test/kotlin/com/neutrine/krate/storage/redis/RedisStateStorageBuilderTest.kt index 0c62283..0b0a3f4 100644 --- a/krate-storage-redis/src/test/kotlin/com/neutrine/krate/storage/redis/RedisStateStorageBuilderTest.kt +++ b/krate-storage-redis/src/test/kotlin/com/neutrine/krate/storage/redis/RedisStateStorageBuilderTest.kt @@ -4,7 +4,6 @@ import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.Test internal class RedisStateStorageBuilderTest { - @Test fun `should return an instance of RedisStateStorageBuilder with default values`() { val stateStorage: RedisStateStorage = redisStateStorage {} @@ -15,10 +14,11 @@ internal class RedisStateStorageBuilderTest { @Test fun `should return an instance of RedisStateStorageBuilder with custom values`() { - val stateStorage: RedisStateStorage = redisStateStorage { - host = "127.0.0.1" - port = 4242 - } + val stateStorage: RedisStateStorage = + redisStateStorage { + host = "127.0.0.1" + port = 4242 + } assertEquals("127.0.0.1", stateStorage.host) assertEquals(4242, stateStorage.port) diff --git a/krate-storage-redis/src/test/kotlin/com/neutrine/krate/storage/redis/RedisStateStorageTest.kt b/krate-storage-redis/src/test/kotlin/com/neutrine/krate/storage/redis/RedisStateStorageTest.kt index b7c2016..51fcc77 100644 --- a/krate-storage-redis/src/test/kotlin/com/neutrine/krate/storage/redis/RedisStateStorageTest.kt +++ b/krate-storage-redis/src/test/kotlin/com/neutrine/krate/storage/redis/RedisStateStorageTest.kt @@ -43,8 +43,9 @@ internal class RedisStateStorageTest { private lateinit var storage: RedisStateStorage @Container - val redis: GenericContainer<*> = GenericContainer(DockerImageName.parse("redis:5.0.3-alpine")) - .withExposedPorts(6379) + val redis: GenericContainer<*> = + GenericContainer(DockerImageName.parse("redis:5.0.3-alpine")) + .withExposedPorts(6379) @BeforeEach fun setup() { @@ -52,73 +53,78 @@ internal class RedisStateStorageTest { } @Test - fun `should compare and set a new state when not exists`() = runTest { - val newState = BucketState(10, clock.instant()) - storage.compareAndSet("42") { current -> - assertNull(current) - newState + fun `should compare and set a new state when not exists`() = + runTest { + val newState = BucketState(10, clock.instant()) + storage.compareAndSet("42") { current -> + assertNull(current) + newState + } + + assertEquals(newState, storage.getBucketState("42")) } - assertEquals(newState, storage.getBucketState("42")) - } - @Test - fun `should compare and set a new state when not exists concurrently`() = runTest { - val key = "42" - val newState = BucketState(10, clock.instant()) - val concurrentState = BucketState(8, clock.instant().plusSeconds(10)) - - launch { - storage.compareAndSet(key) { concurrentState } - } - storage.compareAndSet(key) { current -> - advanceTimeBy(100) - current?.copy(current.remainingTokens - 1) ?: newState + fun `should compare and set a new state when not exists concurrently`() = + runTest { + val key = "42" + val newState = BucketState(10, clock.instant()) + val concurrentState = BucketState(8, clock.instant().plusSeconds(10)) + + launch { + storage.compareAndSet(key) { concurrentState } + } + storage.compareAndSet(key) { current -> + advanceTimeBy(100) + current?.copy(current.remainingTokens - 1) ?: newState + } + + assertEquals(7, storage.getBucketState(key)?.remainingTokens) } - assertEquals(7, storage.getBucketState(key)?.remainingTokens) - } - @Test - fun `should compare and set a new state when exists`() = runTest { - val currentState = BucketState(10, clock.instant()) - val newState = BucketState(9, clock.instant().plusSeconds(10)) - storage.compareAndSet("42") { currentState } - - storage.compareAndSet("42") { current -> - assertEquals(currentState, current) - newState + fun `should compare and set a new state when exists`() = + runTest { + val currentState = BucketState(10, clock.instant()) + val newState = BucketState(9, clock.instant().plusSeconds(10)) + storage.compareAndSet("42") { currentState } + + storage.compareAndSet("42") { current -> + assertEquals(currentState, current) + newState + } + + assertEquals(newState, storage.getBucketState("42")) } - assertEquals(newState, storage.getBucketState("42")) - } - @Test - fun `should compare and set a new state concurrently`() = runTest { - val key = "42" - val currentState = BucketState(10, clock.instant()) - val concurrentState = BucketState(8, clock.instant().plusSeconds(10)) - storage.compareAndSet(key) { currentState } - - launch { - storage.compareAndSet(key) { concurrentState } - } - - storage.compareAndSet(key) { current -> - advanceTimeBy(100) - current!!.copy(current.remainingTokens - 1) + fun `should compare and set a new state concurrently`() = + runTest { + val key = "42" + val currentState = BucketState(10, clock.instant()) + val concurrentState = BucketState(8, clock.instant().plusSeconds(10)) + storage.compareAndSet(key) { currentState } + + launch { + storage.compareAndSet(key) { concurrentState } + } + + storage.compareAndSet(key) { current -> + advanceTimeBy(100) + current!!.copy(current.remainingTokens - 1) + } + + assertEquals(7, storage.getBucketState(key)?.remainingTokens) } - assertEquals(7, storage.getBucketState(key)?.remainingTokens) - } - @Test - fun `should return the current state`() = runTest { - val currentState = BucketState(10, clock.instant()) - storage.compareAndSet("42") { currentState } + fun `should return the current state`() = + runTest { + val currentState = BucketState(10, clock.instant()) + storage.compareAndSet("42") { currentState } - assertEquals(currentState, storage.getBucketState("42")) - } + assertEquals(currentState, storage.getBucketState("42")) + } @Test fun `should return null when the state not exists`() {