Make select/poll interruptible with a dummy wake socket

Author: theuni

src/common/sockman.cpp

@@ -332,7 +332,7 @@ void SockMan::ThreadSocketHandler()
         if (io_readiness.events_per_sock.empty() ||
             // WaitMany() may as well be a static method, the context of the first Sock in the vector is not relevant.
             !io_readiness.events_per_sock.begin()->first->WaitMany(SELECT_TIMEOUT,
-                                                                   io_readiness.events_per_sock)) {
+                                                                   io_readiness.events_per_sock, nullptr)) {
             interruptNet.sleep_for(SELECT_TIMEOUT);
         }
 

src/test/pcp_tests.cpp

@@ -206,7 +206,7 @@ class PCPTestSock final : public Sock
         return true;
     }
 
-    bool WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per_sock) const override
+    bool WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per_sock, std::shared_ptr<const Sock> wakesock) const override
     {
         return false;
     }

src/test/util/net.cpp

@@ -208,7 +208,7 @@ bool ZeroSock::Wait(std::chrono::milliseconds timeout, Event requested, Event* o
     return true;
 }
 
-bool ZeroSock::WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per_sock) const
+bool ZeroSock::WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per_sock, std::shared_ptr<const Sock> wakesock) const
 {
     for (auto& [sock, events] : events_per_sock) {
         (void)sock;
@@ -362,14 +362,14 @@ bool DynSock::Wait(std::chrono::milliseconds timeout,
 {
     EventsPerSock ev;
     ev.emplace(this, Events{requested});
-    const bool ret{WaitMany(timeout, ev)};
+    const bool ret{WaitMany(timeout, ev, nullptr)};
     if (occurred != nullptr) {
         *occurred = ev.begin()->second.occurred;
     }
     return ret;
 }
 
-bool DynSock::WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per_sock) const
+bool DynSock::WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per_sock, std::shared_ptr<const Sock> wakesock) const
 {
     const auto deadline = std::chrono::steady_clock::now() + timeout;
     bool at_least_one_event_occurred{false};
@@ -392,6 +392,10 @@ bool DynSock::WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_
                 }
             }
         }
+        if (wakesock) {
+            std::array<std::byte, 1024> bytes;
+            [[maybe_unused]] size_t read_bytes = wakesock->Recv(bytes.data(), bytes.size(), MSG_DONTWAIT);
+        }
 
         if (at_least_one_event_occurred || std::chrono::steady_clock::now() > deadline) {
             break;

src/test/util/net.h

@@ -180,7 +180,7 @@ class ZeroSock : public Sock
               Event requested,
               Event* occurred = nullptr) const override;
 
-    bool WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per_sock) const override;
+    bool WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per_sock, std::shared_ptr<const Sock> wakesock) const override;
 
 private:
     ZeroSock& operator=(Sock&& other) override;
@@ -333,7 +333,7 @@ class DynSock : public ZeroSock
               Event requested,
               Event* occurred = nullptr) const override;
 
-    bool WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per_sock) const override;
+    bool WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per_sock, std::shared_ptr<const Sock> wakesock) const override;
 
 private:
     DynSock& operator=(Sock&&) override;

src/util/sock.cpp

@@ -145,7 +145,7 @@ bool Sock::Wait(std::chrono::milliseconds timeout, Event requested, Event* occur
 
     EventsPerSock events_per_sock{std::make_pair(shared, Events{requested})};
 
-    if (!WaitMany(timeout, events_per_sock)) {
+    if (!WaitMany(timeout, events_per_sock, nullptr)) {
         return false;
     }
 
@@ -156,7 +156,7 @@ bool Sock::Wait(std::chrono::milliseconds timeout, Event requested, Event* occur
     return true;
 }
 
-bool Sock::WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per_sock) const
+bool Sock::WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per_sock, std::shared_ptr<const Sock> wakesock) const
 {
 #ifdef USE_POLL
     std::vector<pollfd> pfds;
@@ -172,11 +172,27 @@ bool Sock::WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per
         }
     }
 
+    if (wakesock) {
+        pfds.emplace_back();
+        auto& pfd = pfds.back();
+        pfd.fd = wakesock->m_socket;
+        pfd.events |= POLLIN;
+    }
+
     if (poll(pfds.data(), pfds.size(), count_milliseconds(timeout)) == SOCKET_ERROR) {
         return false;
     }
 
-    assert(pfds.size() == events_per_sock.size());
+    if (wakesock) {
+        assert(pfds.size() == events_per_sock.size() + 1);
+        if (pfds.back().revents & POLLIN) {
+            std::array<std::byte, 1024> bytes;
+            [[maybe_unused]] size_t read_bytes = wakesock->Recv(bytes.data(), bytes.size(), MSG_DONTWAIT);
+        }
+    } else {
+        assert(pfds.size() == events_per_sock.size());
+    }
+
     size_t i{0};
     for (auto& [sock, events] : events_per_sock) {
         assert(sock->m_socket == static_cast<SOCKET>(pfds[i].fd));
@@ -218,12 +234,26 @@ bool Sock::WaitMany(std::chrono::milliseconds timeout, EventsPerSock& events_per
         socket_max = std::max(socket_max, s);
     }
 
+    if (wakesock) {
+        if (!wakesock->IsSelectable()) {
+            return false;
+        }
+        const auto& s = wakesock->m_socket;
+        FD_SET(s, &recv);
+        socket_max = std::max(socket_max, s);
+    }
+
     timeval tv = MillisToTimeval(timeout);
 
     if (select(socket_max + 1, &recv, &send, &err, &tv) == SOCKET_ERROR) {
         return false;
     }
 
+    if (wakesock && FD_ISSET(wakesock->m_socket, &recv)) {
+        std::array<std::byte, 1024> bytes;
+        [[maybe_unused]] size_t read_bytes = wakesock->Recv(bytes.data(), bytes.size(), MSG_DONTWAIT);
+    }
+
     for (auto& [sock, events] : events_per_sock) {
         const auto& s = sock->m_socket;
         events.occurred = 0;

src/util/sock.h

@@ -216,7 +216,7 @@ class Sock
      * false otherwise
      */
     [[nodiscard]] virtual bool WaitMany(std::chrono::milliseconds timeout,
-                                        EventsPerSock& events_per_sock) const;
+                                        EventsPerSock& events_per_sock, std::shared_ptr<const Sock> wakesock) const;
 
     /* Higher level, convenience, methods. These may throw. */