Update changelog for version 0.1.2, enhance WebsocketManager to prevent use-after-free issues, and improve test coverage for unsubscribe behavior

Author: kzhdev

CHANGELOG.md

@@ -7,6 +7,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ---
 
+## [0.1.2] - 2026-06-09
+
+### Fixed
+- **`WebsocketManager`: use-after-free / same-dispatch unsubscribe race** — callbacks now dispatch from a lock-free snapshot of shared handler state instead of copying `std::function`s under a mutex. Each callback entry carries an `active` flag and `in_flight` counter: `unsubscribe()` removes the callback from future snapshots, marks that specific callback inactive, and waits on its per-callback in-flight count with C++20 atomic wait/notify. This prevents a later callback from firing after being removed earlier in the same snapshot while still allowing self-unsubscribe without deadlocking.
+
+### Changed
+- **`WebsocketManager`: removed `pending_subs_` pre-connection queue** — the underlying `slick::net::Websocket` already buffers outbound frames until the connection is established, making the manual pending-subscription queue redundant. `subscribe()` and `unsubscribe()` now call `ws_->send()` unconditionally; `flush_pending()`, `writer_mutex_`, and the `pending_subs_` vector are all removed.
+
+### Tests
+- **`UnsubscribeBlocksUntilCallbackCompletes`** (integration) — regression test for the above fix; subscribes with a callback that blocks until explicitly released, calls `unsubscribe()` concurrently from a second thread, and asserts that `unsubscribe()` does not return before the in-flight callback has finished.
+- **`CallbackCanRemoveLaterCallbackFromSameDispatch`** (integration) — regression test for the snapshot hazard where callback A unsubscribes callback B on the same channel before the dispatcher reaches B in the already-loaded snapshot.
+
 ## [0.1.1] - 2026-06-08
 
 ### Added

CMakeLists.txt

@@ -1,5 +1,5 @@
 cmake_minimum_required(VERSION 3.21)
-project(hyperliquid-cpp VERSION 0.1.1 LANGUAGES CXX)
+project(hyperliquid-cpp VERSION 0.1.2 LANGUAGES CXX)
 
 set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)

include/hyperliquid/websocket_manager.hpp

@@ -3,7 +3,6 @@
 #include <atomic>
 #include <functional>
 #include <memory>
-#include <mutex>
 #include <string>
 #include <string_view>
 #include <thread>
@@ -51,25 +50,32 @@ class WebsocketManager {
     void on_error(std::string err);
 
     void ping_loop();        // runs in ping_thread_
-    void flush_pending();    // send queued subs once connected
 
     std::string ws_url_;
     std::unique_ptr<slick::net::Websocket> ws_;
 
-    std::mutex mutex_;
+    // Handler snapshots are published with copy-on-write. Each snapshot holds
+    // shared per-callback state so unsubscribe() can deactivate exactly one
+    // callback and wait for only that callback's in-flight invocation to drain.
+    struct Handler {
+        int subscription_id;
+        std::function<void(const nlohmann::json&)> callback;
+        std::atomic<bool> active{true};
+        std::atomic<unsigned int> in_flight{0};
+
+        Handler(int id, std::function<void(const nlohmann::json&)> cb)
+            : subscription_id(id), callback(std::move(cb)) {}
+    };
+
+    using HandlerPtr = std::shared_ptr<Handler>;
+    using HandlerMap = std::unordered_map<std::string, std::vector<HandlerPtr>>;
+
+    std::atomic<std::shared_ptr<const HandlerMap>> handlers_;
+
     std::atomic<bool> connected_{false};
     std::atomic<bool> running_{true};
     std::atomic<int>  next_id_{1};
 
-    // identifier → list of {subscription_id, callback}
-    std::unordered_map<
-        std::string,
-        std::vector<std::pair<int, std::function<void(const nlohmann::json&)>>>>
-        handlers_;
-
-    // Full subscription JSONs queued before connection ready
-    std::vector<nlohmann::json> pending_subs_;
-
     std::thread ping_thread_;
 };
 

src/websocket_manager.cpp

@@ -7,9 +7,22 @@
 
 #include <chrono>
 #include <iostream>
-#include <stdexcept>
 #include <thread>
 
+namespace {
+    // Points at the callback currently executing on the WebSocket IO thread.
+    // unsubscribe() uses this to avoid deadlocking when a callback removes itself.
+    thread_local const void* dispatching_handler = nullptr;
+
+    void wait_for_zero(std::atomic<unsigned int>& counter) {
+        unsigned int in_flight = counter.load(std::memory_order_acquire);
+        while (in_flight != 0) {
+            counter.wait(in_flight, std::memory_order_relaxed);
+            in_flight = counter.load(std::memory_order_acquire);
+        }
+    }
+}
+
 namespace hyperliquid {
 
 // ── URL helpers ───────────────────────────────────────────────────────────────
@@ -57,6 +70,7 @@ std::string WebsocketManager::to_identifier(const nlohmann::json& sub) {
 
 WebsocketManager::WebsocketManager(std::string_view http_base_url)
     : ws_url_(http_to_ws_url(http_base_url))
+    , handlers_(std::make_shared<HandlerMap>())
 {
     ws_ = std::make_unique<slick::net::Websocket>(
         ws_url_,
@@ -80,12 +94,11 @@ WebsocketManager::~WebsocketManager() {
 // ── Connection callbacks ──────────────────────────────────────────────────────
 
 void WebsocketManager::on_connected() {
-    connected_ = true;
-    flush_pending();
+    connected_.store(true, std::memory_order_release);
 }
 
 void WebsocketManager::on_disconnected() {
-    connected_ = false;
+    connected_.store(false, std::memory_order_release);
 }
 
 void WebsocketManager::on_message(const char* data, std::size_t len) {
@@ -113,21 +126,42 @@ void WebsocketManager::on_message(const char* data, std::size_t len) {
             identifier = channel + ":" + d["user"].get<std::string>();
     }
 
-    std::vector<std::function<void(const nlohmann::json&)>> cbs;
-    {
-        std::lock_guard lock(mutex_);
-        auto it = handlers_.find(identifier);
-        if (it == handlers_.end()) {
-            // Try without suffix (e.g. allMids has no coin)
-            it = handlers_.find(channel);
+    // Lock-free snapshot: atomically borrow a reference to the current map.
+    // Each callback entry carries its own active/in-flight state so a callback
+    // removed mid-dispatch can still be skipped before invocation.
+    auto h = handlers_.load(std::memory_order_acquire);
+    auto it = h->find(identifier);
+    if (it == h->end())
+        it = h->find(channel);   // fallback for channels without a suffix (allMids)
+    if (it == h->end())
+        return;
+
+    for (const auto& handler : it->second) {
+        if (!handler->active.load(std::memory_order_acquire))
+            continue;
+
+        handler->in_flight.fetch_add(1, std::memory_order_acq_rel);
+        if (!handler->active.load(std::memory_order_acquire)) {
+            if (handler->in_flight.fetch_sub(1, std::memory_order_acq_rel) == 1)
+                handler->in_flight.notify_all();
+            continue;
         }
-        if (it != handlers_.end()) {
-            for (const auto& [id, cb] : it->second)
-                cbs.push_back(cb);
+
+        const void* previous = dispatching_handler;
+        dispatching_handler = handler.get();
+        try {
+            handler->callback(msg);
+        } catch (...) {
+            dispatching_handler = previous;
+            if (handler->in_flight.fetch_sub(1, std::memory_order_acq_rel) == 1)
+                handler->in_flight.notify_all();
+            throw;
         }
-    }
+        dispatching_handler = previous;
 
-    for (auto& cb : cbs) cb(msg);
+        if (handler->in_flight.fetch_sub(1, std::memory_order_acq_rel) == 1)
+            handler->in_flight.notify_all();
+    }
 }
 
 void WebsocketManager::on_error(std::string err) {
@@ -139,29 +173,13 @@ void WebsocketManager::on_error(std::string err) {
 void WebsocketManager::ping_loop() {
     while (running_) {
         std::this_thread::sleep_for(std::chrono::milliseconds(50));
-        if (connected_ && ws_) {
+        if (connected_.load(std::memory_order_acquire) && ws_) {
             static const std::string ping_msg = R"({"method":"ping"})";
             ws_->send(ping_msg.c_str(), ping_msg.size());
         }
     }
 }
 
-// ── Pending subscription flush ────────────────────────────────────────────────
-
-void WebsocketManager::flush_pending() {
-    std::vector<nlohmann::json> subs;
-    {
-        std::lock_guard lock(mutex_);
-        subs = std::move(pending_subs_);
-        pending_subs_.clear();
-    }
-    for (const auto& sub : subs) {
-        nlohmann::json msg{{"method", "subscribe"}, {"subscription", sub}};
-        std::string s = msg.dump();
-        ws_->send(s.c_str(), s.size());
-    }
-}
-
 // ── Subscribe / unsubscribe ───────────────────────────────────────────────────
 
 int WebsocketManager::subscribe(
@@ -171,45 +189,82 @@ int WebsocketManager::subscribe(
     int id = next_id_.fetch_add(1);
     std::string ident = to_identifier(subscription);
 
-    {
-        std::lock_guard lock(mutex_);
-        handlers_[ident].emplace_back(id, std::move(callback));
-
-        if (!connected_) {
-            pending_subs_.push_back(subscription);
-        } else {
-            nlohmann::json msg{{"method", "subscribe"}, {"subscription", subscription}};
-            std::string s = msg.dump();
-            ws_->send(s.c_str(), s.size());
-        }
-    }
+    // CAS loop: build a new map, swap it in; retry if another writer raced us.
+    auto old = handlers_.load(std::memory_order_acquire);
+    auto handler = std::make_shared<Handler>(id, std::move(callback));
+    std::shared_ptr<HandlerMap> new_map;
+    do {
+        new_map = std::make_shared<HandlerMap>(*old);
+        (*new_map)[ident].push_back(handler);
+    } while (!handlers_.compare_exchange_weak(
+        old, new_map,
+        std::memory_order_release,
+        std::memory_order_acquire));
+
+    // The WS layer buffers sends made before the connection is established,
+    // so no need to track pending subscriptions ourselves.
+    nlohmann::json msg{{"method", "subscribe"}, {"subscription", subscription}};
+    std::string s = msg.dump();
+    ws_->send(s.c_str(), s.size());
     return id;
 }
 
 void WebsocketManager::unsubscribe(
     const nlohmann::json& subscription, int subscription_id)
 {
     std::string ident = to_identifier(subscription);
-    {
-        std::lock_guard lock(mutex_);
-        auto it = handlers_.find(ident);
-        if (it != handlers_.end()) {
+    HandlerPtr removed_handler;
+    bool last_handler = false;
+
+    // CAS loop: remove the entry from a fresh copy; retry on contention.
+    std::shared_ptr<const HandlerMap> old = handlers_.load(std::memory_order_acquire);
+    for (;;) {
+        auto new_map = std::make_shared<HandlerMap>(*old);
+        removed_handler.reset();
+        last_handler = false;
+
+        auto it = new_map->find(ident);
+        if (it != new_map->end()) {
             auto& vec = it->second;
             vec.erase(
                 std::remove_if(vec.begin(), vec.end(),
-                    [subscription_id](const auto& p) { return p.first == subscription_id; }),
+                    [&](const HandlerPtr& handler) {
+                        if (handler->subscription_id != subscription_id)
+                            return false;
+                        removed_handler = handler;
+                        return true;
+                    }),
                 vec.end());
             if (vec.empty()) {
-                handlers_.erase(it);
-                // Send unsubscribe only when no more handlers remain
-                if (connected_) {
-                    nlohmann::json msg{{"method", "unsubscribe"}, {"subscription", subscription}};
-                    std::string s = msg.dump();
-                    ws_->send(s.c_str(), s.size());
-                }
+                new_map->erase(it);
+                last_handler = (removed_handler != nullptr);
             }
         }
+        if (handlers_.compare_exchange_weak(
+                old, new_map,
+                std::memory_order_release,
+                std::memory_order_acquire))
+            break;
     }
+
+    if (!removed_handler)
+        return;
+
+    removed_handler->active.store(false, std::memory_order_release);
+
+    if (last_handler) {
+        nlohmann::json msg{{"method", "unsubscribe"}, {"subscription", subscription}};
+        std::string s = msg.dump();
+        ws_->send(s.c_str(), s.size());
+    }
+
+    // Self-unsubscribe from inside the callback would deadlock waiting for our
+    // own in-flight invocation to finish. Cross-callback removal remains safe:
+    // later callbacks in the same snapshot re-check `active` before invoking.
+    if (dispatching_handler == removed_handler.get())
+        return;
+
+    wait_for_zero(removed_handler->in_flight);
 }
 
 } // namespace hyperliquid