fix yamux callback

Signed-off-by: turuslan <[email protected]>

Author: turuslan

include/libp2p/muxer/yamux/yamux_stream.hpp

@@ -147,9 +147,6 @@ namespace libp2p::connection {
     /// True if the stream is writable, until FIN sent
     bool is_writable_ = true;
 
-    /// If set to true, then no more callbacks to client
-    bool no_more_callbacks_ = false;
-
     /// True after FIN sent
     bool fin_sent_ = false;
 

src/muxer/yamux/yamux_stream.cpp

@@ -55,16 +55,7 @@ namespace libp2p::connection {
 
   void YamuxStream::deferReadCallback(outcome::result<size_t> res,
                                       ReadCallbackFunc cb) {
-    if (no_more_callbacks_) {
-      log()->debug("{} closed by client, ignoring callback", stream_id_);
-      return;
-    }
-    feedback_.deferCall([wptr = weak_from_this(), res, cb = std::move(cb)]() {
-      auto self = wptr.lock();
-      if (self && !self->no_more_callbacks_) {
-        cb(res);
-      }
-    });
+    feedback_.deferCall([res, cb{std::move(cb)}] { cb(res); });
   }
 
   void YamuxStream::writeSome(BytesIn in, size_t bytes, WriteCallbackFunc cb) {
@@ -73,16 +64,7 @@ namespace libp2p::connection {
 
   void YamuxStream::deferWriteCallback(std::error_code ec,
                                        WriteCallbackFunc cb) {
-    if (no_more_callbacks_) {
-      log()->debug("{} closed by client, ignoring callback", stream_id_);
-      return;
-    }
-    feedback_.deferCall([wptr = weak_from_this(), ec, cb = std::move(cb)]() {
-      auto self = wptr.lock();
-      if (self && !self->no_more_callbacks_) {
-        cb(ec);
-      }
-    });
+    feedback_.deferCall([ec, cb{std::move(cb)}] { cb(ec); });
   }
 
   bool YamuxStream::isClosed() const {
@@ -92,12 +74,8 @@ namespace libp2p::connection {
   void YamuxStream::close(VoidResultHandlerFunc cb) {
     if (isClosed()) {
       if (cb) {
-        feedback_.deferCall([wptr{weak_from_this()}, cb{std::move(cb)}] {
-          auto self = wptr.lock();
-          if (self) {
-            cb(*self->close_reason_);
-          }
-        });
+        feedback_.deferCall(
+            [cb{std::move(cb)}, ec{*close_reason_}] { cb(ec); });
       }
       return;
     }
@@ -137,7 +115,6 @@ namespace libp2p::connection {
   }
 
   void YamuxStream::reset() {
-    no_more_callbacks_ = true;
     feedback_.resetStream(stream_id_);
     doClose(Error::STREAM_RESET_BY_HOST, true);
   }
@@ -163,14 +140,7 @@ namespace libp2p::connection {
     }
 
     if (cb) {
-      feedback_.deferCall([wptr = weak_from_this(), cb = std::move(cb), ec]() {
-        auto self = wptr.lock();
-        if (!self) {
-          return;
-        }
-        if (self->no_more_callbacks_) {
-          return;
-        }
+      feedback_.deferCall([cb{std::move(cb)}, ec] {
         if (!ec) {
           cb(outcome::success());
         } else {
@@ -336,7 +306,7 @@ namespace libp2p::connection {
       return;
     }
 
-    if (result.cb && !no_more_callbacks_) {
+    if (result.cb) {
       result.cb(result.size_to_ack);
     }
   }
@@ -373,37 +343,22 @@ namespace libp2p::connection {
     VoidResultHandlerFunc window_size_cb;
     window_size_cb.swap(window_size_cb_);
 
-    if (no_more_callbacks_) {
-      return;
-    }
-
     // now we are detached from *this* and may be killed from inside callbacks
     // we will call
-    auto wptr = weak_from_this();
+    auto self = shared_from_this();
 
     if (read_cb_and_res.first) {
       read_cb_and_res.first(read_cb_and_res.second);
     }
 
-    if (wptr.expired() || no_more_callbacks_) {
-      return;
-    }
-
     for (const auto &cb : write_callbacks) {
       cb(ec);
-      if (wptr.expired() || no_more_callbacks_) {
-        return;
-      }
     }
 
     if (window_size_cb) {
       window_size_cb(ec);
     }
 
-    if (wptr.expired() || no_more_callbacks_) {
-      return;
-    }
-
     if (close_cb_and_res.first) {
       close_cb_and_res.first(close_cb_and_res.second);
     }

src/muxer/yamux/yamuxed_connection.cpp

@@ -644,7 +644,7 @@ namespace libp2p::connection {
     }
 
     // this instance may be killed inside further callback
-    auto wptr = weak_from_this();
+    auto self = shared_from_this();
 
     if (stream_id != 0) {
       // pass write ack to stream about data size written except header size
@@ -669,11 +669,6 @@ namespace libp2p::connection {
       }
     }
 
-    if (wptr.expired()) {
-      // *this* no longer exists
-      return;
-    }
-
     is_writing_ = false;
 
     if (not write_queue_.empty()) {

src/transport/tcp/tcp_connection.cpp

@@ -92,17 +92,14 @@ namespace libp2p::transport {
     template <typename Callback>
     auto closeOnError(TcpConnection &conn, Callback cb) {
       return [cb{std::move(cb)}, wptr{conn.weak_from_this()}](auto ec,
-                                                              auto result) {
-        if (auto self = wptr.lock()) {
-          if (ec) {
-            cb(ec);
+                                                              size_t result) {
+        if (ec) {
+          cb(ec);
+          if (auto self = wptr.lock()) {
             self->close(ec);
-            return;
           }
-          TRACE("{} {}", self->str(), result);
-          cb(result);
         } else {
-          log().debug("connection wptr expired");
+          cb(result);
         }
       };
     }
@@ -203,42 +200,14 @@ namespace libp2p::transport {
                              closeOnError(*this, std::move(cb)));
   }
 
-  namespace {
-    template <typename Callback, typename Arg>
-    void deferCallback(boost::asio::io_context &ctx,
-                       std::weak_ptr<TcpConnection> wptr,
-                       bool &closed_by_host,
-                       Callback cb,
-                       Arg arg) {
-      // defers callback to the next event loop cycle,
-      // cb will be called iff TcpConnection is still alive
-      // and was not closed by host's side
-      boost::asio::post(
-          ctx,
-          [wptr = std::move(wptr), cb = std::move(cb), arg, &closed_by_host]() {
-            if (!wptr.expired() && !closed_by_host) {
-              cb(arg);
-            }
-          });
-    }
-  }  // namespace
-
   void TcpConnection::deferReadCallback(outcome::result<size_t> res,
                                         ReadCallbackFunc cb) {
-    deferCallback(context_,
-                  weak_from_this(),
-                  std::ref(closed_by_host_),
-                  std::move(cb),
-                  res);
+    boost::asio::post(context_, [res, cb{std::move(cb)}] { cb(res); });
   }
 
   void TcpConnection::deferWriteCallback(std::error_code ec,
                                          WriteCallbackFunc cb) {
-    deferCallback(context_,
-                  weak_from_this(),
-                  std::ref(closed_by_host_),
-                  std::move(cb),
-                  ec);
+    boost::asio::post(context_, [ec, cb{std::move(cb)}] { cb(ec); });
   }
 
   outcome::result<void> TcpConnection::saveMultiaddresses() {