Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- @ExperimentalTime
- @ExperimentalCoroutinesApi
- private class RateLimiterImpl(override val permitsPerSecond: Int, override val timeSource: TimeSource) : RateLimiter {
- init {
- require(permitsPerSecond > 0) { "Rate must be > 0" }
- }
- private val mutex = Mutex()
- private val permitsLeft = atomic(permitsPerSecond)
- private val timeFirstPermitAcquired = atomic<TimeMark?>(null)
- override suspend fun acquire() {
- mutex.withLock {
- if (permitsLeft.getAndUpdate { permits -> if (permits == 0) permitsPerSecond - 1 else permits - 1 } < 1) {
- delay(duration = timeUntilReset())
- resetAndAcquire()
- } else {
- timeFirstPermitAcquired.compareAndSet(null, timeSource.markNow())
- }
- }
- }
- override fun tryAcquire(): Boolean {
- val isAcquired = permitsLeft.getAndDecrement() >= 1
- if (!isAcquired && isTimeToReset()) {
- return runBlocking {
- mutex.withLock {
- if (isTimeToReset()) {
- resetAndAcquire(); true
- } else {
- tryAcquire()
- }
- }
- }
- }
- if (isAcquired) {
- timeFirstPermitAcquired.compareAndSet(null, timeSource.markNow())
- }
- return isAcquired
- }
- private fun resetAndAcquire() {
- permitsLeft.value = permitsPerSecond - 1
- timeFirstPermitAcquired.value = timeSource.markNow()
- }
- private fun isTimeToReset() = timeUntilReset().inWholeMilliseconds <= 0
- private fun timeUntilReset() = Duration.seconds(1).minus(timeFirstPermitAcquired.value!!.elapsedNow())
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement