2 files changed, 34 insertions, 13 deletions

diff --git a/include/salticidae/conn.h b/include/salticidae/conn.h
index aa414e8..f8156ee 100644
--- a/include/salticidae/conn.h
+++ b/include/salticidae/conn.h
@@ -146,6 +146,9 @@ class ConnPool {
     class Conn;
     /** The handle to a bi-directional connection. */
     using conn_t = RcObj<Conn>;
+    /** The type of callback invoked when connection status is changed. */
+    using conn_callback_t = std::function<void(const conn_t &)>;
+
     /** Abstraction for a bi-directional connection. */
     class Conn {
         friend ConnPool;
@@ -224,17 +227,25 @@ class ConnPool {
         }
 
         /** Called when new data is available. */
-        virtual void on_read() = 0;
+        virtual void on_read() {
+            if (cpool->read_cb) cpool->read_cb(self());
+        }
         /** Called when the underlying connection is established. */
-        virtual void on_setup() = 0;
+        virtual void on_setup() {
+            if (cpool->conn_cb) cpool->conn_cb(self());
+        }
         /** Called when the underlying connection breaks. */
-        virtual void on_teardown() = 0;
+        virtual void on_teardown() {
+            if (cpool->conn_cb) cpool->conn_cb(self());
+        }
     };
     
     private:
     int max_listen_backlog;
     double conn_server_timeout;
     size_t seg_buff_size;
+    conn_callback_t read_cb;
+    conn_callback_t conn_cb;
 
     std::unordered_map<int, conn_t> pool;
     int listen_fd;
@@ -275,6 +286,9 @@ class ConnPool {
      * Does not need to be called if do not want to accept any passive
      * connections. */
     void listen(NetAddr listen_addr);
+
+    template<typename Func>
+    void reg_conn_handler(Func cb) { conn_cb = cb; }
 };
 
 }
diff --git a/include/salticidae/network.h b/include/salticidae/network.h
index 2bc445c..8ede187 100644
--- a/include/salticidae/network.h
+++ b/include/salticidae/network.h
@@ -79,9 +79,6 @@ class MsgNetwork: public ConnPool {
 
         Msg msg;
         MsgState msg_state;
-        MsgNetwork *get_net() {
-            return static_cast<MsgNetwork *>(get_pool());
-        }
 
         protected:
 #ifdef SALTICIDAE_MSG_STAT
@@ -96,6 +93,10 @@ class MsgNetwork: public ConnPool {
 #endif
         {}
 
+        MsgNetwork *get_net() {
+            return static_cast<MsgNetwork *>(get_pool());
+        }
+
 #ifdef SALTICIDAE_MSG_STAT
         size_t get_nsent() const { return nsent; }
         size_t get_nrecv() const { return nrecv; }
@@ -105,8 +106,6 @@ class MsgNetwork: public ConnPool {
 
         protected:
         void on_read() override;
-        void on_setup() override {}
-        void on_teardown() override {}
     };
 
     using conn_t = RcObj<Conn>;
@@ -182,13 +181,14 @@ class ClientNetwork: public MsgNetwork<OpcodeType> {
     public:
     class Conn: public MsgNet::Conn {
         friend ClientNetwork;
-        ClientNetwork *get_net() {
-            return static_cast<ClientNetwork *>(ConnPool::Conn::get_pool());
-        }
 
         public:
         Conn() = default;
 
+        ClientNetwork *get_net() {
+            return static_cast<ClientNetwork *>(ConnPool::Conn::get_pool());
+        }
+
         protected:
         void on_setup() override;
         void on_teardown() override;
@@ -234,12 +234,14 @@ class PeerNetwork: public MsgNetwork<OpcodeType> {
         friend PeerNetwork;
         NetAddr peer_id;
         Event ev_timeout;
+
+        public:
+        Conn() = default;
+
         PeerNetwork *get_net() {
             return static_cast<PeerNetwork *>(ConnPool::Conn::get_pool());
         }
 
-        public:
-        Conn() = default;
         const NetAddr &get_peer() { return peer_id; }
 
         protected:
@@ -366,6 +368,7 @@ class PeerNetwork: public MsgNetwork<OpcodeType> {
 
 template<typename OpcodeType>
 void MsgNetwork<OpcodeType>::Conn::on_read() {
+    ConnPool::Conn::on_read();
     auto &recv_buffer = read();
     auto mn = get_net();
     while (get_fd() != -1)
@@ -431,6 +434,7 @@ void PeerNetwork<O, _, __>::Peer::reset_conn(conn_t new_conn) {
 
 template<typename O, O _, O __>
 void PeerNetwork<O, _, __>::Conn::on_setup() {
+    MsgNet::Conn::on_setup();
     auto pn = get_net();
     assert(!ev_timeout);
     ev_timeout = Event(pn->ec, -1, 0, [this](evutil_socket_t, short) {
@@ -445,6 +449,7 @@ void PeerNetwork<O, _, __>::Conn::on_setup() {
 
 template<typename O, O _, O __>
 void PeerNetwork<O, _, __>::Conn::on_teardown() {
+    MsgNet::Conn::on_teardown();
     auto pn = get_net();
     auto it = pn->id2peer.find(peer_id);
     if (it == pn->id2peer.end()) return;
@@ -645,6 +650,7 @@ const std::vector<NetAddr> &PeerNetwork<O, _, __>::all_peers() const {
 
 template<typename OpcodeType>
 void ClientNetwork<OpcodeType>::Conn::on_setup() {
+    MsgNet::Conn::on_setup();
     assert(this->get_mode() == Conn::PASSIVE);
     const auto &addr = this->get_addr();
     auto cn = get_net();
@@ -656,6 +662,7 @@ void ClientNetwork<OpcodeType>::Conn::on_setup() {
 
 template<typename OpcodeType>
 void ClientNetwork<OpcodeType>::Conn::on_teardown() {
+    MsgNet::Conn::on_teardown();
     assert(this->get_mode() == Conn::PASSIVE);
     get_net()->addr2conn.erase(this->get_addr());
 }