Missing write requests (#75)

* Fix missing callbacks in yamux

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

Author: kamilsa

include/libp2p/muxer/mplex/mplex_stream.hpp

@@ -6,6 +6,9 @@
 #ifndef LIBP2P_MPLEX_STREAM_HPP
 #define LIBP2P_MPLEX_STREAM_HPP
 
+#include <deque>
+#include <mutex>
+
 #include <boost/asio/streambuf.hpp>
 #include <boost/noncopyable.hpp>
 #include <libp2p/common/logger.hpp>
@@ -111,6 +114,12 @@ namespace libp2p::connection {
     std::function<void(outcome::result<size_t>)> data_notifyee_;
     bool data_notified_ = false;
 
+    /// Queue of write requests that were received when stream was writing
+    std::deque<std::tuple<gsl::span<const uint8_t>, size_t, WriteCallbackFunc>>
+        write_queue_{};
+
+    mutable std::mutex write_queue_mutex_;
+
     /// is the stream opened for reads?
     bool is_readable_ = true;
     bool is_reading_ = false;

include/libp2p/muxer/yamux/yamux_stream.hpp

@@ -6,6 +6,9 @@
 #ifndef LIBP2P_YAMUX_STREAM_HPP
 #define LIBP2P_YAMUX_STREAM_HPP
 
+#include <deque>
+#include <mutex>
+
 #include <boost/asio/streambuf.hpp>
 #include <boost/noncopyable.hpp>
 #include <libp2p/connection/stream.hpp>
@@ -159,6 +162,13 @@ namespace libp2p::connection {
     /// write callback, non-zero during async sends
     WriteCallbackFunc write_cb_;
 
+    /// Queue of write requests that were received when stream was writing
+    std::deque<
+        std::tuple<gsl::span<const uint8_t>, size_t, WriteCallbackFunc, bool>>
+        write_queue_{};
+
+    mutable std::mutex write_queue_mutex_;
+
     /// starts async write operation
     void beginWrite(WriteCallbackFunc cb);
 

src/muxer/mplex/mplex_stream.cpp

@@ -136,7 +136,9 @@ namespace libp2p::connection {
       return cb(Error::INVALID_ARGUMENT);
     }
     if (is_writing_) {
-      return cb(Error::IS_WRITING);
+      std::lock_guard<std::mutex> lock(write_queue_mutex_);
+      write_queue_.emplace_back(in, bytes, cb);
+      return;
     }
     if (connection_.expired()) {
       return cb(Error::CONNECTION_IS_DEAD);
@@ -153,6 +155,15 @@ namespace libp2p::connection {
                               write_res.error().message());
           }
           cb(std::forward<decltype(write_res)>(write_res));
+
+          std::lock_guard<std::mutex> lock(self->write_queue_mutex_);
+          // check if new write messages were received while stream was writing
+          // and propagate these messages
+          if (not self->write_queue_.empty()) {
+            auto [in, bytes, cb] = self->write_queue_.front();
+            self->write_queue_.pop_front();
+            self->write(in, bytes, cb);
+          }
         });
   }
 

src/muxer/yamux/yamux_stream.cpp

@@ -179,6 +179,15 @@ namespace libp2p::connection {
       write_cb_ = WriteCallbackFunc{};
       cb(result);
     }
+
+    std::lock_guard<std::mutex> lock(write_queue_mutex_);
+    // check if new write messages were received while stream was writing
+    // and propagate these messages
+    if (not write_queue_.empty()) {
+      auto [in, bytes, cb, some] = write_queue_.front();
+      write_queue_.pop_front();
+      write(in, bytes, cb, some);
+    }
   }
 
   void YamuxStream::write(gsl::span<const uint8_t> in, size_t bytes,
@@ -187,7 +196,9 @@ namespace libp2p::connection {
       return cb(Error::NOT_WRITABLE);
     }
     if (is_writing_) {
-      return cb(Error::IS_WRITING);
+      std::lock_guard<std::mutex> lock(write_queue_mutex_);
+      write_queue_.emplace_back(in, bytes, cb, some);
+      return;
     }
 
     beginWrite(std::move(cb));