Update main.cpp

main.cpp

@@ -1,11 +1,12 @@
 /*
- * CollabDocs C++ — Main Server (UPDATED)
+ * CollabDocs C++ — Main Server (v2)
  *
- * Stack:
- * Boost.Asio  — async I/O
- * Boost.Beast — HTTP + WebSocket
- * nlohmann/json — JSON
- * C++17
+ * Improvements over v1:
+ *   1. Strand-driven write queue (no spin-lock atomic)
+ *   2. Op broadcast batching (8 ms / 20 ops window) via Document::enqueue_broadcast_locked
+ *   3. File-descriptor limit raised at startup
+ *   4. "batch" message type forwarded to clients
+ *   5. user_left carries the name correctly (looked up before erase)
  */
 
 #include "ot_engine.hpp"
@@ -18,6 +19,8 @@
 #include <boost/beast/websocket.hpp>
 #include <nlohmann/json.hpp>
 
+#include <sys/resource.h>
+
 #include <iostream>
 #include <fstream>
 #include <sstream>
@@ -27,6 +30,7 @@
 #include <random>
 #include <filesystem>
 #include <functional>
+#include <deque>
 
 namespace asio  = boost::asio;
 namespace beast = boost::beast;
@@ -37,7 +41,7 @@ using json      = nlohmann::json;
 
 // ─── Globals ─────────────────────────────────────────────────────────────────
 
-static std::string g_html_content;     // cached index.html
+static std::string g_html_content;
 static std::atomic<int> g_conn_count{0};
 
 // ─── Helpers ─────────────────────────────────────────────────────────────────
@@ -48,7 +52,7 @@ inline std::string make_op_id() {
     char buf[17];
     std::snprintf(buf, sizeof(buf), "%016llx",
                   static_cast<unsigned long long>(dist(rng)));
-    return std::string(buf, 8); 
+    return std::string(buf, 8);
 }
 
 inline std::string op_type_str(collab::OpType t) {
@@ -56,6 +60,7 @@ inline std::string op_type_str(collab::OpType t) {
 }
 
 namespace collab {
+
 nlohmann::json Document::get_state_json() const {
     return {
         {"doc_id",  doc_id},
@@ -79,7 +84,8 @@ std::vector<nlohmann::json> DocumentStore::list_docs() {
     }
     return out;
 }
-} 
+
+} // namespace collab
 
 // ─── HTTP handler ─────────────────────────────────────────────────────────────
 
@@ -88,12 +94,12 @@ http::response<http::string_body>
 handle_http(http::request<Body, http::basic_fields<Allocator>> const& req)
 {
     auto target = std::string(req.target());
-    auto qpos = target.find('?');
+    auto qpos   = target.find('?');
     std::string path = (qpos != std::string::npos) ? target.substr(0, qpos) : target;
 
     auto make_response = [&](http::status status,
-                             const std::string& content_type,
-                             const std::string& body) {
+                              const std::string& content_type,
+                              const std::string& body) {
         http::response<http::string_body> res{status, req.version()};
         res.set(http::field::server, "CollabDocs/1.0");
         res.set(http::field::content_type, content_type);
@@ -104,12 +110,11 @@ handle_http(http::request<Body, http::basic_fields<Allocator>> const& req)
         return res;
     };
 
-    if (req.method() == http::verb::get && (path == "/" || path == "/index.html")) {
+    if (req.method() == http::verb::get && (path == "/" || path == "/index.html"))
         return make_response(http::status::ok, "text/html; charset=utf-8", g_html_content);
-    }
 
     if (req.method() == http::verb::get && path == "/health") {
-        json j = {{"status", "ok"}, {"docs", collab::g_store.size()}};
+        json j = {{"status", "ok"}, {"docs", collab::g_store.size()}, {"conns", g_conn_count.load()}};
         return make_response(http::status::ok, "application/json", j.dump());
     }
 
@@ -146,18 +151,16 @@ class WsSession : public std::enable_shared_from_this<WsSession> {
 public:
     explicit WsSession(tcp::socket socket, asio::io_context& ioc)
         : ws_(std::move(socket))
-        , strand_(asio::make_strand(ioc))
         , write_strand_(asio::make_strand(ioc))
     {}
 
-    // FIX 1: Accepts the request object to complete the handshake correctly
-    void run(http::request<http::string_body> req, 
-             std::string doc_id, 
-             std::string user_id, 
-             asio::yield_context yield) {
-        
+    void run(http::request<http::string_body> req,
+             std::string doc_id,
+             std::string user_id,
+             asio::yield_context yield)
+    {
         beast::error_code ec;
-        doc_id_ = doc_id;
+        doc_id_  = doc_id;
         user_id_ = user_id;
 
         ws_.set_option(ws::stream_base::timeout::suggested(beast::role_type::server));
@@ -165,44 +168,57 @@ public:
             res.set(http::field::server, "CollabDocs/1.0");
         }));
 
-        // FIX 2: Pass the already-read request into async_accept
         ws_.async_accept(req, yield[ec]);
         if (ec) return;
 
         auto self = shared_from_this();
+
+        // Register with presence; give it a send callback
         collab::g_presence.join(doc_id, user_id, [self](const std::string& msg) {
             self->enqueue_send(msg);
         });
 
         auto doc = collab::g_store.get_or_create(doc_id);
 
-        // Send initialization state
+        // Wire up the document's broadcast function once (idempotent)
+        {
+            std::lock_guard<std::mutex> lk(doc->mu);
+            if (!doc->broadcast_fn) {
+                doc->broadcast_fn = [](const std::string& did,
+                                       const std::string& json_str,
+                                       const std::string& exclude) {
+                    collab::g_presence.broadcast(did, json_str, exclude);
+                };
+            }
+        }
+
+        // Send init state
         {
             auto user_opt = collab::g_presence.get_user(doc_id, user_id);
             json init;
             {
                 std::lock_guard<std::mutex> lk(doc->mu);
                 init = {
-                    {"type", "init"},
-                    {"user_id", user_id},
-                    {"name", user_opt ? user_opt->name : "Unknown"},
-                    {"color", user_opt ? user_opt->color : "#999"},
+                    {"type",      "init"},
+                    {"user_id",   user_id},
+                    {"name",      user_opt ? user_opt->name  : "Unknown"},
+                    {"color",     user_opt ? user_opt->color : "#999"},
                     {"doc_state", doc->get_state_json()},
-                    {"users", collab::g_presence.get_users_json(doc_id)},
+                    {"users",     collab::g_presence.get_users_json(doc_id)},
                 };
             }
             enqueue_send(init.dump());
         }
 
-        // Broadast join to other users in the same doc
+        // Broadcast join to peers
         {
             auto user_opt = collab::g_presence.get_user(doc_id, user_id);
             json joined = {
-                {"type", "user_joined"},
+                {"type",    "user_joined"},
                 {"user_id", user_id},
-                {"name", user_opt ? user_opt->name : "Unknown"},
-                {"color", user_opt ? user_opt->color : "#999"},
-                {"users", collab::g_presence.get_users_json(doc_id)},
+                {"name",    user_opt ? user_opt->name  : "Unknown"},
+                {"color",   user_opt ? user_opt->color : "#999"},
+                {"users",   collab::g_presence.get_users_json(doc_id)},
             };
             collab::g_presence.broadcast(doc_id, joined.dump(), user_id);
         }
@@ -225,106 +241,130 @@ public:
         }
 
         g_conn_count--;
+
+        // Look up name before erasing from presence
+        std::string leaving_name = "Someone";
+        std::string leaving_color = "#999";
+        {
+            auto u = collab::g_presence.get_user(doc_id, user_id);
+            if (u) { leaving_name = u->name; leaving_color = u->color; }
+        }
+
         collab::g_presence.leave(doc_id, user_id);
 
         json left = {
-            {"type", "user_left"},
+            {"type",    "user_left"},
             {"user_id", user_id},
-            {"name", "Someone"},
-            {"users", collab::g_presence.get_users_json(doc_id)},
+            {"name",    leaving_name},
+            {"color",   leaving_color},
+            {"users",   collab::g_presence.get_users_json(doc_id)},
         };
         collab::g_presence.broadcast(doc_id, left.dump());
     }
 
 private:
-    ws::stream<tcp::socket> ws_;
-    asio::strand<asio::io_context::executor_type> strand_;
-    asio::strand<asio::io_context::executor_type> write_strand_;
+    ws::stream<tcp::socket>                              ws_;
+    asio::strand<asio::io_context::executor_type>        write_strand_;
+    std::deque<std::string>                              write_queue_;
 
     std::string doc_id_;
     std::string user_id_;
 
-    std::mutex wq_mu_;
-    std::vector<std::string> write_queue_;
-    std::atomic<bool> writing_{false};
+    // ── Strand-serialised write queue ────────────────────────────────────────
+    // All access to write_queue_ happens only on write_strand_, so no mutex needed.
 
     void enqueue_send(const std::string& msg) {
-        {
-            std::lock_guard<std::mutex> lk(wq_mu_);
-            write_queue_.push_back(msg);
-        }
-        do_write();
+        asio::post(write_strand_, [self = shared_from_this(), msg]() {
+            bool was_empty = self->write_queue_.empty();
+            self->write_queue_.push_back(msg);
+            if (was_empty) self->do_flush();
+        });
     }
 
-    void do_write() {
-        bool expected = false;
-        if (!writing_.compare_exchange_strong(expected, true)) return;
+    void do_flush() {
+        // Must be called on write_strand_
+        if (write_queue_.empty()) return;
 
-        std::string msg;
-        {
-            std::lock_guard<std::mutex> lk(wq_mu_);
-            if (write_queue_.empty()) { writing_ = false; return; }
-            msg = std::move(write_queue_.front());
-            write_queue_.erase(write_queue_.begin());
-        }
+        const std::string& front = write_queue_.front();
+        beast::error_code ec;
+        ws_.text(true);
+        ws_.write(asio::buffer(front), ec);
+        write_queue_.pop_front();
 
-        auto self = shared_from_this();
-        asio::post(write_strand_, [self, m = std::move(msg)]() mutable {
-            beast::error_code ec;
-            self->ws_.text(true);
-            self->ws_.write(asio::buffer(m), ec);
-            self->writing_ = false;
-            if (!ec) self->do_write();
-        });
+        if (ec) {
+            write_queue_.clear();
+            return;
+        }
+        if (!write_queue_.empty()) do_flush();
     }
 
-    void handle_message(const std::string& raw, std::shared_ptr<collab::Document> doc) {
+    // ── Message dispatch ──────────────────────────────────────────────────────
+
+    void handle_message(const std::string& raw,
+                        std::shared_ptr<collab::Document> doc)
+    {
         json msg;
         try { msg = json::parse(raw); } catch (...) { return; }
 
         std::string type = msg.value("type", "");
-        if (type == "operation") {
-            handle_operation(msg, doc);
-        } else if (type == "cursor") {
-            handle_cursor(msg);
-        } else if (type == "title_change") {
-            handle_title_change(msg);
-        } else if (type == "ping") {
-            auto u = collab::g_presence.get_user(doc_id_, user_id_);
-            (void)u;
-        }
+        if      (type == "operation")    handle_operation(msg, doc);
+        else if (type == "cursor")       handle_cursor(msg);
+        else if (type == "title_change") handle_title_change(msg);
+        else if (type == "ping")         { /* keep-alive, no-op */ }
     }
 
-    void handle_operation(const json& msg, std::shared_ptr<collab::Document> doc) {
+    // ── Operation handler (with broadcast batching) ───────────────────────────
+
+    void handle_operation(const json& msg,
+                          std::shared_ptr<collab::Document> doc)
+    {
         std::string op_type_s = msg.value("op_type", "");
         if (op_type_s != "insert" && op_type_s != "delete") return;
 
         collab::Operation op;
         op.type = (op_type_s == "insert") ? collab::OpType::INSERT : collab::OpType::DELETE;
-        op.position = msg.value("position", 0);
-        op.value = msg.value("value", "");
-        op.length = (op.type == collab::OpType::INSERT) ? static_cast<int>(op.value.size()) : msg.value("length", 0);
+        op.position     = msg.value("position", 0);
+        op.value        = msg.value("value", "");
+        op.length       = (op.type == collab::OpType::INSERT)
+                              ? static_cast<int>(op.value.size())
+                              : msg.value("length", 0);
         op.base_version = msg.value("base_version", 0);
-        op.user_id = user_id_;
-        op.op_id = msg.value("op_id", make_op_id());
+        op.user_id      = user_id_;
+        op.op_id        = msg.value("op_id", make_op_id());
 
         collab::Operation transformed;
         int new_version;
+
         {
             std::lock_guard<std::mutex> lk(doc->mu);
             transformed = doc->apply_op_locked(op);
             new_version = doc->version;
+
+            // Build the broadcast payload
+            json bcast = {
+                {"type",           "operation"},
+                {"op_type",        op_type_str(transformed.type)},
+                {"position",       transformed.position},
+                {"value",          transformed.value},
+                {"length",         transformed.length},
+                {"server_version", new_version},
+                {"user_id",        user_id_},
+                {"op_id",          transformed.op_id},
+            };
+
+            // Enqueue for batched broadcast (flushes automatically when thresholds hit)
+            doc->enqueue_broadcast_locked(std::move(bcast), user_id_);
         }
 
-        enqueue_send(json({{"type", "ack"}, {"op_id", op.op_id}, {"server_version", new_version}}).dump());
+        // Fire any deferred broadcasts (called AFTER lock release)
+        doc->fire_pending_broadcasts();
 
-        json bcast = {
-            {"type", "operation"}, {"op_type", op_type_str(transformed.type)},
-            {"position", transformed.position}, {"value", transformed.value},
-            {"length", transformed.length}, {"server_version", new_version},
-            {"user_id", user_id_}, {"op_id", transformed.op_id},
-        };
-        collab::g_presence.broadcast(doc_id_, bcast.dump(), user_id_);
+        // ACK directly back to the sender (not batched — sender needs this immediately)
+        enqueue_send(json({
+            {"type",           "ack"},
+            {"op_id",          op.op_id},
+            {"server_version", new_version},
+        }).dump());
     }
 
     void handle_cursor(const json& msg) {
@@ -332,8 +372,11 @@ private:
         collab::g_presence.update_cursor(doc_id_, user_id_, pos);
         auto u = collab::g_presence.get_user(doc_id_, user_id_);
         json bcast = {
-            {"type", "cursor"}, {"user_id", user_id_}, {"name", u ? u->name : ""},
-            {"color", u ? u->color : "#999"}, {"cursor_pos", pos}
+            {"type",       "cursor"},
+            {"user_id",    user_id_},
+            {"name",       u ? u->name  : ""},
+            {"color",      u ? u->color : "#999"},
+            {"cursor_pos", pos},
         };
         collab::g_presence.broadcast(doc_id_, bcast.dump(), user_id_);
     }
@@ -341,7 +384,10 @@ private:
     void handle_title_change(const json& msg) {
         std::string title = msg.value("title", "Untitled Document");
         collab::g_store.update_title(doc_id_, title);
-        collab::g_presence.broadcast(doc_id_, json({{"type", "title_change"}, {"title", title}, {"user_id", user_id_}}).dump(), user_id_);
+        collab::g_presence.broadcast(
+            doc_id_,
+            json({{"type","title_change"},{"title",title},{"user_id",user_id_}}).dump(),
+            user_id_);
     }
 };
 
@@ -368,55 +414,94 @@ void do_listen(asio::io_context& ioc, tcp::endpoint ep) {
             if (ws::is_upgrade(req)) {
                 std::string target = std::string(req.target());
                 auto qpos = target.find('?');
-                std::string doc_path = (qpos != std::string::npos) ? target.substr(0, qpos) : target;
-                
-                // Extract doc_id and user_id
-                std::string doc_id = (doc_path.size() > 4) ? doc_path.substr(4) : "welcome";
-                
-                // Simple user_id extraction
+                std::string doc_path = (qpos != std::string::npos)
+                                         ? target.substr(0, qpos)
+                                         : target;
+
+                // /ws/<doc_id>
+                std::string doc_id = (doc_path.size() > 4)
+                                       ? doc_path.substr(4)
+                                       : "welcome";
+
+                // Extract user_id from query string
                 std::string user_id = make_op_id();
                 if (qpos != std::string::npos) {
                     std::string query = target.substr(qpos + 1);
                     auto upos = query.find("user_id=");
-                    if (upos != std::string::npos) user_id = query.substr(upos + 8);
+                    if (upos != std::string::npos)
+                        user_id = query.substr(upos + 8);
+                    // Trim any trailing '&' params
+                    auto amp = user_id.find('&');
+                    if (amp != std::string::npos) user_id = user_id.substr(0, amp);
                 }
 
                 auto session = std::make_shared<WsSession>(stream.release_socket(), ioc);
-                asio::spawn(ioc, [session, req = std::move(req), doc_id, user_id](asio::yield_context y) mutable {
+                asio::spawn(ioc, [session, req = std::move(req), doc_id, user_id]
+                                 (asio::yield_context y) mutable {
                     session->run(std::move(req), doc_id, user_id, y);
                 });
             } else {
                 auto res = handle_http(req);
                 http::async_write(stream, res, yield[ec]);
-                if (!req.keep_alive()) stream.socket().shutdown(tcp::socket::shutdown_send, ec);
+                if (!req.keep_alive())
+                    stream.socket().shutdown(tcp::socket::shutdown_send, ec);
             }
         });
     }
 }
 
+// ─── main ─────────────────────────────────────────────────────────────────────
+
 int main(int argc, char* argv[]) {
+    // ── 1. Raise file-descriptor limit so 500+ WebSocket connections work ────
+    {
+        struct rlimit rl{};
+        getrlimit(RLIMIT_NOFILE, &rl);
+        rl.rlim_cur = std::min(static_cast<rlim_t>(65535), rl.rlim_max);
+        if (setrlimit(RLIMIT_NOFILE, &rl) == 0)
+            std::cout << "[CollabDocs] fd limit raised to " << rl.rlim_cur << "\n";
+        else
+            std::cerr << "[CollabDocs] Warning: could not raise fd limit\n";
+    }
+
+    // ── 2. Port ───────────────────────────────────────────────────────────────
     uint16_t port = 7860;
     if (auto* p = std::getenv("PORT")) port = static_cast<uint16_t>(std::stoi(p));
-    
-    std::ifstream f("index.html");
-    if (f) { std::ostringstream ss; ss << f.rdbuf(); g_html_content = ss.str(); }
-
-    auto doc = collab::g_store.get_or_create("welcome");
-    { 
-        std::lock_guard<std::mutex> lk(doc->mu); 
-        if (doc->content.empty()) { doc->title = "Welcome"; doc->content = "Start typing..."; doc->version = 1; }
+
+    // ── 3. Load frontend ──────────────────────────────────────────────────────
+    {
+        std::ifstream f("index.html");
+        if (f) { std::ostringstream ss; ss << f.rdbuf(); g_html_content = ss.str(); }
+        else   { std::cerr << "[CollabDocs] Warning: index.html not found\n"; }
+    }
+
+    // ── 4. Seed welcome document ──────────────────────────────────────────────
+    {
+        auto doc = collab::g_store.get_or_create("welcome");
+        std::lock_guard<std::mutex> lk(doc->mu);
+        if (doc->content.empty()) {
+            doc->title   = "Welcome";
+            doc->content = "Start typing…";
+            doc->version = 1;
+        }
     }
 
+    // ── 5. IO threads ─────────────────────────────────────────────────────────
     unsigned threads = std::max(1u, std::thread::hardware_concurrency());
     asio::io_context ioc(static_cast<int>(threads));
     auto ep = tcp::endpoint(asio::ip::make_address("0.0.0.0"), port);
-    
-    std::cout << "[CollabDocs] Listening on 0.0.0.0:" << port << " (" << threads << " threads)" << std::endl;
+
+    std::cout << "[CollabDocs] Listening on 0.0.0.0:" << port
+              << "  threads=" << threads << "\n";
+
     auto work_guard = asio::make_work_guard(ioc);
-    std::thread listener([&]{ do_listen(ioc, ep); });
+    std::thread listener([&] { do_listen(ioc, ep); });
 
     std::vector<std::thread> pool;
-    for (unsigned i = 0; i < threads - 1; ++i) pool.emplace_back([&]{ ioc.run(); });
+    pool.reserve(threads - 1);
+    for (unsigned i = 0; i < threads - 1; ++i)
+        pool.emplace_back([&] { ioc.run(); });
+
     ioc.run();
     return 0;
 }