Merge pull request #127 from leanix/feature/CID-3926/Full-scan-not-finishing-if-git-data-hub-shuts-downs-mid-full-scan

mohamedlajmileanix · web-flow · commit ce46541eb91f · 2025-05-28T08:55:03.000+02:00
CID-3926: Full scan not finishing if git-data-hub shuts down
diff --git a/src/main/kotlin/net/leanix/githubagent/GitHubAgentApplication.kt b/src/main/kotlin/net/leanix/githubagent/GitHubAgentApplication.kt
@@ -10,9 +10,11 @@ import org.springframework.boot.autoconfigure.SpringBootApplication
 import org.springframework.boot.context.properties.EnableConfigurationProperties
 import org.springframework.boot.runApplication
 import org.springframework.cloud.openfeign.EnableFeignClients
+import org.springframework.scheduling.annotation.EnableScheduling
 
 @SpringBootApplication
 @EnableFeignClients
+@EnableScheduling
 @EnableConfigurationProperties(
     value = [
         net.leanix.githubagent.config.GitHubEnterpriseProperties::class,
diff --git a/src/main/kotlin/net/leanix/githubagent/dto/ConnectionEstablishedEvent.kt b/src/main/kotlin/net/leanix/githubagent/dto/ConnectionEstablishedEvent.kt
@@ -0,0 +1,3 @@
+package net.leanix.githubagent.dto
+
+class ConnectionEstablishedEvent
diff --git a/src/main/kotlin/net/leanix/githubagent/exceptions/Exceptions.kt b/src/main/kotlin/net/leanix/githubagent/exceptions/Exceptions.kt
@@ -14,3 +14,4 @@ class GraphQLApiException(errors: List<GraphQLClientError>) :
 class WebhookSecretNotSetException : RuntimeException("Webhook secret not set")
 class InvalidEventSignatureException : RuntimeException("Invalid event signature")
 class ManifestFileNotFoundException : RuntimeException("Manifest File Not Found")
+class UnableToSendMessageException : RuntimeException("Unable to send message to the backend")
diff --git a/src/main/kotlin/net/leanix/githubagent/handler/BrokerStompSessionHandler.kt b/src/main/kotlin/net/leanix/githubagent/handler/BrokerStompSessionHandler.kt
@@ -1,8 +1,10 @@
 package net.leanix.githubagent.handler
 
+import net.leanix.githubagent.dto.ConnectionEstablishedEvent
 import net.leanix.githubagent.services.WebSocketService
 import org.slf4j.LoggerFactory
 import org.springframework.beans.factory.annotation.Autowired
+import org.springframework.context.ApplicationEventPublisher
 import org.springframework.context.annotation.Lazy
 import org.springframework.messaging.simp.stomp.StompCommand
 import org.springframework.messaging.simp.stomp.StompHeaders
@@ -15,6 +17,7 @@ class BrokerStompSessionHandler(
     private val artifactDownloadHandler: ArtifactDownloadHandler,
     private val repositoryGetHandler: RepositoryGetHandler,
     private val installationGetHandler: InstallationGetHandler,
+    private val eventPublisher: ApplicationEventPublisher
 ) : StompSessionHandlerAdapter() {
     @Lazy
     @Autowired
@@ -30,6 +33,7 @@ class BrokerStompSessionHandler(
         session.subscribe("/user/queue/message/artifact", artifactDownloadHandler)
         session.subscribe("/user/queue/message/repository", repositoryGetHandler)
         session.subscribe("/user/queue/message/installation", installationGetHandler)
+        eventPublisher.publishEvent(ConnectionEstablishedEvent())
     }
 
     override fun handleException(
diff --git a/src/main/kotlin/net/leanix/githubagent/listener/ApplicationReadyListener.kt b/src/main/kotlin/net/leanix/githubagent/listener/ApplicationReadyListener.kt
@@ -1,6 +1,6 @@
 package net.leanix.githubagent.listener
 
-import net.leanix.githubagent.services.GitHubStartService
+import net.leanix.githubagent.services.WebSocketService
 import org.slf4j.LoggerFactory
 import org.springframework.boot.context.event.ApplicationReadyEvent
 import org.springframework.context.ApplicationListener
@@ -10,13 +10,13 @@ import org.springframework.stereotype.Component
 @Component
 @Profile("!test")
 class ApplicationReadyListener(
-    private val gitHubStartService: GitHubStartService,
+    private val webSocketService: WebSocketService,
 ) : ApplicationListener<ApplicationReadyEvent> {
 
     private val logger = LoggerFactory.getLogger(this::class.java)
 
     override fun onApplicationEvent(event: ApplicationReadyEvent) {
         logger.info("Agent started")
-        gitHubStartService.startAgent()
+        webSocketService.initSession()
     }
 }
diff --git a/src/main/kotlin/net/leanix/githubagent/services/GitHubStartService.kt b/src/main/kotlin/net/leanix/githubagent/services/GitHubStartService.kt
@@ -1,8 +1,9 @@
 package net.leanix.githubagent.services
 
-import net.leanix.githubagent.handler.BrokerStompSessionHandler
+import net.leanix.githubagent.dto.ConnectionEstablishedEvent
+import net.leanix.githubagent.exceptions.UnableToSendMessageException
 import net.leanix.githubagent.shared.APP_NAME_TOPIC
-import org.slf4j.LoggerFactory
+import org.springframework.context.event.EventListener
 import org.springframework.scheduling.annotation.Async
 import org.springframework.stereotype.Service
 
@@ -13,26 +14,25 @@ class GitHubStartService(
     private val gitHubScanningService: GitHubScanningService,
     private val gitHubEnterpriseService: GitHubEnterpriseService,
     private val cachingService: CachingService,
-    private val brokerStompSessionHandler: BrokerStompSessionHandler,
     private val syncLogService: SyncLogService
 ) {
+    var requireScan = true
 
-    private val logger = LoggerFactory.getLogger(this::class.java)
-
+    @SuppressWarnings("UnusedParameter")
     @Async
-    fun startAgent() {
-        webSocketService.initSession()
-        if (!brokerStompSessionHandler.isConnected()) {
-            logger.error("Stopping the application as the WebSocket connection could not be established.")
-            return
-        }
-        kotlin.runCatching {
-            syncLogService.sendFullScanStart(null)
-            scanResources()
-        }.onSuccess {
-            syncLogService.sendFullScanSuccess()
-        }.onFailure {
-            syncLogService.sendFullScanFailure(it.message)
+    @EventListener
+    fun startAgent(connectionEstablishedEvent: ConnectionEstablishedEvent) {
+        if (requireScan) {
+            runCatching {
+                requireScan = false
+                syncLogService.sendFullScanStart(null)
+                scanResources()
+            }.onSuccess {
+                syncLogService.sendFullScanSuccess()
+            }.onFailure {
+                if (it is UnableToSendMessageException) requireScan = true
+                syncLogService.sendFullScanFailure(it.message)
+            }
         }
     }
 
diff --git a/src/main/kotlin/net/leanix/githubagent/services/WebSocketService.kt b/src/main/kotlin/net/leanix/githubagent/services/WebSocketService.kt
@@ -1,6 +1,8 @@
 package net.leanix.githubagent.services
 
+import io.github.resilience4j.retry.annotation.Retry
 import net.leanix.githubagent.config.WebSocketClientConfig
+import net.leanix.githubagent.exceptions.UnableToSendMessageException
 import net.leanix.githubagent.handler.BrokerStompSessionHandler
 import net.leanix.githubagent.shared.TOPIC_PREFIX
 import org.slf4j.LoggerFactory
@@ -11,6 +13,7 @@ import org.springframework.stereotype.Service
 class WebSocketService(
     private val webSocketClientConfig: WebSocketClientConfig,
     private val brokerStompSessionHandler: BrokerStompSessionHandler,
+    private val cachingService: CachingService,
 ) {
 
     private val logger = LoggerFactory.getLogger(WebSocketService::class.java)
@@ -27,10 +30,15 @@ class WebSocketService(
         }
     }
 
+    @Retry(name = "sendMessageRetry")
     fun sendMessage(topic: String, data: Any) {
-        if (!brokerStompSessionHandler.isConnected()) {
+        if (!brokerStompSessionHandler.isConnected() && cachingService.get("runId") == null) {
             return
         }
-        stompSession!!.send("$TOPIC_PREFIX$topic", data)
+        runCatching {
+            stompSession!!.send("$TOPIC_PREFIX$topic", data)
+        }.onFailure {
+            throw UnableToSendMessageException()
+        }
     }
 }
diff --git a/src/test/kotlin/net/leanix/githubagent/listener/ApplicationReadyListenerTest.kt b/src/test/kotlin/net/leanix/githubagent/listener/ApplicationReadyListenerTest.kt
@@ -3,30 +3,28 @@ package net.leanix.githubagent.listener
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
-import net.leanix.githubagent.services.GitHubStartService
+import net.leanix.githubagent.services.WebSocketService
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 
 class ApplicationReadyListenerTest {
 
-    private lateinit var gitHubStartService: GitHubStartService
     private lateinit var applicationListener: ApplicationReadyListener
+    private var webSocketService: WebSocketService = mockk()
 
     @BeforeEach
     fun setUp() {
-        gitHubStartService = mockk()
-
         applicationListener = ApplicationReadyListener(
-            gitHubStartService
+            webSocketService
         )
 
-        every { gitHubStartService.startAgent() } returns Unit
+        every { webSocketService.initSession() } returns Unit
     }
 
     @Test
     fun `should start the agent process`() {
         applicationListener.onApplicationEvent(mockk())
 
-        verify { gitHubStartService.startAgent() }
+        verify { webSocketService.initSession() }
     }
 }
diff --git a/src/test/kotlin/net/leanix/githubagent/services/GitHubStartServiceTest.kt b/src/test/kotlin/net/leanix/githubagent/services/GitHubStartServiceTest.kt
@@ -3,6 +3,7 @@ package net.leanix.githubagent.services
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
+import net.leanix.githubagent.dto.ConnectionEstablishedEvent
 import net.leanix.githubagent.dto.GitHubAppResponse
 import net.leanix.githubagent.handler.BrokerStompSessionHandler
 import net.leanix.githubagent.shared.APP_NAME_TOPIC
@@ -36,7 +37,6 @@ class GitHubStartServiceTest {
             gitHubScanningService,
             gitHubEnterpriseService,
             cachingService,
-            brokerStompSessionHandler,
             syncLogService
         )
 
@@ -61,7 +61,7 @@ class GitHubStartServiceTest {
                 gitHubAppName, mapOf(), listOf()
             )
 
-        gitHubStartService.startAgent()
+        gitHubStartService.startAgent(ConnectionEstablishedEvent())
 
         verify { webSocketService.sendMessage(APP_NAME_TOPIC, gitHubAppName) }
         verify { syncLogService.sendFullScanStart(any()) }
diff --git a/src/test/kotlin/net/leanix/githubagent/services/WebSocketServiceTests.kt b/src/test/kotlin/net/leanix/githubagent/services/WebSocketServiceTests.kt
@@ -15,10 +15,11 @@ class WebSocketServiceTests {
     private val webSocketClientConfig: WebSocketClientConfig = mockk()
     private val stompSession: StompSession = mockk()
     private val brokerStompSessionHandler: BrokerStompSessionHandler = mockk()
+    private val cachingService: CachingService = mockk()
 
     @BeforeEach
     fun setUp() {
-        webSocketService = WebSocketService(webSocketClientConfig, brokerStompSessionHandler)
+        webSocketService = WebSocketService(webSocketClientConfig, brokerStompSessionHandler, cachingService)
     }
 
     @Test

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+package net.leanix.githubagent.dto`
	`2`	`+`
	`3`	`+class ConnectionEstablishedEvent`