finish stress test coding

2 files changed, 24 insertions, 6 deletions

diff --git a/include/salticidae/conn.h b/include/salticidae/conn.h
index 3144a68..69402d3 100644
--- a/include/salticidae/conn.h
+++ b/include/salticidae/conn.h
@@ -77,10 +77,11 @@ class ConnPool {
         protected:
         size_t seg_buff_size;
         conn_t self_ref;
+        std::mutex ref_mlock;
         int fd;
         Worker *worker;
         ConnPool *cpool;
-        ConnMode mode;
+        std::atomic<ConnMode> mode;
         NetAddr addr;
 
         MPSCWriteBuffer send_buffer;
@@ -110,7 +111,16 @@ class ConnPool {
         }
 
         /** Get the handle to itself. */
-        conn_t self() { return self_ref; }
+        conn_t self() {
+            mutex_lg_t _(ref_mlock);
+            return self_ref;
+        }
+
+        void release_self() {
+            mutex_lg_t _(ref_mlock);
+            self_ref = nullptr;
+        }
+
         operator std::string() const;
         const NetAddr &get_addr() const { return addr; }
         ConnMode get_mode() const { return mode; }
@@ -380,6 +390,14 @@ class ConnPool {
 
     template<typename Func>
     void reg_conn_handler(Func cb) { conn_cb = cb; }
+
+    void terminate(Conn &_conn) {
+        auto conn = _conn.self();
+        if (!conn) return;
+        disp_tcall->async_call([this, conn](ThreadCall::Handle &) {
+            conn->disp_terminate();
+        });
+    }
 };
 
 }
diff --git a/include/salticidae/network.h b/include/salticidae/network.h
index 78449eb..4e182fc 100644
--- a/include/salticidae/network.h
+++ b/include/salticidae/network.h
@@ -622,11 +622,11 @@ void PeerNetwork<O, _, __>::start_active_conn(const NetAddr &addr) {
 /* begin: functions invoked by the user loop */
 template<typename O, O _, O __>
 void PeerNetwork<O, _, __>::msg_ping(MsgPing &&msg, Conn &_conn) {
-    if (_conn.get_mode() == ConnPool::Conn::DEAD) return;
     auto conn = static_pointer_cast<Conn>(_conn.self());
-    assert(conn);
+    if (!conn) return;
     uint16_t port = msg.port;
     this->disp_tcall->async_call([this, conn, port](ThreadCall::Handle &msg) {
+        if (conn->get_mode() == ConnPool::Conn::DEAD) return;
         SALTICIDAE_LOG_INFO("ping from %s, port %u",
                             std::string(*conn).c_str(), ntohs(port));
         if (check_new_conn(conn, port)) return;
@@ -637,11 +637,11 @@ void PeerNetwork<O, _, __>::msg_ping(MsgPing &&msg, Conn &_conn) {
 
 template<typename O, O _, O __>
 void PeerNetwork<O, _, __>::msg_pong(MsgPong &&msg, Conn &_conn) {
-    if (_conn.get_mode() == ConnPool::Conn::DEAD) return;
     auto conn = static_pointer_cast<Conn>(_conn.self());
-    assert(conn);
+    if (!conn) return;
     uint16_t port = msg.port;
     this->disp_tcall->async_call([this, conn, port](ThreadCall::Handle &msg) {
+        if (conn->get_mode() == ConnPool::Conn::DEAD) return;
         auto it = id2peer.find(conn->peer_id);
         if (it == id2peer.end())
         {