Add a user-provided-buffers mode, and use it.

This is slightly cleaner when dealing with uvw's buffers.  It seems to
make no appreciable performance difference, though.

Author: jagerman

connection.cpp

@@ -12,6 +12,7 @@
 #include "uvw/poll.h"
 #include "uvw/timer.h"
 
+#include <iterator>
 #include <oxenmq/hex.h>
 #include <oxenmq/bt_serialize.h>
 
@@ -516,19 +517,27 @@ void Connection::flush_streams() {
     while (!strs.empty() && stream_packets < 10) {
         for (auto it = strs.begin(); it != strs.end(); ) {
             auto& stream = **it;
-            auto [first, second] = stream.pending();
+            auto bufs = stream.pending();
             if (stream.is_shutdown ||
-                    (first.empty() && !stream.is_new && !(stream.is_closing && !stream.sent_fin))) {
+                    (bufs.empty() && !stream.is_new && !(stream.is_closing && !stream.sent_fin))) {
                 it = strs.erase(it);
                 continue;
             }
-            std::array<ngtcp2_vec, 2> vecs;
-            vecs[0].base = const_cast<uint8_t*>(u8data(first));
-            vecs[0].len = first.size();
-            vecs[1].base = const_cast<uint8_t*>(u8data(second));
-            vecs[1].len = second.size();
-            size_t vecs_size = first.empty() ? 0 : second.empty() ? 1 : 2;
-            Debug("Sending ", vecs[0].len, "+", vecs[1].len, " data for ", stream.id());
+            std::vector<ngtcp2_vec> vecs;
+            vecs.reserve(bufs.size());
+            std::transform(bufs.begin(), bufs.end(), std::back_inserter(vecs),
+                    [](const auto& buf) { return ngtcp2_vec{const_cast<uint8_t*>(u8data(buf)), buf.size()}; });
+
+#ifndef NDEBUG
+            {
+                std::string buf_sizes;
+                for (auto& b : bufs) {
+                    if (!buf_sizes.empty()) buf_sizes += '+';
+                    buf_sizes += std::to_string(b.size());
+                }
+                Debug("Sending ", buf_sizes.empty() ? "no" : buf_sizes, " data for ", stream.id());
+            }
+#endif
 
             uint32_t extra_flags = 0;
             if (stream.is_closing && !stream.sent_fin) {
@@ -539,7 +548,7 @@ void Connection::flush_streams() {
                 stream.is_new = false;
             }
 
-            auto [nwrite, consumed] = add_stream_data(stream.id(), vecs.data(), vecs_size, extra_flags);
+            auto [nwrite, consumed] = add_stream_data(stream.id(), vecs.data(), vecs.size(), extra_flags);
             Debug("add_stream_data for stream ", stream.id(), " returned [", nwrite, ",", consumed, "]");
 
             if (nwrite > 0) {

endpoint.h

@@ -20,7 +20,8 @@
 #include <uvw/poll.h>
 #include <uvw/timer.h>
 
-#if defined(__linux__) && !defined(NO_RECVMMSG)
+// True if we support recvmmsg/sendmmsg
+#if defined(__linux__) && !defined(LOKINET_NO_RECVMMSG)
 #  define LOKINET_HAVE_RECVMMSG
 #endif
 

stream.cpp

@@ -67,19 +67,24 @@ Stream::Stream(Connection& conn, StreamID id, size_t buffer_size)
 void Stream::set_buffer_size(size_t size) {
     if (used() != 0)
         throw std::runtime_error{"Cannot update buffer size while buffer is in use"};
-    if (size == 0)
-        size = 64*1024;
-    else if (size < 2048)
+    if (size > 0 && size < 2048)
         size = 2048;
 
     buffer.resize(size);
     buffer.shrink_to_fit();
     start = size = unacked_size = 0;
 }
 
+size_t Stream::buffer_size() const {
+    return buffer.empty()
+        ? size + start // start is the acked amount of the first buffer
+        : buffer.size();
+}
+
 bool Stream::append(bstring_view data) {
-    size_t avail = available();
-    if (avail < data.size())
+    assert(!buffer.empty());
+
+    if (data.size() > available())
         return false;
 
     // When we are appending we have three cases:
@@ -107,14 +112,20 @@ bool Stream::append(bstring_view data) {
     return true;
 }
 size_t Stream::append_any(bstring_view data) {
-    size_t avail = available();
-    if (data.size() > avail)
+    if (size_t avail = available(); data.size() > avail)
         data.remove_suffix(data.size() - avail);
     [[maybe_unused]] bool appended = append(data);
     assert(appended);
     return data.size();
 }
 
+void Stream::append_buffer(const std::byte* buffer, size_t length) {
+    assert(this->buffer.empty());
+    user_buffers.emplace_back(buffer, length);
+    size += length;
+    conn.io_ready();
+}
+
 void Stream::acknowledge(size_t bytes) {
     // Frees bytes; e.g. acknowledge(3) changes:
     //     [  áaaaaarr  ]  to  [     áaarr  ]
@@ -128,21 +139,61 @@ void Stream::acknowledge(size_t bytes) {
 
     unacked_size -= bytes;
     size -= bytes;
-    start = size == 0 ? 0 : (start + bytes) % buffer.size(); // reset start to 0 (to reduce wrapping buffers) if empty
+    if (!buffer.empty())
+        start = size == 0 ? 0 : (start + bytes) % buffer.size(); // reset start to 0 (to reduce wrapping buffers) if empty
+    else if (size == 0) {
+        user_buffers.clear();
+        start = 0;
+    } else {
+        while (bytes) {
+            assert(!user_buffers.empty());
+            assert(start < user_buffers.front().second);
+            if (size_t remaining = user_buffers.front().second - start;
+                    bytes >= remaining) {
+                user_buffers.pop_front();
+                start = 0;
+                bytes -= remaining;
+            } else {
+                start += bytes;
+                bytes = 0;
+            }
+        }
+    }
+
     if (!unblocked_callbacks.empty())
         available_ready();
 }
 
-std::pair<bstring_view, bstring_view> Stream::pending() {
-    std::pair<bstring_view, bstring_view> bufs;
-    if (size_t rsize = unsent(); rsize > 0) {
+auto get_buffer_it(std::deque<std::pair<std::unique_ptr<const std::byte[]>, size_t>>& bufs, size_t offset) {
+    auto it = bufs.begin();
+    while (offset >= it->second) {
+        offset -= it->second;
+        it++;
+    }
+    return std::make_pair(std::move(it), offset);
+}
+
+std::vector<bstring_view> Stream::pending() {
+    std::vector<bstring_view> bufs;
+    size_t rsize = unsent();
+    if (!rsize) return bufs;
+    if (!buffer.empty()) {
         size_t rpos = (start + unacked_size) % buffer.size();
         if (size_t rend = rpos + rsize; rend <= buffer.size()) {
-            bufs.first = {buffer.data() + rpos, rsize};
+            bufs.emplace_back(buffer.data() + rpos, rsize);
         } else { // wrapping
-            bufs.first = {buffer.data() + rpos, buffer.size() - rpos};
-            bufs.second = {buffer.data(), rend % buffer.size()};
+            bufs.reserve(2);
+            bufs.emplace_back(buffer.data() + rpos, buffer.size() - rpos);
+            bufs.emplace_back(buffer.data(), rend % buffer.size());
         }
+    } else {
+        assert(!user_buffers.empty()); // If empty then unsent() should have been 0
+        auto [it, offset] = get_buffer_it(user_buffers, start + unacked_size);
+        bufs.reserve(std::distance(it, user_buffers.end()));
+        assert(it != user_buffers.end());
+        bufs.emplace_back(it->first.get() + offset, it->second - offset);
+        for (++it; it != user_buffers.end(); ++it)
+            bufs.emplace_back(it->first.get(), it->second);
     }
     return bufs;
 }
@@ -153,6 +204,10 @@ void Stream::when_available(unblocked_callback_t unblocked_cb) {
 }
 
 void Stream::handle_unblocked() {
+    if (buffer.empty()) {
+        while (!unblocked_callbacks.empty() && unblocked_callbacks.front()(*this))
+            unblocked_callbacks.pop();
+    }
     while (!unblocked_callbacks.empty() && available() > 0) {
         if (unblocked_callbacks.front()(*this))
             unblocked_callbacks.pop();
@@ -171,6 +226,7 @@ void Stream::wrote(size_t bytes) {
     //     [  áaarrrrrr  ]  or  [rr     áaar]
     // to:
     //     [  áaaaaarrr  ]  or  [aa     áaaa]
+    Debug("wrote ", bytes, ", unsent=",unsent());
     assert(bytes <= unsent());
     unacked_size += bytes;
 }

stream.h

@@ -67,14 +67,17 @@ class Stream : public std::enable_shared_from_this<Stream> {
 
     // Sets the size of the outgoing data buffer.  This may *only* be used if the buffer is
     // currently entirely empty; otherwise a runtime_error is thrown.  The minimum buffer size is
-    // 2048, the default is 64kiB.
+    // 2048, the default is 64kiB.  A value of 0 puts the Stream into user-provided buffer mode
+    // where only the version of `append` taking ownership of a char* is permitted.
     void set_buffer_size(size_t size);
 
-    // Returns the size of the buffer (including both pending and free space).
-    size_t buffer_size() const { return buffer.size(); }
+    // Returns the size of the buffer (including both pending and free space).  If using
+    // user-provided buffer mode then this is the sum of all held buffers.
+    size_t buffer_size() const;
 
-    // Returns the number of free bytes available in the outgoing stream data buffer
-    size_t available() const { return is_closing ? 0 : buffer.size() - size; }
+    // Returns the number of free bytes available in the outgoing stream data buffer. Always 0 in
+    // user-provided buffer mode.
+    size_t available() const { return is_closing || buffer.empty() ? 0 : buffer.size() - size; }
 
     // Returns the number of bytes currently referenced in the buffer (i.e. pending or
     // sent-but-unacknowledged).
@@ -104,6 +107,11 @@ class Stream : public std::enable_shared_from_this<Stream> {
         return append_any(bstring_view{reinterpret_cast<const std::byte*>(data.data()), data.size()});
     }
 
+    // Takes ownership of the given buffer pointer, queuing it to be sent after any existing buffers
+    // and freed once fully acked.  You *must* have called `set_buffer_size(0)` (or set the
+    // endpoints default_stream_buffer_size to 0) in order to use this.
+    void append_buffer(const std::byte* buf, size_t length);
+
     // Starting closing the stream and prevent any more outgoing data from being appended.  If
     // `error_code` is provided then we close immediately with the given code; if std::nullopt (the
     // default) we close gracefully by sending a FIN bit.
@@ -142,6 +150,11 @@ class Stream : public std::enable_shared_from_this<Stream> {
     // Available callbacks should only be used when the buffer is full, typically immediately after
     // an `append_any` call that returns less than the full write.  Similarly a false return from an
     // unblock function (which keeps the callback alive) should satisfy the same condition.
+    //
+    // In user-provided buffer mode the callback will be invoked after any data has been acked: it
+    // is up to the caller to look at used()/buffer_size()/etc. to decide what to do.  As described
+    // above, return true to remove this callback, false to keep it and try again after the next
+    // ack.
     void when_available(unblocked_callback_t unblocked_cb);
 
     // Calls io_ready() on the stream's connection to scheduling sending outbound data
@@ -188,12 +201,10 @@ class Stream : public std::enable_shared_from_this<Stream> {
     // appending data).
     void acknowledge(size_t bytes);
 
-    // Returns a view into unwritten stream data.  This returns two string_views: if there is no
-    // pending data to write then both are empty; if the data to write does not wrap the buffer then
-    // .second will be empty and .first contains the data; if it wraps then both are non-empty and
-    // the .second buffer data immediately follows the .first buffer data.  After writing any of the
-    // provided data you should call `wrote()` to signal how much data you consumed.
-    std::pair<bstring_view, bstring_view> pending();
+    // Returns a view into unwritten stream data.  This returns a vector of string_views of the data
+    // to write, in order.  After writing any of the provided data you must call `wrote()` to signal
+    // how much of the given data was consumed (to advance the next pending() call).
+    std::vector<bstring_view> pending();
 
     // Called to signal that bytes have been written and should now be considered sent (but still
     // unacknowledged), thereby advancing the initial data position returned by the next `pending()`
@@ -205,10 +216,17 @@ class Stream : public std::enable_shared_from_this<Stream> {
     StreamID stream_id{-1};
 
     // ring buffer of outgoing stream data that has not yet been acknowledged.  This cannot be
-    // resized once used as ngtcp2 will have pointers into the data.
+    // resized once used as ngtcp2 will have pointers into the data.  If this is empty then we are
+    // in user-provided buffer mode.
     std::vector<std::byte> buffer{65536};
 
+    // user-provided buffers; only used when `buffer` is empty (via a `set_buffer_size(0)` or a 0
+    // size given in the constructor).
+    std::deque<std::pair<std::unique_ptr<const std::byte[]>, size_t>> user_buffers;
+
     // Offset of the first used byte in the circular buffer, will always be in [0, buffer.size()).
+    // For user-provided buffers this is the starting offset in the currently sending user-provided
+    // buffer.
     size_t start{0};
 
     // Number of sent-but-unacked packets in the buffer (i.e. [start, start+unacked_size) are sent but

tunnel-client.cpp

@@ -91,6 +91,7 @@ int main(int argc, char *argv[]) {
         loop,
         dest_port // tunnel destination port
         );
+    tunnel_client->default_stream_buffer_size = 0; // We steal uvw's provided buffers
     quic::Debug("Initialized client");
 
     // Start listening for TCP connections:

tunnel-server.cpp

@@ -91,7 +91,8 @@ int main(int argc, char *argv[]) {
                 tcp.erase(error_handler);
                 tunnel::install_stream_forwarding(tcp, *stream);
                 assert(stream->used() == 0);
-                stream->append(quic::bstring_view{&tunnel::CONNECT_INIT, 1});
+
+                stream->append_buffer(new std::byte[1]{tunnel::CONNECT_INIT}, 1);
                 tcp.read();
             });
 
@@ -100,6 +101,7 @@ int main(int argc, char *argv[]) {
             return true;
         }
     };
+    s.default_stream_buffer_size = 0; // We steal uvw's provided buffers
     quic::Debug("Initialized server");
     std::cout << "Listening on localhost:" << listen_port << " with tunnel(s) to localhost port(s):";
     if (allowed_ports.empty())

tunnel.cpp

@@ -1,4 +1,5 @@
 #include "tunnel.h"
+#include "log.h"
 #include "stream.h"
 #include "uvw/tcp.h"
 
@@ -11,33 +12,22 @@ void on_outgoing_data(uvw::DataEvent& event, uvw::TCPHandle& client) {
     std::string_view data{event.data.get(), event.length};
     auto peer = client.peer();
     quic::Debug(peer.ip, ":", peer.port, " → lokinet ", quic::buffer_printer{data});
-    if (auto wrote = stream->append_any(data); wrote < data.size()) {
-        // This gets complicated: we've received some data to forward but the stream's
-        // internal buffer is full of unacknowledged data.  We have to basically pause the
-        // local socket until the situation improves, and keep a sort of "overflow" buffer
-        // here to be reinserted once the stream space frees up.
-        quic::Debug("quic tunnel is congested (wrote ", wrote, " of ", data.size(), " bytes; pausing local tcp connection reads");
+    // Steal the buffer from the DataEvent's unique_ptr<char[]>:
+    stream->append_buffer(reinterpret_cast<const std::byte*>(event.data.release()), event.length);
+    if (stream->used() >= PAUSE_SIZE) {
+        quic::Debug("quic tunnel is congested (have ", stream->used(), " bytes in flight); pausing local tcp connection reads");
         client.stop();
-        data.remove_prefix(wrote);
-        stream->when_available([
-                client=client.shared_from_this(),
-                // Steal the unique_ptr<char[]> from DataEvent (but have to use a shared_ptr because
-                // of std::function).
-                buffer=std::shared_ptr<char[]>(event.data.release()),
-                data=std::move(data)
-        ](quic::Stream& s) mutable {
-            if (auto wrote = s.append_any(data); wrote < data.size()) {
-                quic::Debug("quic tunnel is partially unstuck (wrote ", wrote, " of ", data.size(), " remaining bytes)");
-                data.remove_prefix(wrote);
-                return false; // Not done.
+        stream->when_available([](quic::Stream& s) {
+            auto client = s.data<uvw::TCPHandle>();
+            if (s.used() < PAUSE_SIZE) {
+                quic::Debug("quic tunnel is no longer congested; resuming tcp connection reading");
+                client->read();
+                return true;
             }
-
-            quic::Debug("quic tunnel is no longer congested; resuming tcp connection reading");
-            client->read();
-            return true;
+            return false;
         });
     } else {
-        quic::Debug("Sent ", wrote, " bytes");
+        quic::Debug("Queued ", event.length, " bytes");
     }
 }
 

tunnel.h

@@ -26,6 +26,10 @@ inline constexpr uint64_t ERROR_BAD_INIT{0x5471908};
 // failure)
 inline constexpr uint64_t ERROR_TCP{0x5471909};
 
+// We pause reading from the local TCP socket if we have more than this amount of outstanding
+// unacked data in the quic tunnel, then resume once it drops below this.
+inline constexpr size_t PAUSE_SIZE = 64*1024;
+
 // Callbacks for network events.  The uvw::TCPHandle client must contain a shared pointer to the
 // associated quic::Stream in its data, and the quic::Stream must contain a weak pointer to the
 // uvw::TCPHandle.