From e80921266f88ef3f2bc4319a72b6f1f7c1161096 Mon Sep 17 00:00:00 2001 From: mohamedlajmileanix Date: Tue, 6 May 2025 07:00:20 +0200 Subject: [PATCH 1/5] CID-3792: Implement heartbeat support --- .../config/WebSocketClientConfig.kt | 32 ++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt b/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt index 6817739..449459c 100644 --- a/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt +++ b/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt @@ -5,6 +5,7 @@ import io.github.resilience4j.retry.annotation.Retry import net.leanix.githubagent.handler.BrokerStompSessionHandler import net.leanix.githubagent.services.LeanIXAuthService import net.leanix.githubagent.shared.GitHubAgentProperties.GITHUB_AGENT_VERSION +import org.slf4j.LoggerFactory import org.springframework.context.annotation.Bean import org.springframework.context.annotation.Configuration import org.springframework.messaging.converter.MappingJackson2MessageConverter @@ -16,6 +17,7 @@ import org.springframework.web.socket.client.standard.StandardWebSocketClient import org.springframework.web.socket.messaging.WebSocketStompClient import org.springframework.web.socket.sockjs.client.SockJsClient import org.springframework.web.socket.sockjs.client.WebSocketTransport +import java.util.concurrent.ScheduledFuture @Configuration class WebSocketClientConfig( @@ -25,6 +27,9 @@ class WebSocketClientConfig( private val leanIXProperties: LeanIXProperties, private val gitHubEnterpriseProperties: GitHubEnterpriseProperties ) { + private var heartbeatTask: ScheduledFuture<*>? = null + private val logger = LoggerFactory.getLogger(WebSocketClientConfig::class.java) + @Retry(name = "ws_init_session") fun initSession(): StompSession { val headers = WebSocketHttpHeaders() @@ -32,12 +37,37 @@ class WebSocketClientConfig( stompHeaders["Authorization"] = "Bearer ${leanIXAuthService.getBearerToken()}" stompHeaders["GitHub-Enterprise-URL"] = gitHubEnterpriseProperties.baseUrl stompHeaders["GitHub-Agent-Version"] = GITHUB_AGENT_VERSION - return stompClient().connectAsync( + val session = stompClient().connectAsync( leanIXProperties.wsBaseUrl, headers, stompHeaders, brokerStompSessionHandler, ).get() + + sendHeartbeat(session) + return session + } + + private fun sendHeartbeat(session: StompSession) { + val scheduler = ThreadPoolTaskScheduler() + scheduler.initialize() + heartbeatTask = scheduler.scheduleAtFixedRate({ + kotlin.runCatching { + if (session.isConnected) { + session.send("/app/ghe/heartbeat", "") + logger.debug("Heartbeat sent to /app/heartbeat") + } else { + logger.warn("Session is not connected, stopping heartbeat") + stopHeartbeat() + } + }.onFailure { + logger.error("Failed to send heartbeat: ${it.message}") + } + }, java.time.Duration.ofSeconds(10)) + } + + fun stopHeartbeat() { + heartbeatTask?.cancel(true) } @Bean From 7f3b80779e5c080d7f8527acdbedccfb297f6814 Mon Sep 17 00:00:00 2001 From: mohamedlajmileanix Date: Tue, 6 May 2025 08:29:14 +0200 Subject: [PATCH 2/5] CID-3792: Add tests --- .../githubagent/config/WebSocketClientConfig.kt | 2 +- .../config/WebSocketClientConfigTests.kt | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt b/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt index 449459c..f8adb07 100644 --- a/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt +++ b/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt @@ -48,7 +48,7 @@ class WebSocketClientConfig( return session } - private fun sendHeartbeat(session: StompSession) { + fun sendHeartbeat(session: StompSession) { val scheduler = ThreadPoolTaskScheduler() scheduler.initialize() heartbeatTask = scheduler.scheduleAtFixedRate({ diff --git a/src/test/kotlin/net/leanix/githubagent/config/WebSocketClientConfigTests.kt b/src/test/kotlin/net/leanix/githubagent/config/WebSocketClientConfigTests.kt index 62c596f..86a37f4 100644 --- a/src/test/kotlin/net/leanix/githubagent/config/WebSocketClientConfigTests.kt +++ b/src/test/kotlin/net/leanix/githubagent/config/WebSocketClientConfigTests.kt @@ -2,7 +2,9 @@ package net.leanix.githubagent.config import com.fasterxml.jackson.databind.ObjectMapper import io.mockk.coEvery +import io.mockk.every import io.mockk.mockk +import io.mockk.verify import kotlinx.coroutines.runBlocking import net.leanix.githubagent.handler.BrokerStompSessionHandler import net.leanix.githubagent.services.LeanIXAuthService @@ -14,6 +16,7 @@ import org.springframework.messaging.simp.stomp.StompHeaders import org.springframework.messaging.simp.stomp.StompSession import org.springframework.web.socket.WebSocketHttpHeaders import org.springframework.web.socket.messaging.WebSocketStompClient +import java.util.concurrent.ScheduledFuture class WebSocketClientConfigTests { private lateinit var webSocketClientConfig: WebSocketClientConfig @@ -23,6 +26,7 @@ class WebSocketClientConfigTests { private lateinit var leanIXProperties: LeanIXProperties private lateinit var gitHubEnterpriseProperties: GitHubEnterpriseProperties private lateinit var leanIXAuthService: LeanIXAuthService + private lateinit var scheduledFuture: ScheduledFuture<*> @BeforeEach fun setUp() { @@ -41,6 +45,7 @@ class WebSocketClientConfigTests { leanIXProperties, gitHubEnterpriseProperties ) + scheduledFuture = mockk() GitHubAgentProperties.GITHUB_AGENT_VERSION = "test-version" } @@ -63,4 +68,15 @@ class WebSocketClientConfigTests { assertEquals(null, session) } + + @Test + fun `should send heartbeat when session is connected`() { + val receiptable = mockk() + every { stompSession.isConnected } returns true + every { stompSession.send(any(), any()) } returns receiptable + + webSocketClientConfig.sendHeartbeat(stompSession) + + verify { stompSession.send("/app/ghe/heartbeat", "") } + } } From 28167044a1e4cd009854da1abcbf724d64dc0c1e Mon Sep 17 00:00:00 2001 From: mohamedlajmileanix Date: Tue, 6 May 2025 09:49:16 +0200 Subject: [PATCH 3/5] CID-3792: Simplify implementation --- .../config/WebSocketClientConfig.kt | 33 ++----------------- 1 file changed, 2 insertions(+), 31 deletions(-) diff --git a/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt b/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt index f8adb07..a0a0caa 100644 --- a/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt +++ b/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt @@ -5,7 +5,6 @@ import io.github.resilience4j.retry.annotation.Retry import net.leanix.githubagent.handler.BrokerStompSessionHandler import net.leanix.githubagent.services.LeanIXAuthService import net.leanix.githubagent.shared.GitHubAgentProperties.GITHUB_AGENT_VERSION -import org.slf4j.LoggerFactory import org.springframework.context.annotation.Bean import org.springframework.context.annotation.Configuration import org.springframework.messaging.converter.MappingJackson2MessageConverter @@ -17,7 +16,6 @@ import org.springframework.web.socket.client.standard.StandardWebSocketClient import org.springframework.web.socket.messaging.WebSocketStompClient import org.springframework.web.socket.sockjs.client.SockJsClient import org.springframework.web.socket.sockjs.client.WebSocketTransport -import java.util.concurrent.ScheduledFuture @Configuration class WebSocketClientConfig( @@ -27,9 +25,6 @@ class WebSocketClientConfig( private val leanIXProperties: LeanIXProperties, private val gitHubEnterpriseProperties: GitHubEnterpriseProperties ) { - private var heartbeatTask: ScheduledFuture<*>? = null - private val logger = LoggerFactory.getLogger(WebSocketClientConfig::class.java) - @Retry(name = "ws_init_session") fun initSession(): StompSession { val headers = WebSocketHttpHeaders() @@ -37,37 +32,12 @@ class WebSocketClientConfig( stompHeaders["Authorization"] = "Bearer ${leanIXAuthService.getBearerToken()}" stompHeaders["GitHub-Enterprise-URL"] = gitHubEnterpriseProperties.baseUrl stompHeaders["GitHub-Agent-Version"] = GITHUB_AGENT_VERSION - val session = stompClient().connectAsync( + return stompClient().connectAsync( leanIXProperties.wsBaseUrl, headers, stompHeaders, brokerStompSessionHandler, ).get() - - sendHeartbeat(session) - return session - } - - fun sendHeartbeat(session: StompSession) { - val scheduler = ThreadPoolTaskScheduler() - scheduler.initialize() - heartbeatTask = scheduler.scheduleAtFixedRate({ - kotlin.runCatching { - if (session.isConnected) { - session.send("/app/ghe/heartbeat", "") - logger.debug("Heartbeat sent to /app/heartbeat") - } else { - logger.warn("Session is not connected, stopping heartbeat") - stopHeartbeat() - } - }.onFailure { - logger.error("Failed to send heartbeat: ${it.message}") - } - }, java.time.Duration.ofSeconds(10)) - } - - fun stopHeartbeat() { - heartbeatTask?.cancel(true) } @Bean @@ -80,6 +50,7 @@ class WebSocketClientConfig( val sockJsClient = SockJsClient(transports) val stompClient = WebSocketStompClient(sockJsClient) stompClient.messageConverter = jsonConverter + stompClient.defaultHeartbeat = longArrayOf(3000, 3000) val scheduler = ThreadPoolTaskScheduler() scheduler.initialize() stompClient.taskScheduler = scheduler From 893abae14b247f14770cd14693ae07c2db93e976 Mon Sep 17 00:00:00 2001 From: mohamedlajmileanix Date: Tue, 6 May 2025 09:51:41 +0200 Subject: [PATCH 4/5] CID-3792: Simplify implementation --- .../config/WebSocketClientConfigTests.kt | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/src/test/kotlin/net/leanix/githubagent/config/WebSocketClientConfigTests.kt b/src/test/kotlin/net/leanix/githubagent/config/WebSocketClientConfigTests.kt index 86a37f4..62c596f 100644 --- a/src/test/kotlin/net/leanix/githubagent/config/WebSocketClientConfigTests.kt +++ b/src/test/kotlin/net/leanix/githubagent/config/WebSocketClientConfigTests.kt @@ -2,9 +2,7 @@ package net.leanix.githubagent.config import com.fasterxml.jackson.databind.ObjectMapper import io.mockk.coEvery -import io.mockk.every import io.mockk.mockk -import io.mockk.verify import kotlinx.coroutines.runBlocking import net.leanix.githubagent.handler.BrokerStompSessionHandler import net.leanix.githubagent.services.LeanIXAuthService @@ -16,7 +14,6 @@ import org.springframework.messaging.simp.stomp.StompHeaders import org.springframework.messaging.simp.stomp.StompSession import org.springframework.web.socket.WebSocketHttpHeaders import org.springframework.web.socket.messaging.WebSocketStompClient -import java.util.concurrent.ScheduledFuture class WebSocketClientConfigTests { private lateinit var webSocketClientConfig: WebSocketClientConfig @@ -26,7 +23,6 @@ class WebSocketClientConfigTests { private lateinit var leanIXProperties: LeanIXProperties private lateinit var gitHubEnterpriseProperties: GitHubEnterpriseProperties private lateinit var leanIXAuthService: LeanIXAuthService - private lateinit var scheduledFuture: ScheduledFuture<*> @BeforeEach fun setUp() { @@ -45,7 +41,6 @@ class WebSocketClientConfigTests { leanIXProperties, gitHubEnterpriseProperties ) - scheduledFuture = mockk() GitHubAgentProperties.GITHUB_AGENT_VERSION = "test-version" } @@ -68,15 +63,4 @@ class WebSocketClientConfigTests { assertEquals(null, session) } - - @Test - fun `should send heartbeat when session is connected`() { - val receiptable = mockk() - every { stompSession.isConnected } returns true - every { stompSession.send(any(), any()) } returns receiptable - - webSocketClientConfig.sendHeartbeat(stompSession) - - verify { stompSession.send("/app/ghe/heartbeat", "") } - } } From 1b41a1876a81a15c80e5a6c278a1dd0e238e4b65 Mon Sep 17 00:00:00 2001 From: mohamedlajmileanix Date: Tue, 6 May 2025 09:57:48 +0200 Subject: [PATCH 5/5] CID-3792: Update interval between heartbeats --- .../net/leanix/githubagent/config/WebSocketClientConfig.kt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt b/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt index a0a0caa..d2b6c70 100644 --- a/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt +++ b/src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt @@ -50,7 +50,7 @@ class WebSocketClientConfig( val sockJsClient = SockJsClient(transports) val stompClient = WebSocketStompClient(sockJsClient) stompClient.messageConverter = jsonConverter - stompClient.defaultHeartbeat = longArrayOf(3000, 3000) + stompClient.defaultHeartbeat = longArrayOf(10000, 10000) val scheduler = ThreadPoolTaskScheduler() scheduler.initialize() stompClient.taskScheduler = scheduler