[ntcore] WebSocketConnection: Use weak capture

PeterJohnson · Oct 25, 2023 · 6ed3c37 · 6ed3c37
1 parent 8a8e220
commit 6ed3c37
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 9 deletions.
diff --git a/ntcore/src/main/native/cpp/NetworkClient.cpp b/ntcore/src/main/native/cpp/NetworkClient.cpp
@@ -401,6 +401,7 @@ void NetworkClient::WsConnected(wpi::WebSocket& ws, uv::Tcp& tcp,
 
   m_wire =
       std::make_shared<net::WebSocketConnection>(ws, connInfo.protocol_version);
+  m_wire->Start();
   m_clientImpl = std::make_unique<net::ClientImpl>(
       m_loop.Now().count(), m_inst, *m_wire, m_logger, m_timeSyncUpdated,
       [this](uint32_t repeatMs) {

diff --git a/ntcore/src/main/native/cpp/NetworkServer.cpp b/ntcore/src/main/native/cpp/NetworkServer.cpp
@@ -280,6 +280,7 @@ void NetworkServer::ServerConnection4::ProcessWsUpgrade() {
     INFO("CONNECTED NT4 client '{}' (from {})", dedupName, m_connInfo);
     m_info.remote_id = dedupName;
     m_server.AddConnection(this, m_info);
+    m_wire->Start();
     m_websocket->closed.connect([this](uint16_t, std::string_view reason) {
       auto realReason = m_wire->GetDisconnectReason();
       INFO("DISCONNECTED NT4 client '{}' (from {}): {}", m_info.remote_id,

diff --git a/ntcore/src/main/native/cpp/net/WebSocketConnection.cpp b/ntcore/src/main/native/cpp/net/WebSocketConnection.cpp
@@ -84,15 +84,7 @@ void WebSocketConnection::Stream::write_impl(const char* data, size_t len) {
 
 WebSocketConnection::WebSocketConnection(wpi::WebSocket& ws,
                                          unsigned int version)
-    : m_ws{ws}, m_version{version} {
-  m_ws.pong.connect([this](auto data) {
-    if (data.size() != 8) {
-      return;
-    }
-    m_lastPingResponse =
-        wpi::support::endian::read64<wpi::support::native>(data.data());
-  });
-}
+    : m_ws{ws}, m_version{version} {}
 
 WebSocketConnection::~WebSocketConnection() {
   for (auto&& buf : m_bufs) {
@@ -103,6 +95,18 @@ WebSocketConnection::~WebSocketConnection() {
   }
 }
 
+void WebSocketConnection::Start() {
+  m_ws.pong.connect([selfweak = weak_from_this()](auto data) {
+    if (data.size() != 8) {
+      return;
+    }
+    if (auto self = selfweak.lock()) {
+      self->m_lastPingResponse =
+          wpi::support::endian::read64<wpi::support::native>(data.data());
+    }
+  });
+}
+
 void WebSocketConnection::SendPing(uint64_t time) {
   auto buf = AllocBuf();
   buf.len = 8;

diff --git a/ntcore/src/main/native/cpp/net/WebSocketConnection.h b/ntcore/src/main/native/cpp/net/WebSocketConnection.h
@@ -26,6 +26,8 @@ class WebSocketConnection final
   WebSocketConnection(const WebSocketConnection&) = delete;
   WebSocketConnection& operator=(const WebSocketConnection&) = delete;
 
+  void Start();
+
   unsigned int GetVersion() const final { return m_version; }
 
   void SendPing(uint64_t time) final;