Merge pull request #118 from leanix/feature/CID-3856/Frequent-Websocket-connection-loss

CID-3856: Switching from default heartbeats to manual heartbeats

Author: mohamedlajmileanix

src/main/kotlin/net/leanix/githubagent/config/WebSocketClientConfig.kt

@@ -5,7 +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.springframework.beans.factory.annotation.Value
+import org.slf4j.LoggerFactory
 import org.springframework.context.annotation.Bean
 import org.springframework.context.annotation.Configuration
 import org.springframework.messaging.converter.MappingJackson2MessageConverter
@@ -17,6 +17,8 @@ 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.time.Duration
+import java.util.concurrent.ScheduledFuture
 
 @Configuration
 class WebSocketClientConfig(
@@ -25,21 +27,49 @@ class WebSocketClientConfig(
     private val leanIXAuthService: LeanIXAuthService,
     private val leanIXProperties: LeanIXProperties,
     private val gitHubEnterpriseProperties: GitHubEnterpriseProperties,
-    @Value("\${websocket.heartbeat-interval}") private val heartbeatInterval: Long
 ) {
+
+    private var heartbeatTask: ScheduledFuture<*>? = null
+    private val logger = LoggerFactory.getLogger(WebSocketClientConfig::class.java)
+
     @Retry(name = "ws_init_session")
     fun initSession(): StompSession {
         val headers = WebSocketHttpHeaders()
         val stompHeaders = StompHeaders()
         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
+    }
+
+    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}")
+            }
+        }, Duration.ofSeconds(60))
+    }
+
+    fun stopHeartbeat() {
+        heartbeatTask?.cancel(true)
     }
 
     @Bean
@@ -52,7 +82,6 @@ class WebSocketClientConfig(
         val sockJsClient = SockJsClient(transports)
         val stompClient = WebSocketStompClient(sockJsClient)
         stompClient.messageConverter = jsonConverter
-        stompClient.defaultHeartbeat = longArrayOf(heartbeatInterval, 0)
         val scheduler = ThreadPoolTaskScheduler()
         scheduler.initialize()
         stompClient.taskScheduler = scheduler

src/main/resources/application.yaml

@@ -12,8 +12,6 @@ leanix:
     technical-user-token: ${LEANIX_TECHNICAL_USER_TOKEN}
 webhookEventService:
   waitingTime: 10000
-websocket:
-  heartbeat-interval: 30000
 
 resilience4j.retry:
   configs:

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() {
@@ -34,15 +38,14 @@ class WebSocketClientConfigTests {
         stompSession = mockk()
         authService = mockk()
         leanIXAuthService = mockk()
-        val heartbeatInterval = 10000L
+        scheduledFuture = mockk()
 
         webSocketClientConfig = WebSocketClientConfig(
             brokerStompSessionHandler,
             objectMapper,
             leanIXAuthService,
             leanIXProperties,
-            gitHubEnterpriseProperties,
-            heartbeatInterval
+            gitHubEnterpriseProperties
         )
 
         GitHubAgentProperties.GITHUB_AGENT_VERSION = "test-version"
@@ -66,4 +69,15 @@ class WebSocketClientConfigTests {
 
         assertEquals(null, session)
     }
+
+    @Test
+    fun `should send heartbeat when session is connected`() {
+        val receiptable = mockk<StompSession.Receiptable>()
+        every { stompSession.isConnected } returns true
+        every { stompSession.send(any<String>(), any()) } returns receiptable
+
+        webSocketClientConfig.sendHeartbeat(stompSession)
+
+        verify { stompSession.send("/app/ghe/heartbeat", "") }
+    }
 }

src/test/resources/application.yaml

@@ -12,5 +12,3 @@ leanix:
         technical-user-token: ${LEANIX_TECHNICAL_USER_TOKEN:dummy}
 webhookEventService:
   waitingTime: 100
-websocket:
-  heartbeat-interval: 30000