From 89e738262c7dd80daa3702432250e49fa57d7f55 Mon Sep 17 00:00:00 2001 From: Peter Johnson Date: Sun, 20 Aug 2023 13:53:52 -0700 Subject: [PATCH] [ntcore] Limit buffer pool size to 64KB per connection (#5485) Previously this was unlimited, which could result in holding on to a large amount of memory if the connection got backlogged or a burst of data transmission occurred. --- ntcore/src/main/native/cpp/net/WebSocketConnection.cpp | 9 ++++++++- .../src/main/native/cpp/net3/UvStreamConnection3.cpp | 10 +++++++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/ntcore/src/main/native/cpp/net/WebSocketConnection.cpp b/ntcore/src/main/native/cpp/net/WebSocketConnection.cpp index 2a30e4299b..24da037337 100644 --- a/ntcore/src/main/native/cpp/net/WebSocketConnection.cpp +++ b/ntcore/src/main/native/cpp/net/WebSocketConnection.cpp @@ -16,6 +16,7 @@ using namespace nt::net; static constexpr size_t kAllocSize = 4096; static constexpr size_t kTextFrameRolloverSize = 4096; static constexpr size_t kBinaryFrameRolloverSize = 8192; +static constexpr size_t kMaxPoolSize = 16; WebSocketConnection::WebSocketConnection(wpi::WebSocket& ws) : m_ws{ws}, @@ -53,7 +54,13 @@ void WebSocketConnection::Flush() { ++m_sendsActive; m_ws.SendFrames(m_ws_frames, [selfweak = weak_from_this()](auto bufs, auto) { if (auto self = selfweak.lock()) { - self->m_buf_pool.insert(self->m_buf_pool.end(), bufs.begin(), bufs.end()); + size_t numToPool = + (std::min)(bufs.size(), kMaxPoolSize - self->m_buf_pool.size()); + self->m_buf_pool.insert(self->m_buf_pool.end(), bufs.begin(), + bufs.begin() + numToPool); + for (auto&& buf : bufs.subspan(numToPool)) { + buf.Deallocate(); + } if (self->m_sendsActive > 0) { --self->m_sendsActive; } diff --git a/ntcore/src/main/native/cpp/net3/UvStreamConnection3.cpp b/ntcore/src/main/native/cpp/net3/UvStreamConnection3.cpp index 835950fd70..a2f42d8da9 100644 --- a/ntcore/src/main/native/cpp/net3/UvStreamConnection3.cpp +++ b/ntcore/src/main/native/cpp/net3/UvStreamConnection3.cpp @@ -10,6 +10,8 @@ using namespace nt; using namespace nt::net3; +static constexpr size_t kMaxPoolSize = 16; + UvStreamConnection3::UvStreamConnection3(wpi::uv::Stream& stream) : m_stream{stream}, m_os{m_buffers, [this] { return AllocBuf(); }} {} @@ -26,7 +28,13 @@ void UvStreamConnection3::Flush() { ++m_sendsActive; m_stream.Write(m_buffers, [selfweak = weak_from_this()](auto bufs, auto) { if (auto self = selfweak.lock()) { - self->m_buf_pool.insert(self->m_buf_pool.end(), bufs.begin(), bufs.end()); + size_t numToPool = + (std::min)(bufs.size(), kMaxPoolSize - self->m_buf_pool.size()); + self->m_buf_pool.insert(self->m_buf_pool.end(), bufs.begin(), + bufs.begin() + numToPool); + for (auto&& buf : bufs.subspan(numToPool)) { + buf.Deallocate(); + } if (self->m_sendsActive > 0) { --self->m_sendsActive; }