Enhance TCPServer with poll set management

- Introduced a cached poll set to optimize socket polling, reducing allocations during steady state.
- Added methods to rebuild the poll set on client connections and disconnections.
- Implemented tests to ensure the poll set remains in sync with connected clients and to prevent message loss during concurrent client connections.
- Updated spin() method to utilize the cached poll set for improved performance.

Author: urrsk

include/ur_client_library/comm/tcp_server.h

@@ -40,6 +40,10 @@
 
 #include "ur_client_library/comm/socket_t.h"
 
+#ifndef _WIN32
+#  include <poll.h>
+#endif
+
 namespace urcl
 {
 namespace comm
@@ -198,6 +202,14 @@ class TCPServer
     return port_;
   }
 
+protected:
+  // Test hook: number of file descriptors currently in the cached poll set
+  // (listen socket + one per connected client).
+  size_t getPollSetSize() const
+  {
+    return pollfds_.size();
+  }
+
 private:
   void init();
   void bind(const size_t max_num_tries, const std::chrono::milliseconds reconnection_time);
@@ -217,6 +229,15 @@ class TCPServer
   //! Runs spin() as long as keep_running_ is set to true.
   void worker();
 
+  //! Rebuilds the cached poll set (pollfds_) from the listen socket and client_fds_.
+  void rebuildPollfds();
+
+  // Number of client slots the poll set reserves up front when the client count is unbounded
+  // (max_clients_allowed_ == 0). Bounded servers reserve their exact limit instead. This is just
+  // headroom so the first few connects (incl. transient overlap during a robot reconnect) don't
+  // reallocate; the vector still grows geometrically beyond it.
+  static constexpr uint32_t DEFAULT_RESERVED_CLIENTS = 8;
+
   std::atomic<bool> keep_running_{ false };
   std::thread worker_thread_;
 
@@ -225,6 +246,18 @@ class TCPServer
 
   uint32_t max_clients_allowed_;
   std::vector<socket_t> client_fds_;
+
+#ifdef _WIN32
+  using PollFd = WSAPOLLFD;
+#else
+  using PollFd = struct pollfd;
+#endif
+  // Cached poll set: index 0 is the listen socket, the rest mirror client_fds_. Only touched by
+  // the worker thread (seeded in start() before the thread launches, rebuilt in
+  // handleConnect()/handleDisconnect(), cleared in shutdown() after the worker is joined), so
+  // spin() can poll() on it directly without per-iteration allocation or extra locking.
+  std::vector<PollFd> pollfds_;
+
   std::mutex clients_mutex_;
   std::mutex message_mutex_;
   std::mutex listen_fd_mutex_;

src/comm/tcp_server.cpp

@@ -140,6 +140,8 @@ void TCPServer::shutdown()
   }
   // This will effectively deactivate the disconnection handler.
   client_fds_.clear();
+  // Worker thread has been joined above, so the poll set is no longer in use; keep it consistent.
+  pollfds_.clear();
   ur_close(shutdown_socket);
   ur_close(listen_fd_);
   listen_fd_ = INVALID_SOCKET;
@@ -225,6 +227,7 @@ void TCPServer::handleConnect()
     if (client_fds_.size() < max_clients_allowed_ || max_clients_allowed_ == 0)
     {
       client_fds_.push_back(client_fd);
+      rebuildPollfds();
       accepted = true;
     }
     else
@@ -244,28 +247,32 @@ void TCPServer::handleConnect()
   }
 }
 
-void TCPServer::spin()
+void TCPServer::rebuildPollfds()
 {
-  // Build the poll set fresh each iteration from the listen socket plus all currently connected
-  // clients. poll() is used on both platforms (WSAPoll() on Windows) because it has no
-  // FD_SETSIZE limit on file descriptor numbers, unlike select(). This matters when the hosting
-  // process holds many file descriptors (e.g. a JVM), pushing socket FDs past FD_SETSIZE (1024).
-  std::vector<struct pollfd> pollfds;
-  pollfds.push_back({ static_cast<socket_t>(listen_fd_), POLLIN, 0 });
+  // clear() + push_back() reuses the vector's existing capacity, so this does not allocate once
+  // the set has reached a given size.
+  pollfds_.clear();
+  pollfds_.push_back({ static_cast<socket_t>(listen_fd_), POLLIN, 0 });
+  for (const auto& client_fd : client_fds_)
   {
-    std::lock_guard<std::mutex> lk(clients_mutex_);
-    for (const auto& client_fd : client_fds_)
-    {
-      pollfds.push_back({ client_fd, POLLIN, 0 });
-    }
+    pollfds_.push_back({ client_fd, POLLIN, 0 });
   }
+}
 
+void TCPServer::spin()
+{
+  // Poll the cached set (pollfds_) directly. It is kept in sync with client_fds_ via
+  // rebuildPollfds() on every connect/disconnect, so no allocation happens here in steady state.
+  // poll() is used on both platforms (WSAPoll() on Windows) because it has no FD_SETSIZE limit on
+  // file descriptor numbers, unlike select(). This matters when the hosting process holds many
+  // file descriptors (e.g. a JVM), pushing socket FDs past FD_SETSIZE (1024).
+  //
   // Block for up to 1 s waiting for activity on any socket. A shutdown wakes this immediately by
   // connecting to the listen socket (see shutdown()).
 #ifdef _WIN32
-  int ready = ::WSAPoll(pollfds.data(), static_cast<ULONG>(pollfds.size()), 1000);
+  int ready = ::WSAPoll(pollfds_.data(), static_cast<ULONG>(pollfds_.size()), 1000);
 #else
-  int ready = ::poll(pollfds.data(), pollfds.size(), 1000);
+  int ready = ::poll(pollfds_.data(), pollfds_.size(), 1000);
 #endif
   if (ready < 0)
   {
@@ -279,24 +286,29 @@ void TCPServer::spin()
     return;
   }
 
-  if (pollfds[0].revents & POLLIN)
-  {
-    URCL_LOG_DEBUG("Activity on listen FD %d", (int)listen_fd_);
-    handleConnect();
-  }
+  // Snapshot the poll results into locals BEFORE handleConnect()/handleDisconnect() rebuild
+  // pollfds_, otherwise activity on an existing client that arrives in the same poll cycle as a
+  // new connection would be lost.
+  const bool listen_activity = (pollfds_[0].revents & POLLIN) != 0;
 
   std::vector<socket_t> disconnected_clients;
   std::vector<socket_t> client_fds_with_activity;
 
-  // pollfds[0] is the listen socket; client entries start at index 1.
-  for (size_t i = 1; i < pollfds.size(); ++i)
+  // pollfds_[0] is the listen socket; client entries start at index 1.
+  for (size_t i = 1; i < pollfds_.size(); ++i)
   {
-    if (pollfds[i].revents & (POLLIN | POLLHUP | POLLERR))
+    if (pollfds_[i].revents & (POLLIN | POLLHUP | POLLERR))
     {
-      URCL_LOG_DEBUG("Activity on client FD %d", (int)pollfds[i].fd);
-      client_fds_with_activity.push_back(static_cast<socket_t>(pollfds[i].fd));
+      URCL_LOG_DEBUG("Activity on client FD %d", (int)pollfds_[i].fd);
+      client_fds_with_activity.push_back(static_cast<socket_t>(pollfds_[i].fd));
     }
   }
+
+  if (listen_activity)
+  {
+    URCL_LOG_DEBUG("Activity on listen FD %d", (int)listen_fd_);
+    handleConnect();
+  }
   // We handle client activity outside the clients_mutex_ lock to avoid holding it during potentially slow I/O and
   // message callbacks.
   // The clients_mutex_ lock is only needed to protect the client_fds_ vector, but once we have copied the FDs with
@@ -329,6 +341,7 @@ void TCPServer::handleDisconnect(const socket_t fd)
         break;
       }
     }
+    rebuildPollfds();
   }
 
   {
@@ -394,6 +407,16 @@ void TCPServer::start()
 {
   URCL_LOG_DEBUG("Starting worker thread");
   keep_running_ = true;
+  // Seed the poll set with the listen socket before the worker thread starts polling it. Reserve
+  // room for the listen socket plus the expected number of clients up front so the first
+  // connections don't reallocate. A bounded server reserves its exact client limit; an unbounded
+  // one (max_clients_allowed_ == 0) reserves DEFAULT_RESERVED_CLIENTS as headroom.
+  {
+    std::lock_guard<std::mutex> lk(clients_mutex_);
+    const uint32_t expected_clients = max_clients_allowed_ > 0 ? max_clients_allowed_ : DEFAULT_RESERVED_CLIENTS;
+    pollfds_.reserve(static_cast<size_t>(expected_clients) + 1);
+    rebuildPollfds();
+  }
   worker_thread_ = std::thread(&TCPServer::worker, this);
 }
 

tests/test_tcp_server.cpp

@@ -631,6 +631,58 @@ TEST_F(TCPServerTest, services_client_with_high_fd_number)
 }
 #endif
 
+// White-box regression test: the cached poll set must stay in sync with the connected client set
+// as clients connect and disconnect. rebuildPollfds() runs before the connect/disconnect callbacks
+// fire, so once waitForConnectionCallback()/waitForDisconnectionCallback() returns the poll set is
+// already rebuilt and the invariant can be checked deterministically.
+TEST_F(TCPServerTest, poll_set_tracks_client_set)
+{
+  TestableTcpServer server(port_);
+  server.start();
+  EXPECT_EQ(server.getPollSetSize(), 1u);  // listen socket only
+
+  std::vector<std::unique_ptr<Client>> clients;
+  for (int i = 0; i < 5; ++i)
+  {
+    clients.push_back(std::make_unique<Client>(port_));
+    ASSERT_TRUE(server.waitForConnectionCallback());
+    EXPECT_EQ(server.getPollSetSize(), server.getClientFDs().size() + 1);
+  }
+
+  while (!clients.empty())
+  {
+    clients.back()->close();
+    clients.pop_back();
+    ASSERT_TRUE(server.waitForDisconnectionCallback());
+    EXPECT_EQ(server.getPollSetSize(), server.getClientFDs().size() + 1);
+  }
+  EXPECT_EQ(server.getPollSetSize(), 1u);  // back to listen socket only
+}
+
+// Regression test for the snapshot-before-mutate ordering in spin(): a new client connecting must
+// not cause an already-connected client's message that arrives in the same poll cycle to be
+// dropped. If the revents were read after handleConnect() rebuilt the poll set, the existing
+// client's activity would be lost and waitForMessageCallback() would time out.
+TEST_F(TCPServerTest, message_not_lost_when_client_connects_concurrently)
+{
+  TestableTcpServer server(0);
+  server.start();
+  Client existing(server.getPort());
+  ASSERT_TRUE(server.waitForConnectionCallback());
+
+  for (int i = 0; i < 50; ++i)
+  {
+    std::unique_ptr<Client> fresh;
+    std::thread connector([&]() { fresh = std::make_unique<Client>(server.getPort()); });
+    existing.send("ping\n");
+    ASSERT_TRUE(server.waitForMessageCallback(2000));
+    connector.join();
+    ASSERT_TRUE(server.waitForConnectionCallback(2000));
+    fresh->close();
+    ASSERT_TRUE(server.waitForDisconnectionCallback(2000));
+  }
+}
+
 int main(int argc, char* argv[])
 {
   ::testing::InitGoogleTest(&argc, argv);

tests/test_utils.h

@@ -165,6 +165,11 @@ class TestableTcpServer : public urcl::comm::TCPServer
     return client_fds_;
   }
 
+  size_t getPollSetSize()
+  {
+    return TCPServer::getPollSetSize();
+  }
+
 private:
   std::vector<socket_t> client_fds_;
   std::condition_variable connect_cv_;