Sockets: add an optional deadline to ClientSocket::send()

Common/Cpp/Sockets/AbstractClientSocket.h

@@ -7,6 +7,7 @@
 #ifndef PokemonAutomation_Sockets_AbstractClientSocket_H
 #define PokemonAutomation_Sockets_AbstractClientSocket_H
 
+#include <chrono>
 #include <string>
 #include "Common/Cpp/ListenerSet.h"
 
@@ -63,7 +64,16 @@ class AbstractClientSocket{
     virtual void close() noexcept = 0;
     virtual void connect(const std::string& address, uint16_t port) = 0;
 
-    virtual size_t send(const void* data, size_t bytes) = 0;
+    //  Sends data on the active connection. This blocks until every byte
+    //  has been written, the connection is lost, or "timeout" elapses.
+    //  "std::chrono::milliseconds::max()" (the default) waits indefinitely.
+    //  Returns the number of bytes actually written, which may be less
+    //  than "bytes" if the deadline elapsed or the connection was lost.
+    virtual size_t send(
+        const void* data,
+        size_t bytes,
+        std::chrono::milliseconds timeout = std::chrono::milliseconds::max()
+    ) = 0;
 
 
 protected:

Common/Cpp/Sockets/ClientSocket_POSIX.cpp

@@ -119,11 +119,23 @@ void ClientSocket_POSIX::close_socket_locked(){
     }
 }
 
-size_t ClientSocket_POSIX::send(const void* data, size_t bytes){
+size_t ClientSocket_POSIX::send(const void* data, size_t bytes, std::chrono::milliseconds timeout){
     const char* ptr = (const char*)data;
     size_t sent = 0;
 
+    //  "time_point::max()" (the infinite default) is never reached; using
+    //  it also avoids overflowing "now() + timeout".
+    const std::chrono::steady_clock::time_point deadline =
+        (timeout == std::chrono::milliseconds::max())
+            ? std::chrono::steady_clock::time_point::max()
+            : std::chrono::steady_clock::now() + timeout;
+
     while (bytes > 0 && state() == State::CONNECTED){
+        if (std::chrono::steady_clock::now() >= deadline){
+            //  The deadline elapsed with bytes still unwritten; the
+            //  connection is still fine, so just report the partial count.
+            return sent;
+        }
         int socket = m_socket.load(std::memory_order_acquire);
         if (socket == -1){
             return sent;

Common/Cpp/Sockets/ClientSocket_POSIX.h

@@ -8,6 +8,7 @@
 #define PokemonAutomation_ClientSocket_POSIX_H
 
 #include <atomic>
+#include <chrono>
 #include "Common/Cpp/Concurrency/Mutex.h"
 #include "Common/Cpp/Concurrency/ConditionVariable.h"
 #include "Common/Cpp/Concurrency/AsyncTask.h"
@@ -38,7 +39,11 @@ class ClientSocket_POSIX final : public AbstractClientSocket{
 
     virtual void connect(const std::string& address, uint16_t port) override;
 
-    virtual size_t send(const void* data, size_t bytes) override;
+    virtual size_t send(
+        const void* data,
+        size_t bytes,
+        std::chrono::milliseconds timeout = std::chrono::milliseconds::max()
+    ) override;
 
 
 private:

Common/Cpp/Sockets/ClientSocket_Qt.cpp

@@ -84,7 +84,7 @@ void ClientSocket_Qt::connect(const std::string& address, uint16_t port){
     emit internal_connect(address, port);
 }
 
-size_t ClientSocket_Qt::send(const void* data, size_t bytes){
+size_t ClientSocket_Qt::send(const void* data, size_t bytes, std::chrono::milliseconds timeout){
     QTcpSocket* socket;
     {
         std::lock_guard<Mutex> lg(m_lock);
@@ -96,20 +96,31 @@ size_t ClientSocket_Qt::send(const void* data, size_t bytes){
 
     QSharedPointer<SendData> send_data = QSharedPointer<SendData>::create();
     send_data->packet.assign((const char*)data, bytes);
+    //  "time_point::max()" (the infinite default) is never reached; using
+    //  it also avoids overflowing "now() + timeout".
+    send_data->deadline =
+        (timeout == std::chrono::milliseconds::max())
+            ? std::chrono::steady_clock::time_point::max()
+            : std::chrono::steady_clock::now() + timeout;
 
     emit internal_send(send_data);
 
     //  Wait for the socket's event thread to finish writing the packet.
-    //  The timeout covers the case where "close()" destroys the socket
-    //  before the queued call ever runs, in which case it will never
-    //  complete on its own. The shared ownership makes the timeout
-    //  safe: if it expires, the queued call's own copy keeps the
-    //  payload alive until it runs.
+    //  The write loop itself stops at "deadline", so this bound only
+    //  covers the queued call's start-up delay. The infinite default
+    //  keeps a bounded 5 s safety: if "close()" destroys the socket
+    //  before the queued call ever runs, it will never complete on its
+    //  own. The shared ownership makes the timeout safe: if it expires,
+    //  the queued call's own copy keeps the payload alive until it runs.
+    const std::chrono::milliseconds wait_time =
+        (timeout == std::chrono::milliseconds::max())
+            ? std::chrono::milliseconds(5000)
+            : timeout + std::chrono::milliseconds(1000);
     {
         std::unique_lock<Mutex> lg(send_data->lock);
         send_data->cv.wait_for(
             lg,
-            std::chrono::seconds(5),
+            wait_time,
             [&send_data]{ return send_data->done; }
         );
     }
@@ -252,7 +263,14 @@ std::unique_ptr<QTcpSocket> ClientSocket_Qt::make_socket(){
             {
                 const char* ptr = send_data->packet.data();
                 size_t remaining = send_data->packet.size();
-                while (remaining > 0 && sock->state() == QAbstractSocket::ConnectedState){
+                //  The deadline keeps a peer that never reads from pinning
+                //  this event thread forever (which would also block
+                //  "close()" on it).
+                while (
+                    remaining > 0
+                    && sock->state() == QAbstractSocket::ConnectedState
+                    && std::chrono::steady_clock::now() < send_data->deadline
+                ){
                     qint64 current_sent = sock->write(ptr, remaining);
                     if (current_sent <= 0){
                         //  The kernel's send buffer is full. Wait for

Common/Cpp/Sockets/ClientSocket_Qt.h

@@ -7,6 +7,7 @@
 #ifndef PokemonAutomation_ClientSocket_Qt_H
 #define PokemonAutomation_ClientSocket_Qt_H
 
+#include <chrono>
 #include <string>
 #include <atomic>
 #include <QObject>
@@ -28,6 +29,11 @@ class ClientSocket_Qt final : public QObject, public AbstractClientSocket{
 
     struct SendData{
         std::string packet;
+        //  The event-thread write loop stops at this deadline, so a send
+        //  that can never complete (a peer that never reads) cannot pin
+        //  the socket's event thread forever and block "close()". "max()"
+        //  means no deadline (the infinite default).
+        std::chrono::steady_clock::time_point deadline = std::chrono::steady_clock::time_point::max();
         size_t bytes_sent = 0;
         bool done = false;
 
@@ -56,7 +62,11 @@ class ClientSocket_Qt final : public QObject, public AbstractClientSocket{
 
     virtual void connect(const std::string& address, uint16_t port) override;
 
-    virtual size_t send(const void* data, size_t bytes) override;
+    virtual size_t send(
+        const void* data,
+        size_t bytes,
+        std::chrono::milliseconds timeout = std::chrono::milliseconds::max()
+    ) override;
 
 
 signals:

Common/Cpp/Sockets/ClientSocket_WinSocket.cpp

@@ -100,11 +100,23 @@ void ClientSocket_WinSocket::connect(const std::string& address, uint16_t port){
 }
 
 
-size_t ClientSocket_WinSocket::send(const void* data, size_t bytes){
+size_t ClientSocket_WinSocket::send(const void* data, size_t bytes, std::chrono::milliseconds timeout){
     const char* ptr = (const char*)data;
     size_t sent = 0;
 
+    //  "time_point::max()" (the infinite default) is never reached; using
+    //  it also avoids overflowing "now() + timeout".
+    const std::chrono::steady_clock::time_point deadline =
+        (timeout == std::chrono::milliseconds::max())
+            ? std::chrono::steady_clock::time_point::max()
+            : std::chrono::steady_clock::now() + timeout;
+
     while (bytes > 0 && state() == State::CONNECTED){
+        if (std::chrono::steady_clock::now() >= deadline){
+            //  The deadline elapsed with bytes still unwritten; the
+            //  connection is still fine, so just report the partial count.
+            return sent;
+        }
         SOCKET socket = m_socket.load(std::memory_order_acquire);
         if (socket == INVALID_SOCKET){
             return sent;

Common/Cpp/Sockets/ClientSocket_WinSocket.h

@@ -8,6 +8,7 @@
 #define PokemonAutomation_ClientSocket_WinSocket_H
 
 #include <atomic>
+#include <chrono>
 #include <mutex>
 #include <winsock2.h>
 #include "Common/Cpp/Concurrency/Mutex.h"
@@ -42,7 +43,11 @@ class ClientSocket_WinSocket final : public AbstractClientSocket{
 
     virtual void connect(const std::string& address, uint16_t port) override;
 
-    virtual size_t send(const void* data, size_t bytes) override;
+    virtual size_t send(
+        const void* data,
+        size_t bytes,
+        std::chrono::milliseconds timeout = std::chrono::milliseconds::max()
+    ) override;
 
 
 private:

Common/Cpp/Sockets/Sockets_Tests.cpp

@@ -20,6 +20,7 @@
 #include <random>
 #include <string>
 #include <thread>
+#include <type_traits>
 #include <vector>
 #include "Common/Cpp/Concurrency/AsyncTask.h"
 #include "Common/Cpp/Concurrency/ConditionVariable.h"
@@ -843,6 +844,166 @@ class Test_Sockets_Qt_CompletionLifetime : public UnitTest{
 };
 
 
+//  The peer accepts the connection but stalls reading: a send() with a
+//  deadline must stop at the deadline with a partial count, and
+//  close() must still be able to finish afterwards (on the Qt backend
+//  the in-flight write stops at the deadline, releasing the socket's
+//  event thread, which is what close() has to run on).
+//
+//  The stall lasts 12 s - longer than the 10 s send threshold below - so
+//  the peer is still stalled for the whole duration of every send in the
+//  retry loop. It is interruptible ("stop"), so shutting the peer down
+//  does not wait for it.
+//
+//  An OS with unusually large autotuned socket buffers (observed on
+//  Windows CI, where the kernel buffered the entire 256 MB payload
+//  against a non-reading peer) lets the first send "complete" by parking
+//  the data in the kernel. Those buffers are now full, so a follow-up
+//  send to the same stalled peer cannot complete and must stop at the
+//  deadline with a partial count; the loop retries until that happens
+//  (up to 4 times, i.e. unless the OS can buffer over 1 GiB).
+//
+//  Note: QTcpSocket::write() accepts the whole payload into Qt's internal
+//  buffer and reports it as written, so kernel backpressure never reaches
+//  the Qt write loop and the send legitimately "completes" there; a
+//  retry cannot exercise the deadline on that backend. On that backend
+//  only the timing and close() behavior are asserted.
+template <typename ClientSocketType, typename ServerSocketType>
+class Test_Sockets_ClientSendTimeout : public TestSockets<ClientSocketType, ServerSocketType>{
+public:
+    using TestSockets<ClientSocketType, ServerSocketType>::TestSockets;
+
+    virtual UnitTestResult run(Logger& logger, CancellableScope& scope) const override{
+        constexpr bool QT = std::is_same_v<ClientSocketType, ClientSocket_Qt>;
+
+        //  Stalls the server's reading for 12 s on the first data (or
+        //  until "stop" is set), keeping the client's send buffer full
+        //  for the duration of the send test.
+        struct StallingServerListener : ServerListener{
+            std::atomic<bool> stalled{false};
+            std::atomic<bool> stop{false};
+
+            virtual void on_receive_data(const void* data, size_t bytes) override{
+                if (!stalled.exchange(true)){
+                    const auto end = std::chrono::steady_clock::now() + std::chrono::seconds(12);
+                    while (std::chrono::steady_clock::now() < end
+                        && !stop.load(std::memory_order_acquire)){
+                        std::this_thread::sleep_for(std::chrono::milliseconds(10));
+                    }
+                }
+                ServerListener::on_receive_data(data, bytes);
+            }
+        };
+
+        ServerSocketType server(this->m_thread_pool);
+        StallingServerListener server_listener;
+        server.add_listener(server_listener);
+
+        uint16_t port;
+        if (!this->start_server(server, server_listener, port)){
+            return "Failed to start the server";
+        }
+
+        bool ok = false;
+        std::string error;
+        {
+            ClientSocketType client(this->m_thread_pool);
+            ClientListener client_listener;
+            client.add_listener(client_listener);
+
+            if (!this->connect_client(client, client_listener, port, 1)){
+                error = "The client did not finish connecting";
+            }else if (!server_listener.waiter.wait_for(
+                std::chrono::seconds(5),
+                [&server_listener]{ return server_listener.connected.load() >= 1; }
+            )){
+                error = "The server did not call on_connect";
+            }else{
+                //  Far larger than the kernel send/receive buffering that
+                //  the send loop can rely on being absent, so the send
+                //  must stop on the stalled reader (see the comment on
+                //  the class for the large-buffer OS case).
+                constexpr size_t SIZE = 1 << 28;
+                std::vector<uint8_t> payload(SIZE, 0xAB);
+
+                //  The deadline is 100 ms, but this test also allocates
+                //  and moves 256 MB (the Qt backend additionally copies it
+                //  into its packet buffer), and on a loaded CI runner that
+                //  page traffic - or the calling thread being descheduled
+                //  - can take seconds. The pre-fix behavior was an infinite
+                //  hang, so the threshold only has to prove boundedness;
+                //  10 s keeps the test meaningful while tolerating CI
+                //  load.
+                int completed = 0;
+                while (true){
+                    auto t0 = std::chrono::steady_clock::now();
+                    size_t sent = client.send(
+                        payload.data(),
+                        payload.size(),
+                        std::chrono::milliseconds(100)
+                    );
+                    auto send_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
+                        std::chrono::steady_clock::now() - t0
+                    ).count();
+
+                    if (send_ms > 10000){
+                        error = "send() did not return within 10 s of its 100 ms deadline (" + std::to_string(send_ms) + " ms)";
+                        break;
+                    }else if (sent < payload.size()){
+                        //  The deadline engaged with bytes left unwritten.
+                        break;
+                    }else if (QT){
+                        //  Qt's "completion" is its userspace buffer, not
+                        //  kernel delivery; a retry cannot exercise the
+                        //  deadline on that backend.
+                        break;
+                    }
+
+                    //  The whole payload was written, so the kernel's
+                    //  socket buffers absorbed it (unusually large
+                    //  autotuned buffers). They are now full, so the
+                    //  next send to the same stalled peer cannot
+                    //  complete.
+                    if (++completed >= 4){
+                        error = "All 4 sends of 256 MB completed: the OS buffered over 1 GiB against the stalled peer, so the deadline cannot be verified";
+                        break;
+                    }
+                }
+
+                if (error.empty()){
+                    auto c0 = std::chrono::steady_clock::now();
+                    client.close();
+                    auto close_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
+                        std::chrono::steady_clock::now() - c0
+                    ).count();
+                    if (close_ms > 5000){
+                        error = "close() did not finish after the timed-out send (" + std::to_string(close_ms) + " ms)";
+                    }else if (client.state() != AbstractClientSocket::State::NOT_RUNNING){
+                        error = "The client is not NOT_RUNNING after close()";
+                    }else{
+                        ok = true;
+                    }
+                }
+            }
+            client.close();
+        }
+
+        //  Release the stalled reader so the server's close() (which
+        //  joins its connection thread, currently inside the stall) does
+        //  not wait for it.
+        server_listener.stop.store(true, std::memory_order_release);
+
+        server.close();
+
+        if (!ok && !error.empty()){
+            return error;
+        }
+        return true;
+    }
+};
+
+
+
 template <typename ClientSocketType, typename ServerSocketType>
 void add_socket_tests(UnitTestDatabase& database, ThreadPool& thread_pool, const char* backend_name){
     const std::string prefix = "Sockets (" + std::string(backend_name) + ") - ";
@@ -868,6 +1029,9 @@ void add_socket_tests(UnitTestDatabase& database, ThreadPool& thread_pool, const
     database.add<Test_Sockets_PortInUse<ClientSocketType, ServerSocketType>>(
         thread_pool, prefix + "Port in use"
     );
+    database.add<Test_Sockets_ClientSendTimeout<ClientSocketType, ServerSocketType>>(
+        thread_pool, prefix + "Client send() timeout against a stalled peer"
+    );
 }