NOVA - Better connection handling

Fixes #305
JetBrains · Jul 27, 2024 · 1c3ed5c · 1c3ed5c
1 parent 1954391
commit 1c3ed5c
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 21 deletions.
diff --git a/.gitignore b/.gitignore
@@ -5,3 +5,4 @@
 build/
 venv/
 tmp/
+out/
diff --git a/src/main/kotlin/com/jetbrains/micropython/nova/WebSocketComm.kt b/src/main/kotlin/com/jetbrains/micropython/nova/WebSocketComm.kt
@@ -170,6 +170,7 @@ with open('$fullName', 'wb') as f:
     inner class MpyWebSocketClient(uri: URI) : WebSocketClient(uri) {
         init {
             isTcpNoDelay = true
+            connectionLostTimeout = 0
         }
 
         override fun onOpen(handshakedata: ServerHandshake) = Unit //Nothing to do
@@ -257,47 +258,47 @@ with open('$fullName', 'wb') as f:
             connected = false
             ttySuspended = true
             client?.close()
-            client = MpyWebSocketClient(uri).apply {
-                connectBlocking()
+            client = MpyWebSocketClient(uri).also { newClient ->
+                newClient.connectBlocking()
                 try {
                     withTimeout(TIMEOUT.toLong()) {
-                        while (true) {
+                        while (newClient.isOpen) {
                             when {
                                 offTtyBuffer.length < PASSWORD_PROMPT.length -> delay(SHORT_DELAY)
                                 offTtyBuffer.length > PASSWORD_PROMPT.length * 2 -> {
                                     offTtyBuffer.setLength(PASSWORD_PROMPT.length * 2)
                                     throw ConnectException("Password exchange error. Received prompt: $offTtyBuffer")
                                 }
-
-                                else -> {
-                                    if (offTtyBuffer.toString().contains(PASSWORD_PROMPT)) {
-                                        offTtyBuffer.clear()
-                                        send("$password\n")
-                                        while (!connected) {
-                                            when {
-                                                offTtyBuffer.contains(LOGIN_SUCCESS) -> connected = true
-                                                offTtyBuffer.contains(LOGIN_FAIL) -> throw ConnectException("Access denied")
-                                                else -> delay(SHORT_DELAY)
-                                            }
-                                        }
-                                    } else {
-                                        throw ConnectException("Password exchange error. Received prompt: $offTtyBuffer")
-                                    }
-                                    return@withTimeout
+                                offTtyBuffer.toString().contains(PASSWORD_PROMPT) -> break
+                                else -> throw ConnectException("Password exchange error. Received prompt: $offTtyBuffer")
                                 }
                             }
+                        offTtyBuffer.clear()
+                        newClient.send("$password\n")
+                        while (!connected && newClient.isOpen) {
+                            when {
+                                offTtyBuffer.contains(LOGIN_SUCCESS) -> connected = true
+                                offTtyBuffer.contains(LOGIN_FAIL) -> throw ConnectException("Access denied")
+                                else -> delay(SHORT_DELAY)
                         }
                     }
                     ttySuspended = false
+                    }
                 } catch (e: TimeoutCancellationException) {
+                    try {
+                        newClient.close()
+                    } catch (_: IOException) {}
                     throw ConnectException("Password exchange timeout. Received prompt: $offTtyBuffer")
                 } finally {
                     offTtyBuffer.clear()
                 }
             }
-
         }
+    }
 
-
+    fun ping() {
+        if(isConnected()) {
+            client?.sendPing()
+        }
     }
 }
diff --git a/src/test/kotlin/com/jetbrains/micropython/nova/RealConnect.kt b/src/test/kotlin/com/jetbrains/micropython/nova/RealConnect.kt
@@ -1,6 +1,7 @@
 package com.jetbrains.micropython.nova
 
 import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.Assertions.*
@@ -84,4 +85,29 @@ class RealConnect {
         }
     }
 
+    @Test
+    @Disabled
+    fun longRunConnection() {
+        runBlocking {
+            comm.connect(URI.create(URL), PASSWORD)
+            assertTrue(comm.isConnected())
+            assertFalse(comm.isTtySuspended())
+            println("Connected")
+            launch {
+                while (comm.isConnected()) {
+                    delay(20000)
+                    println("ping")
+                    comm.ping()
+                }
+            }
+            val start = System.currentTimeMillis()
+            while (comm.isConnected()) {
+                delay(1000)
+                println("Still alive")
+            }
+            println("Disconnect time: ${System.currentTimeMillis() - start}ms")
+
+        }
+    }
+
 }