Skip to content

Commit 7f56f31

Browse files
mohamedlajmileanixmohamedlajmileanix
committed
CID-2911: Expand rate limit handling to GitHub GraphQL calls
1 parent 3c57f0a commit 7f56f31

File tree

9 files changed

+118
-27
lines changed

9 files changed

+118
-27
lines changed

‎src/main/kotlin/net/leanix/githubagent/dto/RateLimitType.kt‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package net.leanix.githubagent.dto
2+
3+
enum class RateLimitType {
4+
GRAPHQL,
5+
REST,
6+
SEARCH
7+
}

‎src/main/kotlin/net/leanix/githubagent/handler/RateLimitHandler.kt‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package net.leanix.githubagent.handler
22

3+
import net.leanix.githubagent.dto.RateLimitType
34
import net.leanix.githubagent.services.SyncLogService
45
import net.leanix.githubagent.shared.RateLimitMonitor
56
import org.springframework.stereotype.Component
@@ -9,11 +10,14 @@ class RateLimitHandler(
910
private val syncLogService: SyncLogService,
1011
) {
1112

12-
fun <T> executeWithRateLimitHandler(block: () -> T): T {
13+
fun <T> executeWithRateLimitHandler(rateLimitType: RateLimitType, block: () -> T): T {
1314
while (true) {
14-
val waitTimeSeconds = RateLimitMonitor.shouldThrottle()
15+
val waitTimeSeconds = RateLimitMonitor.shouldThrottle(rateLimitType)
1516
if (waitTimeSeconds > 0) {
16-
syncLogService.sendInfoLog("Approaching rate limit. Waiting for $waitTimeSeconds seconds.")
17+
syncLogService.sendInfoLog(
18+
"Approaching rate limit for $rateLimitType calls. " +
19+
"Waiting for $waitTimeSeconds seconds."
20+
)
1721
Thread.sleep(waitTimeSeconds * 1000)
1822
}
1923
return block()

‎src/main/kotlin/net/leanix/githubagent/interceptor/RateLimitInterceptor.kt‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
package net.leanix.githubagent.interceptor
2+
3+
import net.leanix.githubagent.shared.RateLimitMonitor
4+
import net.leanix.githubagent.shared.determineRateLimitType
5+
import org.springframework.web.reactive.function.client.ClientRequest
6+
import org.springframework.web.reactive.function.client.ClientResponse
7+
import org.springframework.web.reactive.function.client.ExchangeFilterFunction
8+
import org.springframework.web.reactive.function.client.ExchangeFunction
9+
import reactor.core.publisher.Mono
10+
11+
class RateLimitInterceptor : ExchangeFilterFunction {
12+
13+
override fun filter(request: ClientRequest, next: ExchangeFunction): Mono<ClientResponse> {
14+
return next.exchange(request).flatMap { response ->
15+
val headers = response.headers().asHttpHeaders()
16+
val rateLimitRemaining = headers["X-RateLimit-Remaining"]?.firstOrNull()?.toIntOrNull()
17+
val rateLimitReset = headers["X-RateLimit-Reset"]?.firstOrNull()?.toLongOrNull()
18+
19+
if (rateLimitRemaining != null && rateLimitReset != null) {
20+
val rateLimitType = determineRateLimitType(request.url().toString())
21+
RateLimitMonitor.updateRateLimitInfo(rateLimitType, rateLimitRemaining, rateLimitReset)
22+
}
23+
Mono.just(response)
24+
}
25+
}
26+
}

‎src/main/kotlin/net/leanix/githubagent/interceptor/RateLimitResponseInterceptor.kt‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package net.leanix.githubagent.interceptor
33
import feign.InvocationContext
44
import feign.ResponseInterceptor
55
import net.leanix.githubagent.shared.RateLimitMonitor
6+
import net.leanix.githubagent.shared.determineRateLimitType
67

78
class RateLimitResponseInterceptor : ResponseInterceptor {
89

@@ -19,7 +20,8 @@ class RateLimitResponseInterceptor : ResponseInterceptor {
1920
val rateLimitReset = headers["x-ratelimit-reset"]?.firstOrNull()?.toLongOrNull()
2021

2122
if (rateLimitRemaining != null && rateLimitReset != null) {
22-
RateLimitMonitor.updateRateLimitInfo(rateLimitRemaining, rateLimitReset)
23+
val rateLimitType = determineRateLimitType(response.request().url())
24+
RateLimitMonitor.updateRateLimitInfo(rateLimitType, rateLimitRemaining, rateLimitReset)
2325
}
2426
}
2527

‎src/main/kotlin/net/leanix/githubagent/services/GitHubGraphQLService.kt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import net.leanix.githubagent.dto.RepositoryDto
77
import net.leanix.githubagent.exceptions.GraphQLApiException
88
import net.leanix.githubagent.graphql.data.GetRepositories
99
import net.leanix.githubagent.graphql.data.GetRepositoryManifestContent
10+
import net.leanix.githubagent.interceptor.RateLimitInterceptor
1011
import org.slf4j.LoggerFactory
1112
import org.springframework.stereotype.Component
1213
import org.springframework.web.reactive.function.client.WebClient
@@ -97,6 +98,7 @@ class GitHubGraphQLService(
9798
GraphQLWebClient(
9899
url = "${cachingService.get("baseUrl")}/api/graphql",
99100
builder = WebClient.builder().defaultHeaders { it.setBearerAuth(token) }
101+
.filter(RateLimitInterceptor())
100102
)
101103
}
102104

‎src/main/kotlin/net/leanix/githubagent/services/GitHubScanningService.kt‎

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import net.leanix.githubagent.dto.ManifestFileDTO
77
import net.leanix.githubagent.dto.ManifestFilesDTO
88
import net.leanix.githubagent.dto.Organization
99
import net.leanix.githubagent.dto.OrganizationDto
10+
import net.leanix.githubagent.dto.RateLimitType
1011
import net.leanix.githubagent.dto.RepositoryDto
1112
import net.leanix.githubagent.exceptions.JwtTokenNotFound
1213
import net.leanix.githubagent.exceptions.ManifestFileNotFoundException
@@ -59,14 +60,16 @@ class GitHubScanningService(
5960
return
6061
}
6162
val installationToken = cachingService.get("installationToken:${installations.first().id}")
62-
val organizations = gitHubClient.getOrganizations("Bearer $installationToken")
63-
.map { organization ->
64-
if (installations.find { it.account.login == organization.login } != null) {
65-
OrganizationDto(organization.id, organization.login, true)
66-
} else {
67-
OrganizationDto(organization.id, organization.login, false)
63+
val organizations = rateLimitHandler.executeWithRateLimitHandler(RateLimitType.REST) {
64+
gitHubClient.getOrganizations("Bearer $installationToken")
65+
.map { organization ->
66+
if (installations.find { it.account.login == organization.login } != null) {
67+
OrganizationDto(organization.id, organization.login, true)
68+
} else {
69+
OrganizationDto(organization.id, organization.login, false)
70+
}
6871
}
69-
}
72+
}
7073
logger.info("Sending organizations data")
7174
syncLogService.sendInfoLog(
7275
"The connector found ${organizations.filter { it.installed }.size} " +
@@ -82,10 +85,12 @@ class GitHubScanningService(
8285
var page = 1
8386
val repositories = mutableListOf<RepositoryDto>()
8487
do {
85-
val repositoriesPage = gitHubGraphQLService.getRepositories(
86-
token = installationToken,
87-
cursor = cursor
88-
)
88+
val repositoriesPage = rateLimitHandler.executeWithRateLimitHandler(RateLimitType.GRAPHQL) {
89+
gitHubGraphQLService.getRepositories(
90+
token = installationToken,
91+
cursor = cursor
92+
)
93+
}
8994
webSocketService.sendMessage(
9095
"${cachingService.get("runId")}/repositories",
9196
repositoriesPage.repositories.filter { !it.archived }
@@ -120,7 +125,7 @@ class GitHubScanningService(
120125

121126
private fun fetchManifestFiles(installation: Installation, repositoryName: String) = runCatching {
122127
val installationToken = cachingService.get("installationToken:${installation.id}").toString()
123-
rateLimitHandler.executeWithRateLimitHandler {
128+
rateLimitHandler.executeWithRateLimitHandler(RateLimitType.SEARCH) {
124129
gitHubClient.searchManifestFiles(
125130
"Bearer $installationToken",
126131
"" +

‎src/main/kotlin/net/leanix/githubagent/shared/RateLimitHelper.kt‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
package net.leanix.githubagent.shared
2+
3+
import net.leanix.githubagent.dto.RateLimitType
4+
5+
fun determineRateLimitType(requestUrl: String): RateLimitType {
6+
return if (requestUrl.contains("/graphql")) {
7+
RateLimitType.GRAPHQL
8+
} else if (requestUrl.contains("/search")) {
9+
RateLimitType.SEARCH
10+
} else {
11+
RateLimitType.REST
12+
}
13+
}

‎src/main/kotlin/net/leanix/githubagent/shared/RateLimitMonitor.kt‎

Lines changed: 41 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,73 @@
11
package net.leanix.githubagent.shared
22

3+
import net.leanix.githubagent.dto.RateLimitType
34
import org.slf4j.LoggerFactory
45

56
object RateLimitMonitor {
67

78
private val logger = LoggerFactory.getLogger(RateLimitMonitor::class.java)
89

910
@Volatile
10-
private var rateLimitRemaining: Int = Int.MAX_VALUE
11+
private var graphqlRateLimitRemaining: Int = Int.MAX_VALUE
1112

1213
@Volatile
13-
private var rateLimitResetTime: Long = 0
14+
private var graphqlRateLimitResetTime: Long = 0
15+
16+
@Volatile
17+
private var restRateLimitRemaining: Int = Int.MAX_VALUE
18+
19+
@Volatile
20+
private var restRateLimitResetTime: Long = 0
21+
22+
@Volatile
23+
private var searchRateLimitRemaining: Int = Int.MAX_VALUE
24+
25+
@Volatile
26+
private var searchRateLimitResetTime: Long = 0
1427

1528
private val lock = Any()
1629

17-
fun updateRateLimitInfo(remaining: Int, resetTimeEpochSeconds: Long) {
30+
fun updateRateLimitInfo(rateLimitType: RateLimitType, remaining: Int, resetTimeEpochSeconds: Long) {
1831
synchronized(lock) {
19-
rateLimitRemaining = remaining
20-
rateLimitResetTime = resetTimeEpochSeconds
32+
when (rateLimitType) {
33+
RateLimitType.GRAPHQL -> {
34+
graphqlRateLimitRemaining = remaining
35+
graphqlRateLimitResetTime = resetTimeEpochSeconds
36+
}
37+
RateLimitType.REST -> {
38+
restRateLimitRemaining = remaining
39+
restRateLimitResetTime = resetTimeEpochSeconds
40+
}
41+
RateLimitType.SEARCH -> {
42+
searchRateLimitRemaining = remaining
43+
searchRateLimitResetTime = resetTimeEpochSeconds
44+
}
45+
}
2146
}
2247
}
2348

24-
fun shouldThrottle(): Long {
49+
fun shouldThrottle(rateLimitType: RateLimitType): Long {
2550
synchronized(lock) {
51+
val (rateLimitRemaining, rateLimitResetTime) = when (rateLimitType) {
52+
RateLimitType.GRAPHQL -> graphqlRateLimitRemaining to graphqlRateLimitResetTime
53+
RateLimitType.REST -> restRateLimitRemaining to restRateLimitResetTime
54+
RateLimitType.SEARCH -> searchRateLimitRemaining to searchRateLimitResetTime
55+
}
56+
2657
if (rateLimitRemaining <= THRESHOLD) {
2758
val currentTimeSeconds = System.currentTimeMillis() / 1000
2859
val waitTimeSeconds = rateLimitResetTime - currentTimeSeconds + 5
2960

3061
val adjustedWaitTime = if (waitTimeSeconds > 0) waitTimeSeconds else 0
3162
logger.warn(
32-
"Rate limit remaining ($rateLimitRemaining) is at or below threshold ($THRESHOLD)." +
33-
" Throttling for $adjustedWaitTime seconds."
63+
"Rate limit remaining ($rateLimitRemaining) for $rateLimitType calls, is at or below " +
64+
"threshold ($THRESHOLD). Throttling for $adjustedWaitTime seconds."
3465
)
3566
return adjustedWaitTime
3667
} else {
3768
logger.debug(
38-
"Rate limit remaining ($rateLimitRemaining) is above threshold ($THRESHOLD). No need to throttle."
69+
"Rate limit remaining ($rateLimitRemaining) for $rateLimitType calls, is above threshold" +
70+
" ($THRESHOLD). No need to throttle."
3971
)
4072
}
4173
return 0

‎src/test/kotlin/net/leanix/githubagent/services/GitHubScanningServiceTest.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,8 @@ class GitHubScanningServiceTest {
6464
every { gitHubAuthenticationService.generateAndCacheInstallationTokens(any(), any()) } returns Unit
6565
every { syncLogService.sendErrorLog(any()) } returns Unit
6666
every { syncLogService.sendInfoLog(any()) } returns Unit
67-
every { rateLimitHandler.executeWithRateLimitHandler(any<() -> Any>()) } answers
68-
{ firstArg<() -> Any>().invoke() }
67+
every { rateLimitHandler.executeWithRateLimitHandler(any(), any<() -> Any>()) } answers
68+
{ secondArg<() -> Any>().invoke() }
6969
}
7070

7171
@Test

0 commit comments

Comments
 (0)