Merge pull request #25 from leanix/feature/CID-2939/improve-connectio…

…n-closed-error-log

CID-2939: Improve connection error logs. Fix jwtToken variable name. …

src/main/kotlin/net/leanix/githubagent/handler/BrokerStompSessionHandler.kt

@@ -9,6 +9,7 @@ import org.springframework.messaging.simp.stomp.StompHeaders
 import org.springframework.messaging.simp.stomp.StompSession
 import org.springframework.messaging.simp.stomp.StompSessionHandlerAdapter
 import org.springframework.stereotype.Component
+import java.util.concurrent.CountDownLatch
 
 @Component
 class BrokerStompSessionHandler : StompSessionHandlerAdapter() {
@@ -20,9 +21,12 @@ class BrokerStompSessionHandler : StompSessionHandlerAdapter() {
 
     private var isConnected = false
 
+    private val latch = CountDownLatch(1)
+
     override fun afterConnected(session: StompSession, connectedHeaders: StompHeaders) {
         logger.info("connected to the server: ${session.sessionId}")
         isConnected = true
+        latch.countDown()
         session.subscribe("/user/queue/repositories-string", this)
     }
 
@@ -38,11 +42,23 @@ class BrokerStompSessionHandler : StompSessionHandlerAdapter() {
     }
 
     override fun handleTransportError(session: StompSession, exception: Throwable) {
-        logger.error("handleTransportError", exception)
+        logger.error("Connection error: ${exception.message}")
         if (isConnected) {
-            logger.info("session closed")
+            logger.error("Session closed. This could be due to a network error or the server closing the connection.")
             isConnected = false
+            logger.info("Reconnecting...")
             webSocketService.initSession()
+        } else {
+            if (latch.count != 0L) latch.countDown()
         }
     }
+
+    fun isConnected(): Boolean {
+        awaitConnection()
+        return isConnected
+    }
+
+    private fun awaitConnection() {
+        latch.await()
+    }
 }

src/main/kotlin/net/leanix/githubagent/runners/PostStartupRunner.kt

@@ -1,11 +1,13 @@
 package net.leanix.githubagent.runners
 
+import net.leanix.githubagent.handler.BrokerStompSessionHandler
 import net.leanix.githubagent.services.CachingService
 import net.leanix.githubagent.services.GitHubAuthenticationService
 import net.leanix.githubagent.services.GitHubEnterpriseService
 import net.leanix.githubagent.services.GitHubScanningService
 import net.leanix.githubagent.services.WebSocketService
 import net.leanix.githubagent.shared.AGENT_METADATA_TOPIC
+import org.slf4j.LoggerFactory
 import org.springframework.boot.ApplicationArguments
 import org.springframework.boot.ApplicationRunner
 import org.springframework.context.annotation.Profile
@@ -18,13 +20,20 @@ class PostStartupRunner(
     private val webSocketService: WebSocketService,
     private val gitHubScanningService: GitHubScanningService,
     private val gitHubEnterpriseService: GitHubEnterpriseService,
-    private val cachingService: CachingService
+    private val cachingService: CachingService,
+    private val brokerStompSessionHandler: BrokerStompSessionHandler
 ) : ApplicationRunner {
 
+    private val logger = LoggerFactory.getLogger(PostStartupRunner::class.java)
+
     override fun run(args: ApplicationArguments?) {
         webSocketService.initSession()
+        if (!brokerStompSessionHandler.isConnected()) {
+            logger.error("Stopping the application as the WebSocket connection could not be established.")
+            return
+        }
         githubAuthenticationService.generateAndCacheJwtToken()
-        val jwt = cachingService.get("jwt") as String
+        val jwt = cachingService.get("jwtToken") as String
         webSocketService.sendMessage(
             AGENT_METADATA_TOPIC,
             gitHubEnterpriseService.getGitHubApp(jwt).name

src/main/kotlin/net/leanix/githubagent/services/WebSocketService.kt

@@ -15,8 +15,12 @@ class WebSocketService(
     var stompSession: StompSession? = null
 
     fun initSession() {
-        logger.info("init session")
-        stompSession = webSocketClientConfig.initSession()
+        logger.info("Initializing websocket session")
+        kotlin.runCatching {
+            stompSession = webSocketClientConfig.initSession()
+        }.onFailure {
+            logger.error("Failed to initialize WebSocket session")
+        }
     }
 
     fun sendMessage(topic: String, data: Any) {

src/test/kotlin/net/leanix/githubagent/runners/PostStartupRunnerTest.kt

@@ -4,6 +4,7 @@ import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
 import net.leanix.githubagent.dto.GitHubAppResponse
+import net.leanix.githubagent.handler.BrokerStompSessionHandler
 import net.leanix.githubagent.services.CachingService
 import net.leanix.githubagent.services.GitHubAuthenticationService
 import net.leanix.githubagent.services.GitHubEnterpriseService
@@ -21,6 +22,7 @@ class PostStartupRunnerTest {
     private lateinit var gitHubEnterpriseService: GitHubEnterpriseService
     private lateinit var cachingService: CachingService
     private lateinit var postStartupRunner: PostStartupRunner
+    private lateinit var brokerStompSessionHandler: BrokerStompSessionHandler
 
     @BeforeEach
     fun setUp() {
@@ -29,20 +31,23 @@ class PostStartupRunnerTest {
         gitHubScanningService = mockk()
         gitHubEnterpriseService = mockk()
         cachingService = mockk()
+        brokerStompSessionHandler = mockk()
 
         postStartupRunner = PostStartupRunner(
             githubAuthenticationService,
             webSocketService,
             gitHubScanningService,
             gitHubEnterpriseService,
-            cachingService
+            cachingService,
+            brokerStompSessionHandler
         )
 
         every { webSocketService.initSession() } returns Unit
         every { webSocketService.sendMessage(any(), any()) } returns Unit
         every { githubAuthenticationService.generateAndCacheJwtToken() } returns Unit
-        every { cachingService.get("jwt") } returns "jwt"
+        every { cachingService.get("jwtToken") } returns "jwt"
         every { gitHubScanningService.scanGitHubResources() } returns Unit
+        every { brokerStompSessionHandler.isConnected() } returns true
     }
 
     @Test