CID-3277: Add rate limit handling for manifest file searching

leanix · Dec 13, 2024 · a109679 · a109679
1 parent 9dd45ea
commit a109679
Show file tree

Hide file tree

Showing 7 changed files with 136 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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 <T> 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
@@ -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
@@ -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<Installation> {
-        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
@@ -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
@@ -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<GitHubGraphQLService>()
     private val gitHubAuthenticationService = mockk<GitHubAuthenticationService>()
     private val syncLogService = mockk<SyncLogService>(relaxUnitFun = true)
+    private val rateLimitHandler = mockk<RateLimitHandler>(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