Skip to content

Commit f7c60e6

Browse files
mohamedlajmileanixmohamedlajmileanix
authored
Merge pull request #70 from leanix/feature/CID-2911/Expand-rate-limit-handling-to-GitHub-GraphQL-calls
CID-2911: Expand rate limit handling to GitHub GraphQL calls
2 parents 979d012 + 5ca9182 commit f7c60e6

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.GitHubAppInsufficientPermissionsException
1213
import net.leanix.githubagent.exceptions.JwtTokenNotFound
@@ -82,14 +83,16 @@ class GitHubScanningService(
8283
return
8384
}
8485
val installationToken = cachingService.get("installationToken:${installations.first().id}")
85-
val organizations = gitHubAPIService.getPaginatedOrganizations(installationToken.toString())
86-
.map { organization ->
87-
if (installations.find { it.account.login == organization.login } != null) {
88-
OrganizationDto(organization.id, organization.login, true)
89-
} else {
90-
OrganizationDto(organization.id, organization.login, false)
86+
val organizations = rateLimitHandler.executeWithRateLimitHandler(RateLimitType.REST) {
87+
gitHubAPIService.getPaginatedOrganizations(installationToken.toString())
88+
.map { organization ->
89+
if (installations.find { it.account.login == organization.login } != null) {
90+
OrganizationDto(organization.id, organization.login, true)
91+
} else {
92+
OrganizationDto(organization.id, organization.login, false)
93+
}
9194
}
92-
}
95+
}
9396
logger.info("Sending organizations data")
9497
syncLogService.sendInfoLog(
9598
"The connector found ${organizations.filter { it.installed }.size} " +
@@ -105,10 +108,12 @@ class GitHubScanningService(
105108
var page = 1
106109
val repositories = mutableListOf<RepositoryDto>()
107110
do {
108-
val repositoriesPage = gitHubGraphQLService.getRepositories(
109-
token = installationToken,
110-
cursor = cursor
111-
)
111+
val repositoriesPage = rateLimitHandler.executeWithRateLimitHandler(RateLimitType.GRAPHQL) {
112+
gitHubGraphQLService.getRepositories(
113+
token = installationToken,
114+
cursor = cursor
115+
)
116+
}
112117
webSocketService.sendMessage(
113118
"${cachingService.get("runId")}/repositories",
114119
repositoriesPage.repositories.filter { !it.archived }
@@ -144,7 +149,7 @@ class GitHubScanningService(
144149

145150
private fun fetchManifestFiles(installation: Installation, repositoryName: String) = runCatching {
146151
val installationToken = cachingService.get("installationToken:${installation.id}").toString()
147-
rateLimitHandler.executeWithRateLimitHandler {
152+
rateLimitHandler.executeWithRateLimitHandler(RateLimitType.SEARCH) {
148153
gitHubClient.searchManifestFiles(
149154
"Bearer $installationToken",
150155
"" +

‎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
@@ -72,8 +72,8 @@ class GitHubScanningServiceTest {
7272
every { gitHubAuthenticationService.generateAndCacheInstallationTokens(any(), any()) } returns Unit
7373
every { syncLogService.sendErrorLog(any()) } returns Unit
7474
every { syncLogService.sendInfoLog(any()) } returns Unit
75-
every { rateLimitHandler.executeWithRateLimitHandler(any<() -> Any>()) } answers
76-
{ firstArg<() -> Any>().invoke() }
75+
every { rateLimitHandler.executeWithRateLimitHandler(any(), any<() -> Any>()) } answers
76+
{ secondArg<() -> Any>().invoke() }
7777
every { gitHubClient.getApp(any()) } returns GitHubAppResponse("testApp", permissions, events)
7878
}
7979

0 commit comments

Comments
 (0)