From cb432be3272a5663a2e18f2119b1df61ee7ca96f Mon Sep 17 00:00:00 2001
From: Michael Jansen <mjansen4857@gmail.com>
Date: Fri, 22 Dec 2023 14:07:42 -0800
Subject: [PATCH] Fix crash on reconnect & ignore rtt topic ID (#8)

* Ensure proper cleanup when reconnecting

* wait for rtt to connect before continuing

* ignore rtt message topic ID

* update version and changelog
---
 CHANGELOG.md            |  5 +++++
 lib/src/nt4_client.dart | 40 ++++++++++++++++++++++++++--------------
 pubspec.yaml            |  2 +-
 3 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ae204c1..71bc678 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,8 @@
+## 1.3.2
+
+- Fix crash on reconnecting to server
+- Ignore rtt protocol topic ID
+
 ## 1.3.1
 
 - Fix connections being timed out too quickly
diff --git a/lib/src/nt4_client.dart b/lib/src/nt4_client.dart
index e64a58c..c06a105 100644
--- a/lib/src/nt4_client.dart
+++ b/lib/src/nt4_client.dart
@@ -45,7 +45,9 @@ class NT4Client {
   int _timeoutInterval = _pingTimeoutMsV40;
 
   WebSocketChannel? _mainWebsocket;
+  StreamSubscription<dynamic>? _mainWebsocketSub;
   WebSocketChannel? _rttWebsocket;
+  StreamSubscription<dynamic>? _rttWebsocketSub;
 
   /// Create an NT4 client. This will connect to an NT4 server running at
   /// [serverBaseAddress]. This should be either 'localhost' if running a
@@ -450,14 +452,14 @@ class NT4Client {
       _useRTT = true;
       _pingInterval = _pingIntervalMsV41;
       _timeoutInterval = _pingTimeoutMsV41;
-      _rttConnect();
+      await _rttConnect();
     } else {
       _useRTT = false;
       _pingInterval = _pingIntervalMsV40;
       _timeoutInterval = _pingTimeoutMsV40;
     }
 
-    _mainWebsocket!.stream.listen(
+    _mainWebsocketSub = _mainWebsocket!.stream.listen(
       (data) {
         if (!_serverConnectionActive) {
           _lastAnnouncedValues.clear();
@@ -470,10 +472,11 @@ class NT4Client {
       },
       onDone: _wsOnClose,
       onError: (err) {},
+      cancelOnError: true,
     );
 
     NT4Topic timeTopic = NT4Topic(
-        name: "Time",
+        name: 'Time',
         type: NT4TypeStr.typeInt,
         id: -1,
         pubUID: -1,
@@ -500,7 +503,7 @@ class NT4Client {
     }
   }
 
-  void _rttConnect() async {
+  Future<void> _rttConnect() async {
     if (!_useRTT || _rttConnectionActive) {
       return;
     }
@@ -523,7 +526,7 @@ class NT4Client {
 
     _rttConnectionActive = true;
 
-    _rttWebsocket!.stream.listen(
+    _rttWebsocketSub = _rttWebsocket!.stream.listen(
       (data) {
         if (data is! List<int>) {
           return;
@@ -531,7 +534,7 @@ class NT4Client {
 
         var msg = Unpacker.fromList(data).unpackList();
 
-        int topicID = msg[0] as int;
+        // rtt socket will only send timestamps, we can ignore the topic ID
         int timestampUS = msg[1] as int;
         var value = msg[3];
 
@@ -539,16 +542,18 @@ class NT4Client {
           return;
         }
 
-        if (topicID == -1) {
-          _wsHandleRecieveTimestamp(timestampUS, value);
-        }
+        _wsHandleRecieveTimestamp(timestampUS, value);
       },
       onDone: _rttOnClose,
+      onError: (error) {},
+      cancelOnError: true,
     );
   }
 
-  void _rttOnClose() {
-    _rttWebsocket?.sink.close();
+  void _rttOnClose() async {
+    await _rttWebsocketSub?.cancel();
+    _rttWebsocketSub = null;
+    await _rttWebsocket?.sink.close();
     _rttWebsocket = null;
 
     _lastReceivedTime = 0;
@@ -557,9 +562,16 @@ class NT4Client {
     _useRTT = false;
   }
 
-  void _wsOnClose() {
-    _mainWebsocket?.sink.close();
-    _rttWebsocket?.sink.close();
+  void _wsOnClose() async {
+    _pingTimer?.cancel();
+    _pongTimer?.cancel();
+
+    await _mainWebsocketSub?.cancel();
+    _mainWebsocketSub = null;
+    await _mainWebsocket?.sink.close();
+    await _rttWebsocketSub?.cancel();
+    _rttWebsocketSub = null;
+    await _rttWebsocket?.sink.close();
 
     _mainWebsocket = null;
     _rttWebsocket = null;
diff --git a/pubspec.yaml b/pubspec.yaml
index 823329f..eba9146 100644
--- a/pubspec.yaml
+++ b/pubspec.yaml
@@ -1,6 +1,6 @@
 name: nt4
 description: Dart implementation of the WPILib NT4 protocol
-version: 1.3.1
+version: 1.3.2
 repository: https://github.com/mjansen4857/nt4
 
 environment: