Adding pending data check so users can throttle sending of data

Author: smasherprog

Test/main.cpp

@@ -116,7 +116,7 @@ void generaltest()
     auto lastheard = std::chrono::high_resolution_clock::now();
 
     SL::WS_LITE::PortNumber port(3002);
-    auto listenerct =
+    auto listenerctx =
         SL::WS_LITE::CreateContext(SL::WS_LITE::ThreadCount(1))
             ->NoTLS()
             ->CreateListener(port)
@@ -158,6 +158,10 @@ void generaltest()
     while (std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::high_resolution_clock::now() - lastheard).count() < 2000) {
         std::this_thread::sleep_for(200ms);
     }
+
+    std::this_thread::sleep_for(200ms);
+    assert(clientctx->get_DataPending() == 0);
+    assert(listenerctx->get_DataPending() == 0);
 }
 void generalTLStest()
 {
@@ -259,6 +263,9 @@ void generalTLStest()
     while (std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::high_resolution_clock::now() - lastheard).count() < 2000) {
         std::this_thread::sleep_for(200ms);
     }
+    std::this_thread::sleep_for(200ms);
+    assert(clientctx->get_DataPending() == 0);
+    assert(listenerctx->get_DataPending() == 0);
 }
 
 void multithreadtest()
@@ -321,6 +328,11 @@ void multithreadtest()
     while (std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::high_resolution_clock::now() - lastheard).count() < 2000) {
         std::this_thread::sleep_for(200ms);
     }
+    std::this_thread::sleep_for(200ms);
+    for (auto &c : clients) {
+        assert(c->get_DataPending() == 0);
+    }
+    assert(listenerctx->get_DataPending() == 0);
 }
 const auto bufferesize = 1024 * 1024 * 10;
 void multithreadthroughputtest()
@@ -392,10 +404,14 @@ void multithreadthroughputtest()
                 ->connect("localhost", port);
         clients.push_back(c);
     }
-
     while (std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::high_resolution_clock::now() - lastheard).count() < 5000) {
         std::this_thread::sleep_for(200ms);
     }
+    std::this_thread::sleep_for(200ms);
+    for (auto &c : clients) {
+        assert(c->get_DataPending() == 0);
+    }
+    assert(listenerctx->get_DataPending() == 0);
     std::cout << "Received " << mbsreceived << "  bytes" << std::endl;
 }
 int main(int argc, char *argv[])

include/WS_Lite.h

@@ -3,8 +3,8 @@
 #include <functional>
 #include <memory>
 #include <string>
-#include <unordered_map>
 #include <system_error>
+#include <unordered_map>
 typedef struct x509_store_ctx_st X509_STORE_CTX;
 
 #if defined(WINDOWS) || defined(WIN32)
@@ -167,6 +167,8 @@ namespace WS_LITE {
         virtual bool is_v4() const = 0;
         virtual bool is_v6() const = 0;
         virtual bool is_loopback() const = 0;
+        // this will return the amount of data queued up but has not yet been submitted to the OS to be sent
+        virtual size_t get_DataPending() const = 0;
         virtual void send(const WSMessage &msg, bool compressmessage) = 0;
         // send a close message and close the socket
         virtual void close(unsigned short code = 1000, const std::string &msg = "") = 0;
@@ -186,6 +188,8 @@ namespace WS_LITE {
         virtual void set_WriteTimeout(std::chrono::seconds seconds) = 0;
         // get the current write timeout in seconds
         virtual std::chrono::seconds get_WriteTimeout() = 0;
+        // this will return the amount of data queued up but has not yet been submitted to the OS to be sent
+        virtual size_t get_DataPending() const = 0;
     };
     class WS_LITE_EXTERN IWSListener_Configuration {
       public:

include/internal/DataStructures.h

@@ -104,6 +104,7 @@ namespace WS_LITE {
         std::chrono::seconds WriteTimeout = std::chrono::seconds(30);
         size_t MaxPayload = 1024 * 1024 * 20; // 20 MB
         bool TLSEnabled = false;
+        size_t DataPending = 0;
 
         std::function<void(const std::shared_ptr<IWSocket> &, const std::unordered_map<std::string, std::string> &)> onConnection;
         std::function<void(const std::shared_ptr<IWSocket> &, const WSMessage &)> onMessage;
@@ -163,7 +164,7 @@ namespace WS_LITE {
                     sendclosemessage<isServer>(p, self, code, msg);
             }
         }
-
+        virtual size_t get_DataPending() const override { return DataPending; }
         void canceltimers()
         {
             std::error_code ec;
@@ -188,6 +189,7 @@ namespace WS_LITE {
         asio::basic_waitable_timer<std::chrono::steady_clock> write_deadline;
         asio::strand strand;
         std::deque<SendQueueItem> SendMessageQueue;
+        size_t DataPending = 0;
     };
 
 } // namespace WS_LITE

include/internal/WebSocketProtocol.h

@@ -195,12 +195,15 @@ namespace WS_LITE {
         }
 
         socket->strand.post([socket, msg, parent, compressmessage]() {
+
             if (socket->SocketStatus_ == SocketStatus::CONNECTED) {
                 // update the socket status to reflect it is closing to prevent other messages from being sent.. this is the last valid message
                 // make sure to do this after a call to write so the write process sends the close message, but no others
                 if (msg.code == OpCode::CLOSE) {
                     socket->SocketStatus_ = SocketStatus::CLOSING;
                 }
+                socket->DataPending += msg.len;
+                socket->Parent->DataPending += msg.len;
                 socket->SendMessageQueue.emplace_back(SendQueueItem{msg, compressmessage});
                 SL::WS_LITE::startwrite<isServer>(parent, socket);
             }
@@ -254,6 +257,8 @@ namespace WS_LITE {
                                   return handleclose(parent, socket, 1002, "write header failed " + ec.message());
                               }
                               assert(msg.len == bytes_transferred);
+                              socket->DataPending -= msg.len;
+                              socket->Parent->DataPending -= msg.len;
                               startwrite<isServer>(parent, socket);
 
                           }));
@@ -552,6 +557,7 @@ namespace WS_LITE {
         virtual std::chrono::seconds get_ReadTimeout() override;
         virtual void set_WriteTimeout(std::chrono::seconds seconds) override;
         virtual std::chrono::seconds get_WriteTimeout() override;
+        virtual size_t get_DataPending() const override { return Impl_->DataPending; }
     };
     class WSListener : public IWSHub {
         std::shared_ptr<WSContextImpl> Impl_;
@@ -565,6 +571,7 @@ namespace WS_LITE {
         virtual std::chrono::seconds get_ReadTimeout() override;
         virtual void set_WriteTimeout(std::chrono::seconds seconds) override;
         virtual std::chrono::seconds get_WriteTimeout() override;
+        virtual size_t get_DataPending() const override { return Impl_->DataPending; }
     };
 
     class WSListener_Configuration : public IWSListener_Configuration {