Prevent exceptions due to heartbeats

Resolves: #173
joffrey-bion · Jan 16, 2022 · 604c28a · 604c28a
1 parent fe099b9
commit 604c28a
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/krossbow-stomp-core/src/commonMain/kotlin/org/hildan/krossbow/stomp/BaseStompSession.kt b/krossbow-stomp-core/src/commonMain/kotlin/org/hildan/krossbow/stomp/BaseStompSession.kt
@@ -13,6 +13,7 @@ import org.hildan.krossbow.stomp.headers.*
 import org.hildan.krossbow.stomp.heartbeats.HeartBeater
 import org.hildan.krossbow.stomp.heartbeats.NO_HEART_BEATS
 import org.hildan.krossbow.utils.generateUuid
+import org.hildan.krossbow.websocket.WebSocketException
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.coroutines.cancellation.CancellationException
@@ -35,7 +36,15 @@ internal class BaseStompSession(
         HeartBeater(
             heartBeat = heartBeat,
             tolerance = config.heartBeatTolerance,
-            sendHeartBeat = { stompSocket.sendHeartBeat() },
+            sendHeartBeat = {
+                // The web socket could have errored or be closed, and the heart beater's job not yet be cancelled.
+                // In this case, we don't want the heart beater to crash
+                try {
+                    stompSocket.sendHeartBeat()
+                } catch(e : WebSocketException) {
+                    shutdown("STOMP session failed: heart beat couldn't be sent", cause = e)
+                }
+            },
             onMissingHeartBeat = {
                 val cause = MissingHeartBeatException(heartBeat.expectedPeriod)
                 sharedStompEvents.emit(StompEvent.Error(cause))