diff --git a/src/main/kotlin/net/leanix/githubagent/handler/BrokerStompSessionHandler.kt b/src/main/kotlin/net/leanix/githubagent/handler/BrokerStompSessionHandler.kt index 60e5c7d..37118f1 100644 --- a/src/main/kotlin/net/leanix/githubagent/handler/BrokerStompSessionHandler.kt +++ b/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() + } } diff --git a/src/main/kotlin/net/leanix/githubagent/runners/PostStartupRunner.kt b/src/main/kotlin/net/leanix/githubagent/runners/PostStartupRunner.kt index 1abb4da..0434c74 100644 --- a/src/main/kotlin/net/leanix/githubagent/runners/PostStartupRunner.kt +++ b/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 diff --git a/src/main/kotlin/net/leanix/githubagent/services/WebSocketService.kt b/src/main/kotlin/net/leanix/githubagent/services/WebSocketService.kt index 48bf42a..9f23aeb 100644 --- a/src/main/kotlin/net/leanix/githubagent/services/WebSocketService.kt +++ b/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) { diff --git a/src/test/kotlin/net/leanix/githubagent/runners/PostStartupRunnerTest.kt b/src/test/kotlin/net/leanix/githubagent/runners/PostStartupRunnerTest.kt index 4c30788..3b3309a 100644 --- a/src/test/kotlin/net/leanix/githubagent/runners/PostStartupRunnerTest.kt +++ b/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