From a109679060c803ed252e1745e7449219f13790fb Mon Sep 17 00:00:00 2001 From: mohamedlajmileanix Date: Fri, 13 Dec 2024 16:01:40 +0100 Subject: [PATCH] CID-3277: Add rate limit handling for manifest file searching --- .../leanix/githubagent/client/GitHubClient.kt | 7 ++- .../githubagent/config/FeignClientConfig.kt | 15 ++++++ .../githubagent/handler/RateLimitHandler.kt | 22 +++++++++ .../RateLimitResponseInterceptor.kt | 28 +++++++++++ .../services/GitHubScanningService.kt | 20 +++++--- .../githubagent/shared/RateLimitMonitor.kt | 46 +++++++++++++++++++ .../services/GitHubScanningServiceTest.kt | 7 ++- 7 files changed, 136 insertions(+), 9 deletions(-) create mode 100644 src/main/kotlin/net/leanix/githubagent/config/FeignClientConfig.kt create mode 100644 src/main/kotlin/net/leanix/githubagent/handler/RateLimitHandler.kt create mode 100644 src/main/kotlin/net/leanix/githubagent/interceptor/RateLimitResponseInterceptor.kt create mode 100644 src/main/kotlin/net/leanix/githubagent/shared/RateLimitMonitor.kt diff --git a/src/main/kotlin/net/leanix/githubagent/client/GitHubClient.kt b/src/main/kotlin/net/leanix/githubagent/client/GitHubClient.kt index d333087..bb3b5b3 100644 --- a/src/main/kotlin/net/leanix/githubagent/client/GitHubClient.kt +++ b/src/main/kotlin/net/leanix/githubagent/client/GitHubClient.kt @@ -1,5 +1,6 @@ package net.leanix.githubagent.client +import net.leanix.githubagent.config.FeignClientConfig import net.leanix.githubagent.dto.GitHubAppResponse import net.leanix.githubagent.dto.GitHubSearchResponse import net.leanix.githubagent.dto.Installation @@ -13,7 +14,11 @@ import org.springframework.web.bind.annotation.PostMapping import org.springframework.web.bind.annotation.RequestHeader import org.springframework.web.bind.annotation.RequestParam -@FeignClient(name = "githubClient", url = "\${github-enterprise.baseUrl}") +@FeignClient( + name = "githubClient", + url = "\${github-enterprise.baseUrl}", + configuration = [FeignClientConfig::class] +) interface GitHubClient { @GetMapping("/api/v3/app") diff --git a/src/main/kotlin/net/leanix/githubagent/config/FeignClientConfig.kt b/src/main/kotlin/net/leanix/githubagent/config/FeignClientConfig.kt new file mode 100644 index 0000000..67ee070 --- /dev/null +++ b/src/main/kotlin/net/leanix/githubagent/config/FeignClientConfig.kt @@ -0,0 +1,15 @@ +package net.leanix.githubagent.config + +import feign.ResponseInterceptor +import net.leanix.githubagent.interceptor.RateLimitResponseInterceptor +import org.springframework.context.annotation.Bean +import org.springframework.context.annotation.Configuration + +@Configuration +class FeignClientConfig { + + @Bean + fun rateLimitResponseInterceptor(): ResponseInterceptor { + return RateLimitResponseInterceptor() + } +} diff --git a/src/main/kotlin/net/leanix/githubagent/handler/RateLimitHandler.kt b/src/main/kotlin/net/leanix/githubagent/handler/RateLimitHandler.kt new file mode 100644 index 0000000..d051679 --- /dev/null +++ b/src/main/kotlin/net/leanix/githubagent/handler/RateLimitHandler.kt @@ -0,0 +1,22 @@ +package net.leanix.githubagent.handler + +import net.leanix.githubagent.services.SyncLogService +import net.leanix.githubagent.shared.RateLimitMonitor +import org.springframework.stereotype.Component + +@Component +class RateLimitHandler( + private val syncLogService: SyncLogService, +) { + + fun executeWithRateLimitHandler(block: () -> T): T { + while (true) { + val waitTimeSeconds = RateLimitMonitor.shouldThrottle() + if (waitTimeSeconds > 0) { + syncLogService.sendInfoLog("Approaching rate limit. Waiting for $waitTimeSeconds seconds.") + Thread.sleep(waitTimeSeconds * 1000) + } + return block() + } + } +} diff --git a/src/main/kotlin/net/leanix/githubagent/interceptor/RateLimitResponseInterceptor.kt b/src/main/kotlin/net/leanix/githubagent/interceptor/RateLimitResponseInterceptor.kt new file mode 100644 index 0000000..ebdcb1c --- /dev/null +++ b/src/main/kotlin/net/leanix/githubagent/interceptor/RateLimitResponseInterceptor.kt @@ -0,0 +1,28 @@ +package net.leanix.githubagent.interceptor + +import feign.InvocationContext +import feign.ResponseInterceptor +import net.leanix.githubagent.shared.RateLimitMonitor + +class RateLimitResponseInterceptor : ResponseInterceptor { + + override fun intercept( + invocationContext: InvocationContext, + chain: ResponseInterceptor.Chain + ): Any { + val result = chain.next(invocationContext) + + val response = invocationContext.response() + if (response != null) { + val headers = response.headers().mapKeys { it.key.lowercase() } + val rateLimitRemaining = headers["x-ratelimit-remaining"]?.firstOrNull()?.toIntOrNull() + val rateLimitReset = headers["x-ratelimit-reset"]?.firstOrNull()?.toLongOrNull() + + if (rateLimitRemaining != null && rateLimitReset != null) { + RateLimitMonitor.updateRateLimitInfo(rateLimitRemaining, rateLimitReset) + } + } + + return result + } +} diff --git a/src/main/kotlin/net/leanix/githubagent/services/GitHubScanningService.kt b/src/main/kotlin/net/leanix/githubagent/services/GitHubScanningService.kt index 7b5e3b2..cb27268 100644 --- a/src/main/kotlin/net/leanix/githubagent/services/GitHubScanningService.kt +++ b/src/main/kotlin/net/leanix/githubagent/services/GitHubScanningService.kt @@ -10,6 +10,7 @@ import net.leanix.githubagent.dto.OrganizationDto import net.leanix.githubagent.dto.RepositoryDto import net.leanix.githubagent.exceptions.JwtTokenNotFound import net.leanix.githubagent.exceptions.ManifestFileNotFoundException +import net.leanix.githubagent.handler.RateLimitHandler import net.leanix.githubagent.shared.MANIFEST_FILE_NAME import net.leanix.githubagent.shared.fileNameMatchRegex import net.leanix.githubagent.shared.generateFullPath @@ -23,7 +24,8 @@ class GitHubScanningService( private val webSocketService: WebSocketService, private val gitHubGraphQLService: GitHubGraphQLService, private val gitHubAuthenticationService: GitHubAuthenticationService, - private val syncLogService: SyncLogService + private val syncLogService: SyncLogService, + private val rateLimitHandler: RateLimitHandler, ) { private val logger = LoggerFactory.getLogger(GitHubScanningService::class.java) @@ -43,7 +45,9 @@ class GitHubScanningService( } private fun getInstallations(jwtToken: String): List { - val installations = gitHubClient.getInstallations("Bearer $jwtToken") + val installations = rateLimitHandler.executeWithRateLimitHandler { + gitHubClient.getInstallations("Bearer $jwtToken") + } gitHubAuthenticationService.generateAndCacheInstallationTokens(installations, jwtToken) return installations } @@ -118,11 +122,13 @@ class GitHubScanningService( private fun fetchManifestFiles(installation: Installation, repositoryName: String) = runCatching { val installationToken = cachingService.get("installationToken:${installation.id}").toString() - gitHubClient.searchManifestFiles( - "Bearer $installationToken", - "" + - "repo:${installation.account.login}/$repositoryName filename:$MANIFEST_FILE_NAME" - ) + rateLimitHandler.executeWithRateLimitHandler { + gitHubClient.searchManifestFiles( + "Bearer $installationToken", + "" + + "repo:${installation.account.login}/$repositoryName filename:$MANIFEST_FILE_NAME" + ) + } } private fun fetchManifestContents( installation: Installation, diff --git a/src/main/kotlin/net/leanix/githubagent/shared/RateLimitMonitor.kt b/src/main/kotlin/net/leanix/githubagent/shared/RateLimitMonitor.kt new file mode 100644 index 0000000..c0cd666 --- /dev/null +++ b/src/main/kotlin/net/leanix/githubagent/shared/RateLimitMonitor.kt @@ -0,0 +1,46 @@ +package net.leanix.githubagent.shared + +import org.slf4j.LoggerFactory + +object RateLimitMonitor { + + private val logger = LoggerFactory.getLogger(RateLimitMonitor::class.java) + + @Volatile + private var rateLimitRemaining: Int = Int.MAX_VALUE + + @Volatile + private var rateLimitResetTime: Long = 0 + + private val lock = Any() + + fun updateRateLimitInfo(remaining: Int, resetTimeEpochSeconds: Long) { + synchronized(lock) { + rateLimitRemaining = remaining + rateLimitResetTime = resetTimeEpochSeconds + } + } + + fun shouldThrottle(): Long { + synchronized(lock) { + if (rateLimitRemaining <= THRESHOLD) { + val currentTimeSeconds = System.currentTimeMillis() / 1000 + val waitTimeSeconds = rateLimitResetTime - currentTimeSeconds + 5 + + val adjustedWaitTime = if (waitTimeSeconds > 0) waitTimeSeconds else 0 + logger.warn( + "Rate limit remaining ($rateLimitRemaining) is at or below threshold ($THRESHOLD)." + + " Throttling for $adjustedWaitTime seconds." + ) + return adjustedWaitTime + } else { + logger.debug( + "Rate limit remaining ($rateLimitRemaining) is above threshold ($THRESHOLD). No need to throttle." + ) + } + return 0 + } + } + + private const val THRESHOLD = 2 +} diff --git a/src/test/kotlin/net/leanix/githubagent/services/GitHubScanningServiceTest.kt b/src/test/kotlin/net/leanix/githubagent/services/GitHubScanningServiceTest.kt index 3a2164e..2ec89e7 100644 --- a/src/test/kotlin/net/leanix/githubagent/services/GitHubScanningServiceTest.kt +++ b/src/test/kotlin/net/leanix/githubagent/services/GitHubScanningServiceTest.kt @@ -17,6 +17,7 @@ import net.leanix.githubagent.dto.RepositoryDto import net.leanix.githubagent.dto.RepositoryItemResponse import net.leanix.githubagent.exceptions.JwtTokenNotFound import net.leanix.githubagent.graphql.data.enums.RepositoryVisibility +import net.leanix.githubagent.handler.RateLimitHandler import net.leanix.githubagent.shared.MANIFEST_FILE_NAME import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.BeforeEach @@ -32,13 +33,15 @@ class GitHubScanningServiceTest { private val gitHubGraphQLService = mockk() private val gitHubAuthenticationService = mockk() private val syncLogService = mockk(relaxUnitFun = true) + private val rateLimitHandler = mockk(relaxUnitFun = true) private val gitHubScanningService = GitHubScanningService( gitHubClient, cachingService, webSocketService, gitHubGraphQLService, gitHubAuthenticationService, - syncLogService + syncLogService, + rateLimitHandler ) private val runId = UUID.randomUUID() @@ -61,6 +64,8 @@ class GitHubScanningServiceTest { every { gitHubAuthenticationService.generateAndCacheInstallationTokens(any(), any()) } returns Unit every { syncLogService.sendErrorLog(any()) } returns Unit every { syncLogService.sendInfoLog(any()) } returns Unit + every { rateLimitHandler.executeWithRateLimitHandler(any<() -> Any>()) } answers + { firstArg<() -> Any>().invoke() } } @Test