@@ -5,44 +5,15 @@ import kotlinx.coroutines.sync.Semaphore
5
5
import kotlinx.coroutines.sync.withPermit
6
6
import kotlin.time.Duration
7
7
8
- class Ratelimit ( val burst : Int , val duration : Duration ) {
8
+ class SuspendingRatelimit ( private val burst : Int , private val duration : Duration ) {
9
9
@Volatile
10
10
private var remainingQuota: Int = burst
11
- val remaining: Int get() = remainingQuota
12
11
13
12
@Volatile
14
13
private var resetTimestamp: Long = 0
15
- val resetAt: Long get() = resetTimestamp
16
14
17
15
private val quotaRequestSem = Semaphore (1 )
18
16
19
- suspend fun requestQuota () {
20
- quotaRequestSem.withPermit {
21
- if (remainingQuota <= 0 ) {
22
- tryResetQuota()
23
- val waitTime = calculateWaitTime()
24
- delay(waitTime)
25
- }
26
- tryResetQuota()
27
- check(remainingQuota > 0 )
28
- remainingQuota--
29
- }
30
- }
31
-
32
- fun tryRequestQuota (): Pair <Boolean , Long > {
33
- tryResetQuota()
34
- if (remainingQuota <= 0 ) {
35
- return false to calculateWaitTime()
36
- }
37
- check(remainingQuota > 0 )
38
- remainingQuota--
39
- return true to calculateWaitTime()
40
- }
41
-
42
- fun addQuota (amount : Int ) {
43
- remainingQuota + = amount
44
- }
45
-
46
17
fun overrideRatelimit (
47
18
remainingQuota : Int ,
48
19
resetTimestamp : Long ,
@@ -61,4 +32,29 @@ class Ratelimit(val burst: Int, val duration: Duration) {
61
32
resetTimestamp = System .currentTimeMillis() + duration.inWholeMilliseconds
62
33
}
63
34
}
35
+
36
+ suspend fun requestQuota () {
37
+ quotaRequestSem.withPermit {
38
+ if (remainingQuota <= 0 ) {
39
+ val waitTime = calculateWaitTime()
40
+ delay(waitTime)
41
+ }
42
+ tryResetQuota()
43
+
44
+ check(remainingQuota > 0 )
45
+ remainingQuota--
46
+ }
47
+ }
48
+
49
+ fun tryRequestQuota (): Pair <Boolean , Long ?> {
50
+ if (remainingQuota <= 0 ) {
51
+ val waitTime = calculateWaitTime()
52
+ return false to waitTime
53
+ }
54
+ tryResetQuota()
55
+
56
+ check(remainingQuota > 0 )
57
+ remainingQuota--
58
+ return true to null
59
+ }
64
60
}
0 commit comments