TCPImagePusher: Further improvement

image_pusher/TCPStreamPusher.cpp

@@ -10,11 +10,12 @@
 #include <arpa/inet.h>
 #include <netinet/tcp.h>
 #include <unistd.h>
+#include <chrono>
 #if defined(MSG_ZEROCOPY)
 #include <linux/errqueue.h>
 #endif
 
-std::pair<std::string, uint16_t> TCPStreamPusher::ParseTcpAddress(const std::string& addr) {
+std::pair<std::string, std::optional<uint16_t>> TCPStreamPusher::ParseTcpAddress(const std::string& addr) {
     const std::string prefix = "tcp://";
     if (addr.rfind(prefix, 0) != 0)
         throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Invalid TCP address: " + addr);
@@ -27,6 +28,9 @@ std::pair<std::string, uint16_t> TCPStreamPusher::ParseTcpAddress(const std::str
     const auto host = hp.substr(0, p);
     const auto port_str = hp.substr(p + 1);
 
+    if (port_str == "*")
+        return {host, std::nullopt};
+
     int port_i = 0;
     try {
         size_t parsed = 0;
@@ -43,8 +47,8 @@ std::pair<std::string, uint16_t> TCPStreamPusher::ParseTcpAddress(const std::str
     return {host, static_cast<uint16_t>(port_i)};
 }
 
-int TCPStreamPusher::OpenListenSocket(const std::string& addr) {
-    auto [host, port] = ParseTcpAddress(addr);
+std::pair<int, std::string> TCPStreamPusher::OpenListenSocket(const std::string& addr) {
+    auto [host, port_opt] = ParseTcpAddress(addr);
 
     int listen_fd = ::socket(AF_INET, SOCK_STREAM, 0);
     if (listen_fd < 0)
@@ -55,7 +59,7 @@ int TCPStreamPusher::OpenListenSocket(const std::string& addr) {
 
     sockaddr_in sin{};
     sin.sin_family = AF_INET;
-    sin.sin_port = htons(port);
+    sin.sin_port = htons(port_opt.has_value() ? port_opt.value() : 0);
     if (host == "*" || host == "0.0.0.0")
         sin.sin_addr.s_addr = htonl(INADDR_ANY);
     else if (inet_pton(AF_INET, host.c_str(), &sin.sin_addr) != 1)
@@ -67,7 +71,16 @@ int TCPStreamPusher::OpenListenSocket(const std::string& addr) {
     if (listen(listen_fd, 64) != 0)
         throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "listen() failed on " + addr);
 
-    return listen_fd;
+    sockaddr_in actual{};
+    socklen_t actual_len = sizeof(actual);
+    if (getsockname(listen_fd, reinterpret_cast<sockaddr*>(&actual), &actual_len) != 0)
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "getsockname() failed on " + addr);
+
+    const uint16_t bound_port = ntohs(actual.sin_port);
+    const std::string normalized_host = (host == "*") ? "0.0.0.0" : host;
+    const std::string bound_endpoint = "tcp://" + normalized_host + ":" + std::to_string(bound_port);
+
+    return {listen_fd, bound_endpoint};
 }
 
 int TCPStreamPusher::AcceptOne(int listen_fd, std::chrono::milliseconds timeout) {
@@ -105,7 +118,10 @@ TCPStreamPusher::TCPStreamPusher(const std::string& addr,
     if (max_connections == 0)
         throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Max TCP connections cannot be zero");
 
-    listen_fd.store(OpenListenSocket(endpoint));
+    auto [lfd, bound_endpoint] = OpenListenSocket(endpoint);
+    listen_fd.store(lfd);
+    endpoint = bound_endpoint;
+
     acceptor_running = true;
     acceptor_future = std::async(std::launch::async, &TCPStreamPusher::AcceptorThread, this);
     keepalive_future = std::async(std::launch::async, &TCPStreamPusher::KeepaliveThread, this);
@@ -128,15 +144,18 @@ TCPStreamPusher::~TCPStreamPusher() {
 
     // 2. Now no background threads touch connections_mutex. Tear down connections.
     //    We do NOT hold the mutex while joining futures, to avoid deadlock.
-    std::vector<std::unique_ptr<Connection>> local_connections;
+    std::vector<std::shared_ptr<Connection>> local_connections;
     {
         std::lock_guard lg(connections_mutex);
         local_connections = std::move(connections);
         connections.clear();
+        session_connections.clear();
     }
 
-    for (auto& c : local_connections)
-        TearDownConnection(*c);
+    for (auto& c : local_connections) {
+        if (c)
+            TearDownConnection(*c);
+    }
 }
 
 void TCPStreamPusher::TearDownConnection(Connection& c) {
@@ -179,6 +198,7 @@ bool TCPStreamPusher::SendAll(Connection& c, const void* buf, size_t len, bool a
     bool zc_used = false;
     uint32_t zc_first = 0;
     uint32_t zc_last = 0;
+    const auto deadline = std::chrono::steady_clock::now() + send_total_timeout;
 
     bool try_zerocopy = false;
 #if defined(MSG_ZEROCOPY)
@@ -194,6 +214,40 @@ bool TCPStreamPusher::SendAll(Connection& c, const void* buf, size_t len, bool a
             return false;
         }
 
+        if (std::chrono::steady_clock::now() >= deadline) {
+            c.broken = true;
+            CloseFd(c.fd);
+            if (zc_used_out) *zc_used_out = zc_used;
+            if (zc_first_out) *zc_first_out = zc_first;
+            if (zc_last_out) *zc_last_out = zc_last;
+            return false;
+        }
+
+        pollfd pfd{};
+        pfd.fd = local_fd;
+        pfd.events = POLLOUT;
+        const int prc = poll(&pfd, 1, static_cast<int>(send_poll_timeout.count()));
+        if (prc < 0) {
+            if (errno == EINTR)
+                continue;
+            c.broken = true;
+            CloseFd(c.fd);
+            if (zc_used_out) *zc_used_out = zc_used;
+            if (zc_first_out) *zc_first_out = zc_first;
+            if (zc_last_out) *zc_last_out = zc_last;
+            return false;
+        }
+        if (prc == 0)
+            continue;
+        if ((pfd.revents & (POLLERR | POLLHUP | POLLNVAL)) != 0 && !(pfd.revents & POLLOUT)) {
+            c.broken = true;
+            CloseFd(c.fd);
+            if (zc_used_out) *zc_used_out = zc_used;
+            if (zc_first_out) *zc_first_out = zc_first;
+            if (zc_last_out) *zc_last_out = zc_last;
+            return false;
+        }
+
         int flags = MSG_NOSIGNAL;
 #if defined(MSG_ZEROCOPY)
         if (try_zerocopy)
@@ -202,7 +256,7 @@ bool TCPStreamPusher::SendAll(Connection& c, const void* buf, size_t len, bool a
 
         ssize_t rc = ::send(local_fd, p + sent, len - sent, flags);
         if (rc < 0) {
-            if (errno == EINTR)
+            if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK)
                 continue;
 
 #if defined(MSG_ZEROCOPY)
@@ -283,6 +337,7 @@ bool TCPStreamPusher::SendFrame(Connection& c, const uint8_t* data, size_t size,
     return true;
 }
 
+// Caller must hold c.zc_mutex.
 void TCPStreamPusher::ReleaseCompletedZeroCopy(Connection& c) {
     while (!c.zc_pending.empty()) {
         const auto& front = c.zc_pending.front();
@@ -630,7 +685,7 @@ void TCPStreamPusher::AcceptorThread() {
 }
 
 void TCPStreamPusher::SetupNewConnection(int new_fd, uint32_t socket_number) {
-    auto c = std::make_unique<Connection>(send_queue_size);
+    auto c = std::make_shared<Connection>(send_queue_size);
     c->socket_number = socket_number;
     c->fd.store(new_fd);
 
@@ -676,23 +731,17 @@ void TCPStreamPusher::RemoveDeadConnections() {
     // doesn't take connections_mutex, so no deadlock.
     auto it = connections.begin();
     while (it != connections.end()) {
-        auto& c = **it;
-        if (c.broken || !c.connected || !IsConnectionAlive(c)) {
-            // Mark dead first
-            c.connected = false;
-            c.broken = true;
+        auto c = *it;
+        if (c->broken || !c->connected || !IsConnectionAlive(*c)) {
+            c->connected = false;
+            c->broken = true;
+            StopDataCollectionThreads(*c);
+            CloseFd(c->fd);
 
-            // Stop data collection threads (they don't take connections_mutex)
-            StopDataCollectionThreads(c);
+            if (c->persistent_ack_future.valid())
+                c->persistent_ack_future.get();
 
-            // Close fd to unblock PersistentAckThread
-            CloseFd(c.fd);
-
-            // Join persistent ack thread — safe because it doesn't take connections_mutex
-            if (c.persistent_ack_future.valid())
-                c.persistent_ack_future.get();
-
-            logger.Info("Removed dead connection (socket_number=" + std::to_string(c.socket_number) + ")");
+            logger.Info("Removed dead connection (socket_number=" + std::to_string(c->socket_number) + ")");
             it = connections.erase(it);
         } else {
             ++it;
@@ -777,7 +826,15 @@ void TCPStreamPusher::StopDataCollectionThreads(Connection& c) {
         return;
 
     c.active = false;
-    c.queue.PutBlocking({.end = true});
+
+    // Avoid potential shutdown deadlock if queue is full and writer is stalled.
+    if (!c.queue.PutTimeout({.end = true}, std::chrono::milliseconds(200))) {
+        c.broken = true;
+        CloseFd(c.fd);
+        c.queue.Clear();
+        (void)c.queue.Put({.end = true});
+    }
+
     c.ack_cv.notify_all();
     c.zc_cv.notify_all();
 
@@ -846,33 +903,30 @@ void TCPStreamPusher::StartDataCollection(StartMessage& message) {
     total_data_acked_bad.store(0, std::memory_order_relaxed);
     total_data_acked_total.store(0, std::memory_order_relaxed);
 
-    // Stop any leftover data-collection threads and clean up dead connections
+    std::vector<std::shared_ptr<Connection>> local_connections;
     {
         std::lock_guard lg(connections_mutex);
         for (auto& c : connections)
             StopDataCollectionThreads(*c);
         RemoveDeadConnections();
+        if (connections.empty())
+            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No writers connected to " + endpoint);
+
+        session_connections = connections;
+        local_connections = session_connections;
     }
 
-    std::lock_guard lg(connections_mutex);
-
-    if (connections.empty())
-        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
-                              "No writers connected to " + endpoint);
-
-    logger.Info("Starting data collection with " + std::to_string(connections.size()) + " connected writers");
-
+    logger.Info("Starting data collection with " + std::to_string(local_connections.size()) + " connected writers");
     data_collection_active = true;
 
-    // Start writer + zerocopy threads for each connection
-    for (auto& c : connections)
+    for (auto& c : local_connections)
         StartDataCollectionThreads(*c);
 
-    std::vector<bool> started(connections.size(), false);
+    std::vector<bool> started(local_connections.size(), false);
 
     auto rollback_cancel = [&]() {
-        for (size_t i = 0; i < connections.size(); i++) {
-            auto& c = *connections[i];
+        for (size_t i = 0; i < local_connections.size(); i++) {
+            auto& c = *local_connections[i];
             if (!started[i] || c.broken)
                 continue;
 
@@ -883,16 +937,20 @@ void TCPStreamPusher::StartDataCollection(StartMessage& message) {
             (void)WaitForAck(c, TCPFrameType::CANCEL, std::chrono::milliseconds(500), &cancel_ack_err);
         }
 
-        for (auto& c : connections)
+        for (auto& c : local_connections)
             StopDataCollectionThreads(*c);
 
+        {
+            std::lock_guard lg(connections_mutex);
+            session_connections.clear();
+        }
         data_collection_active = false;
     };
 
-    for (size_t i = 0; i < connections.size(); i++) {
-        auto& c = *connections[i];
+    for (size_t i = 0; i < local_connections.size(); i++) {
+        auto& c = *local_connections[i];
 
-        message.socket_number = static_cast<int64_t>(i);
+        message.socket_number = static_cast<int64_t>(c.socket_number);
         message.write_master_file = (i == 0);
 
         serializer.SerializeSequenceStart(message);
@@ -918,13 +976,19 @@ void TCPStreamPusher::StartDataCollection(StartMessage& message) {
 }
 
 bool TCPStreamPusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number) {
-    std::lock_guard lg(connections_mutex);
-    if (connections.empty())
-        return false;
-
-    auto idx = static_cast<size_t>((image_number / images_per_file) % static_cast<int64_t>(connections.size()));
-    auto& c = *connections[idx];
+    std::shared_ptr<Connection> target;
+    size_t conn_count = 0;
+    {
+        std::lock_guard lg(connections_mutex);
+        const auto& use = (!session_connections.empty() ? session_connections : connections);
+        if (use.empty())
+            return false;
+        conn_count = use.size();
+        auto idx = static_cast<size_t>((image_number / images_per_file) % static_cast<int64_t>(conn_count));
+        target = use[idx];
+    }
 
+    auto& c = *target;
     if (c.broken || !IsConnectionAlive(c))
         return false;
 
@@ -936,39 +1000,56 @@ void TCPStreamPusher::SendImage(ZeroCopyReturnValue &z) {
     // Look up the target connection while holding the mutex, but do NOT call
     // PutBlocking while holding it — that can block indefinitely and deadlock
     // against AcceptorThread/KeepaliveThread.
-    Connection* target = nullptr;
+    std::shared_ptr<Connection> target;
     {
         std::lock_guard lg(connections_mutex);
-        if (connections.empty()) {
+        const auto& use = (!session_connections.empty() ? session_connections : connections);
+        if (use.empty()) {
             z.release();
             return;
         }
 
-        auto idx = static_cast<size_t>((z.GetImageNumber() / images_per_file) % static_cast<int64_t>(connections.size()));
-        auto& c = *connections[idx];
-
-        if (c.broken) {
-            z.release();
-            return;
-        }
-        target = &c;
+        auto idx = static_cast<size_t>((z.GetImageNumber() / images_per_file) % static_cast<int64_t>(use.size()));
+        target = use[idx];
     }
 
-    target->queue.PutBlocking(ImagePusherQueueElement{
-        .image_data = static_cast<uint8_t *>(z.GetImage()),
-        .z = &z,
-        .end = false
-    });
+    if (!target || target->broken || !target->active) {
+        z.release();
+        return;
+    }
+
+    const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2);
+    while (std::chrono::steady_clock::now() < deadline) {
+        if (target->broken || !target->active) {
+            z.release();
+            return;
+        }
+
+        if (target->queue.PutTimeout(ImagePusherQueueElement{
+                .image_data = static_cast<uint8_t *>(z.GetImage()),
+                .z = &z,
+                .end = false
+            }, std::chrono::milliseconds(50))) {
+            return;
+            }
+    }
+
+    target->broken = true;
+    z.release();
 }
 
 bool TCPStreamPusher::EndDataCollection(const EndMessage& message) {
     serializer.SerializeSequenceEnd(message);
 
     bool ret = true;
+    std::vector<std::shared_ptr<Connection>> local_connections;
 
-    std::lock_guard lg(connections_mutex);
+    {
+        std::lock_guard lg(connections_mutex);
+        local_connections = (!session_connections.empty() ? session_connections : connections);
+    }
 
-    for (auto& cptr : connections) {
+    for (auto& cptr : local_connections) {
         auto& c = *cptr;
         if (c.broken) {
             ret = false;
@@ -988,28 +1069,35 @@ bool TCPStreamPusher::EndDataCollection(const EndMessage& message) {
             ret = false;
     }
 
-    // Stop only data-collection threads, keep connections alive
-    for (auto& c : connections)
+    for (auto& c : local_connections)
         StopDataCollectionThreads(*c);
 
+    {
+        std::lock_guard lg(connections_mutex);
+        session_connections.clear();
+    }
+
     data_collection_active = false;
     transmission_error = !ret;
     return ret;
 }
 
 bool TCPStreamPusher::SendCalibration(const CompressedImage& message) {
-    std::lock_guard lg(connections_mutex);
-    if (connections.empty())
+    std::shared_ptr<Connection> target;
+    {
+        std::lock_guard lg(connections_mutex);
+        if (connections.empty())
+            return false;
+        target = connections[0];
+    }
+
+    if (!target || target->broken)
         return false;
 
     serializer.SerializeCalibration(message);
 
-    auto& c = *connections[0];
-    if (c.broken)
-        return false;
-
-    std::unique_lock ul(c.send_mutex);
-    return SendFrame(c, serialization_buffer.data(), serializer.GetBufferSize(), TCPFrameType::CALIBRATION, -1, nullptr);
+    std::unique_lock ul(target->send_mutex);
+    return SendFrame(*target, serialization_buffer.data(), serializer.GetBufferSize(), TCPFrameType::CALIBRATION, -1, nullptr);
 }
 
 std::string TCPStreamPusher::Finalize() {