Skip to content

Commit 958317b

Browse files
rubengeesrubengees
committed
Improve rate limiting behaviour in regards to thread safety and the cache
1 parent 5aaf305 commit 958317b

3 files changed

Lines changed: 130 additions & 34 deletions

File tree

‎library/src/main/kotlin/me/proxer/library/ProxerApi.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -260,7 +260,7 @@ class ProxerApi private constructor(retrofit: Retrofit) {
260260
)
261261
)
262262

263-
if (enableRateLimitProtection) add(0, RateLimitInterceptor(moshi))
263+
if (enableRateLimitProtection) add(4, RateLimitInterceptor(moshi, client?.cache))
264264
}
265265

266266
certificatePinner(constructCertificatePinner())

‎library/src/main/kotlin/me/proxer/library/internal/interceptor/RateLimitInterceptor.kt‎

Lines changed: 40 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -5,18 +5,16 @@ import com.squareup.moshi.Types
55
import me.proxer.library.ProxerException.ServerErrorType.RATE_LIMIT
66
import me.proxer.library.internal.ProxerResponse
77
import me.proxer.library.util.ProxerUrls
8+
import okhttp3.Cache
89
import okhttp3.Interceptor
910
import okhttp3.MediaType.Companion.toMediaTypeOrNull
1011
import okhttp3.Protocol
1112
import okhttp3.Request
1213
import okhttp3.Response
1314
import okhttp3.ResponseBody.Companion.toResponseBody
1415
import java.util.Date
15-
import java.util.concurrent.locks.ReentrantReadWriteLock
16-
import kotlin.concurrent.read
17-
import kotlin.concurrent.write
1816

19-
internal class RateLimitInterceptor(private val moshi: Moshi) : Interceptor {
17+
internal class RateLimitInterceptor(private val moshi: Moshi, private val cache: Cache?) : Interceptor {
2018

2119
private companion object {
2220
private const val WAIT_THRESHOLD = 2_000
@@ -47,8 +45,6 @@ internal class RateLimitInterceptor(private val moshi: Moshi) : Interceptor {
4745

4846
private var state = limits.mapValues { State(0, null) }
4947

50-
private val lock = ReentrantReadWriteLock()
51-
5248
@Suppress("ReturnCount")
5349
override fun intercept(chain: Interceptor.Chain): Response {
5450
val oldRequest = chain.request()
@@ -74,35 +70,47 @@ internal class RateLimitInterceptor(private val moshi: Moshi) : Interceptor {
7470
chain: Interceptor.Chain,
7571
oldRequest: Request
7672
): Response {
77-
val (requests, firstRequest) = lock.read { state.getValue(apiClass) }
78-
val timeDifferenceMillis = if (firstRequest != null) Date().time - firstRequest.time else null
79-
80-
if (timeDifferenceMillis == null || timeDifferenceMillis > limit.millis) {
81-
lock.write { state = state.update(apiClass, 1, Date()) }
73+
synchronized(limit) {
74+
val (requests, firstRequest) = state.getValue(apiClass)
75+
val timeDifferenceMillis = if (firstRequest != null) Date().time - firstRequest.time else null
76+
77+
if (timeDifferenceMillis == null || timeDifferenceMillis > limit.millis) {
78+
state = state.update(apiClass, 1, Date())
79+
80+
return chain.proceed(oldRequest)
81+
} else if (requests + 2 >= limit.maxRequests) {
82+
// Do nothing if the request is in the cache.
83+
if (cache?.urls()?.asSequence()?.contains(oldRequest.url.toString()) == true) {
84+
return chain.proceed(oldRequest)
85+
}
86+
87+
return if ((limit.millis - timeDifferenceMillis) < WAIT_THRESHOLD) {
88+
Thread.sleep(timeDifferenceMillis)
89+
90+
intercept(chain)
91+
} else {
92+
val errorBody = moshi
93+
.adapter<ProxerResponse<Unit?>>(responseType)
94+
.toJson(ProxerResponse(true, "Rate limit", RATE_LIMIT.code, null))
95+
96+
Response.Builder()
97+
.body(errorBody.toResponseBody(errorMediaType))
98+
.protocol(Protocol.HTTP_1_1)
99+
.request(oldRequest)
100+
.code(200)
101+
.message("")
102+
.build()
103+
}
104+
} else {
105+
val response = chain.proceed(oldRequest)
82106

83-
return chain.proceed(oldRequest)
84-
} else if (requests + 2 >= limit.maxRequests) {
85-
return if ((limit.millis - timeDifferenceMillis) < WAIT_THRESHOLD) {
86-
Thread.sleep(timeDifferenceMillis)
107+
// Only increment on responses not coming from cache.
108+
if (response.networkResponse != null) {
109+
state = state.update(apiClass, requests + 1, firstRequest)
110+
}
87111

88-
intercept(chain)
89-
} else {
90-
val errorBody = moshi
91-
.adapter<ProxerResponse<Unit?>>(responseType)
92-
.toJson(ProxerResponse(true, "Rate limit", RATE_LIMIT.code, null))
93-
94-
Response.Builder()
95-
.body(errorBody.toResponseBody(errorMediaType))
96-
.protocol(Protocol.HTTP_1_1)
97-
.request(oldRequest)
98-
.code(200)
99-
.message("")
100-
.build()
112+
return response
101113
}
102-
} else {
103-
lock.write { state = state.update(apiClass, requests + 1, firstRequest) }
104-
105-
return chain.proceed(oldRequest)
106114
}
107115
}
108116

‎library/src/test/kotlin/me/proxer/library/internal/interceptor/RateLimitInterceptorTest.kt‎

Lines changed: 89 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,18 +5,27 @@ import me.proxer.library.ProxerException
55
import me.proxer.library.ProxerException.ServerErrorType
66
import me.proxer.library.ProxerTest
77
import me.proxer.library.enums.Language
8+
import me.proxer.library.fromResource
89
import me.proxer.library.runRequest
10+
import okhttp3.Cache
11+
import okhttp3.mockwebserver.MockResponse
12+
import org.amshove.kluent.AnyException
913
import org.amshove.kluent.invoking
1014
import org.amshove.kluent.shouldBe
15+
import org.amshove.kluent.shouldNotBeNull
16+
import org.amshove.kluent.shouldNotThrow
1117
import org.amshove.kluent.shouldThrow
1218
import org.junit.jupiter.api.Disabled
19+
import org.junit.jupiter.api.RepeatedTest
1320
import org.junit.jupiter.api.Test
21+
import java.util.concurrent.CountDownLatch
22+
import java.util.concurrent.TimeUnit
1423

1524
class RateLimitInterceptorTest : ProxerTest() {
1625

1726
private val rateLimitApi = ProxerApi.Builder("mock-key")
1827
.enableRateLimitProtection()
19-
.client(client)
28+
.client(client.newBuilder().cache(Cache(createTempDir(), 1_024L * 1_024L)).build())
2029
.build()
2130

2231
@Test
@@ -45,6 +54,85 @@ class RateLimitInterceptorTest : ProxerTest() {
4554
result.exception.serverErrorType shouldBe ServerErrorType.RATE_LIMIT
4655
}
4756

57+
@Test
58+
fun testCachedResponse() {
59+
val response = MockResponse()
60+
.setBody(fromResource("chapter.json"))
61+
.setHeader("Cache-Control", "public, max-age=31536000")
62+
63+
// Limit 10.
64+
repeat(8) {
65+
server.enqueue(response)
66+
rateLimitApi.manga.chapter("123", 123, Language.ENGLISH).build().execute()
67+
}
68+
69+
invoking {
70+
server.enqueue(response)
71+
rateLimitApi.manga.chapter("123", 123, Language.ENGLISH).build().execute()
72+
} shouldNotThrow AnyException
73+
}
74+
75+
@Test
76+
@RepeatedTest(50)
77+
fun testParallelRequests() {
78+
// Limit 10.
79+
repeat(7) {
80+
server.runRequest("chapter.json") {
81+
rateLimitApi.manga.chapter("123", 123, Language.ENGLISH).build().execute()
82+
}
83+
}
84+
85+
server.enqueue(MockResponse().setBody(fromResource("chapter.json")))
86+
server.enqueue(MockResponse().setBody(fromResource("chapter.json")))
87+
88+
val lock = CountDownLatch(2)
89+
var error: ProxerException? = null
90+
91+
rateLimitApi.manga.chapter("123", 123, Language.ENGLISH).build().enqueue(
92+
{ lock.countDown() },
93+
{
94+
error = it
95+
lock.countDown()
96+
}
97+
)
98+
99+
rateLimitApi.manga.chapter("123", 123, Language.ENGLISH).build().enqueue(
100+
{ lock.countDown() },
101+
{
102+
error = it
103+
lock.countDown()
104+
}
105+
)
106+
107+
lock.await(3_000L, TimeUnit.MILLISECONDS)
108+
109+
error.shouldNotBeNull()
110+
error!!.serverErrorType shouldBe ServerErrorType.RATE_LIMIT
111+
}
112+
113+
@Test
114+
@Disabled // Does not work because of real urls being passed in and mock urls in cache.
115+
fun testLimitButCached() {
116+
// Limit 10.
117+
repeat(7) {
118+
server.runRequest("chapter.json") {
119+
rateLimitApi.manga.chapter("123", 123, Language.ENGLISH).build().execute()
120+
}
121+
}
122+
123+
val response = MockResponse()
124+
.setBody(fromResource("chapter.json"))
125+
.setHeader("Cache-Control", "public, max-age=31536000")
126+
127+
server.enqueue(response)
128+
rateLimitApi.manga.chapter("123", 123, Language.ENGLISH).build().execute()
129+
130+
invoking {
131+
server.enqueue(response)
132+
rateLimitApi.manga.chapter("123", 123, Language.ENGLISH).build().execute()
133+
} shouldNotThrow AnyException
134+
}
135+
48136
@Test
49137
@Disabled
50138
fun testResetStateAfterTimeLimit() {

0 commit comments

Comments
 (0)