Files
Jungfraujoch/image_pusher/TCPStreamPusher.cpp
T
leonarski_fandClaude Opus 5 41983ca589 Broker: refuse a run the writer cannot write before arming the detector
A writer refusing to overwrite an existing file was reported badly, and the
earlier fix only covered half of it. The commit message of a6a715703 states the
assumption that broke: "SendStartMessage precedes every std::async in the
receiver constructor". That is true of JFJochReceiverFPGA, where the start
message goes out from the constructor body, and false of JFJochReceiverLite,
where it goes out from MeasurementThread - launched by std::async from the
constructor, which waits only on data_analysis_started. On the DECTRIS path the
writer's refusal therefore cannot reach services.Start(); it arrives from
services.Stop(), i.e. the second try-block of the state machine's
MeasurementThread, which never touched start_exception. /wait_till_done answered
200 for a run that never wrote anything.

Two changes, in front of the problem and behind it.

In front: a PREFLIGHT frame (TCP protocol version 3 -> 4), sent to every writer
before the receiver is built and before the detector is armed. It carries the
same CBOR start message a START does and asks one question: could this run be
written? The writer runs FileWriter::Preflight - the first three statements of
the FileWriter constructor, stopping short of the only one that creates
anything - and answers with an ACK. Shared code, so a pre-flight that passes
cannot be contradicted by the start behind it. A refusal is not fatal: nothing
was started, the connection stays usable, and above all no series is left armed
waiting for a trigger that will not come. Both ends already reject a frame of
another protocol version, so the bump costs no compatibility.

The check now covers the data files, not only the master. A stale
_data_000001.h5 used to sail through and fail when the file was renamed into
place at the end of the run, after the whole dataset had been collected. The
writer holding the master file checks for all of them - they share a directory
by construction, the master linking them by relative name - so the siblings do
not race each other statting files the others are creating.

Behind: /wait_till_done maps an Idle state to 200 or 500 by the severity of the
end-of-run message, so a collection that lost packets, or whose writer could not
write, is no longer indistinguishable from a good one to a script. A cancelled
collection is not affected - MeasurementThread checks the cancel before it looks
at the completeness of the data, and reports it with Info severity - and a
warning is still a 200. start_exception is untouched.

The PREFLIGHT payload is built from FillMessage alone, which is what keeps it
small: the per-pixel arrays are added by SendStartMessage and are not in it.
Measured on a JUNGFRAU 9M, 1.4 kB against 543 kB for the start message. It is
deliberately the whole of FillMessage rather than the fields the check reads
today - a hand-picked subset would answer for a different set of file names the
day the check learns to read one more field, and would do it silently.

Three defects found on the way:

TCPImagePuller republished a PREFLIGHT as a real start message, announcing to
downstream consumers a run that was not going to happen.

TCPStreamPusher::run_number was a plain uint64_t written by the control thread
and read by every ACK thread; Preflight adds a second writer to it.

StreamWriter::CollectImages reprocessed the last frame whenever WaitForImage
returned false - run is only tested at the top of the loop - so an abort wrote
the last image a second time. Pre-existing; the new branch made it visible.

Verified against a running broker on the local configuration: a colliding master
file and a colliding data file are both refused by /start, /wait_until_running
and /wait_till_done with the writer's own message, and the log shows the receiver
was never started for either; a cancelled collection still answers 200.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016WmryXe8ASbNi632sUMfsa
2026-08-27 22:09:50 +02:00

1244 lines
44 KiB
C++

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "TCPStreamPusher.h"
#include <limits>
#include <poll.h>
#include <cerrno>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/tcp.h>
#include <unistd.h>
#include <chrono>
namespace {
int64_t SteadyNowNs() {
return std::chrono::duration_cast<std::chrono::nanoseconds>(
std::chrono::steady_clock::now().time_since_epoch()).count();
}
}
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);
auto hp = addr.substr(prefix.size());
auto p = hp.find_last_of(':');
if (p == std::string::npos)
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Invalid TCP address: " + addr);
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;
port_i = std::stoi(port_str, &parsed);
if (parsed != port_str.size())
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Invalid TCP port in address: " + addr);
} catch (...) {
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Invalid TCP port in address: " + addr);
}
if (port_i < 1 || port_i > static_cast<int>(std::numeric_limits<uint16_t>::max()))
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "TCP port out of range in address: " + addr);
return {host, static_cast<uint16_t>(port_i)};
}
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)
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "socket(listen) failed");
int one = 1;
setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
sockaddr_in sin{};
sin.sin_family = AF_INET;
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)
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "inet_pton failed for " + host);
if (bind(listen_fd, reinterpret_cast<sockaddr*>(&sin), sizeof(sin)) != 0)
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "bind() failed to " + addr);
if (listen(listen_fd, 64) != 0)
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "listen() failed on " + addr);
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) {
pollfd pfd{};
pfd.fd = listen_fd;
pfd.events = POLLIN;
const int prc = poll(&pfd, 1, static_cast<int>(timeout.count()));
if (prc <= 0)
return -1;
if ((pfd.revents & (POLLERR | POLLHUP | POLLNVAL)) != 0)
return -1;
return accept(listen_fd, nullptr, nullptr);
}
void TCPStreamPusher::CloseFd(std::atomic<int>& fd) {
int old_fd = fd.exchange(-1);
if (old_fd >= 0) {
shutdown(old_fd, SHUT_RDWR);
close(old_fd);
}
}
TCPStreamPusher::TCPStreamPusher(const std::string& addr,
size_t in_max_connections,
std::optional<int32_t> in_send_buffer_size)
: endpoint(addr),
max_connections(in_max_connections),
send_buffer_size(in_send_buffer_size) {
if (endpoint.empty())
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No TCP writer address provided");
if (max_connections == 0)
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Max TCP connections cannot be zero");
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);
logger.Info("TCPStreamPusher listening on " + endpoint + " (max " + std::to_string(max_connections) + " connections)");
}
TCPStreamPusher::~TCPStreamPusher() {
// 1. Stop acceptor + keepalive threads (they take connections_mutex briefly)
acceptor_running = false;
int lfd = listen_fd.exchange(-1);
if (lfd >= 0) {
shutdown(lfd, SHUT_RDWR);
close(lfd);
}
if (acceptor_future.valid())
acceptor_future.get();
if (keepalive_future.valid())
keepalive_future.get();
// 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::shared_ptr<Connection>> local_connections;
{
std::lock_guard lg(connections_mutex);
local_connections = std::move(connections);
connections.clear();
session_connections.clear();
calibration_connection.reset();
}
for (auto& c : local_connections) {
if (c)
TearDownConnection(*c);
}
}
void TCPStreamPusher::TearDownConnection(Connection& c) {
StopDataCollectionThreads(c);
c.connected = false;
c.broken = true;
CloseFd(c.fd);
JoinPersistentAck(c);
}
// The persistent ack reader is joined both here and by the acceptor when it reaps a dead
// connection, so the join takes the same per-connection lock the writer join does.
void TCPStreamPusher::JoinPersistentAck(Connection& c) {
std::lock_guard teardown(c.teardown_mutex);
if (c.persistent_ack_future.valid())
c.persistent_ack_future.get();
}
bool TCPStreamPusher::IsConnectionAlive(const Connection& c) const {
if (c.broken)
return false;
const int local_fd = c.fd.load();
if (local_fd < 0)
return false;
pollfd pfd{};
pfd.fd = local_fd;
pfd.events = POLLOUT;
if (poll(&pfd, 1, 0) < 0)
return false;
if ((pfd.revents & (POLLHUP | POLLNVAL)) != 0)
return false;
int so_error = 0;
socklen_t len = sizeof(so_error);
if (getsockopt(local_fd, SOL_SOCKET, SO_ERROR, &so_error, &len) != 0)
return false;
return so_error == 0;
}
bool TCPStreamPusher::SendAll(Connection& c, const void* buf, size_t len) {
const auto* p = static_cast<const uint8_t*>(buf);
size_t sent = 0;
int64_t last_progress_ns = SteadyNowNs();
while (sent < len) {
const int local_fd = c.fd.load();
if (local_fd < 0 || c.broken)
return false;
const int64_t now_ns = SteadyNowNs();
// Treat backpressure as fatal only if the peer shows NO sign of life. A busy
// writer keeps refreshing last_peer_activity_ns via BUSY heartbeats / ACKs, so
// we wait through long stalls; only a silent peer trips this timeout.
const int64_t silent_ns = now_ns - c.last_peer_activity_ns.load(std::memory_order_relaxed);
if (silent_ns > std::chrono::duration_cast<std::chrono::nanoseconds>(peer_liveness_timeout).count()) {
logger.Warning("No liveness from writer on socket " + std::to_string(c.socket_number) +
" for " + std::to_string(silent_ns / 1000000) + " ms while sending; marking broken");
c.broken = true;
CloseFd(c.fd);
return false;
}
// Hard backpressure cap: a writer may keep heartbeating yet never drain (e.g. a
// permanently wedged filesystem). If no bytes leave the socket for this long, give
// up regardless of heartbeats so the run -- and its finalization -- cannot hang.
if (now_ns - last_progress_ns >
std::chrono::duration_cast<std::chrono::nanoseconds>(max_backpressure_timeout).count()) {
logger.Warning("No send progress to writer on socket " + std::to_string(c.socket_number) +
" for " + std::to_string((now_ns - last_progress_ns) / 1000000) +
" ms despite heartbeats; marking broken");
c.broken = true;
CloseFd(c.fd);
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);
return false;
}
if (prc == 0)
continue;
if ((pfd.revents & (POLLERR | POLLHUP | POLLNVAL)) != 0 && !(pfd.revents & POLLOUT)) {
c.broken = true;
CloseFd(c.fd);
return false;
}
ssize_t rc = ::send(local_fd, p + sent, len - sent, MSG_NOSIGNAL);
if (rc < 0) {
if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK)
continue;
if (errno == EPIPE || errno == ECONNRESET || errno == ENOTCONN) {
c.broken = true;
CloseFd(c.fd);
}
return false;
}
sent += static_cast<size_t>(rc);
last_progress_ns = SteadyNowNs();
}
return true;
}
bool TCPStreamPusher::SendFrame(Connection& c, const uint8_t* data, size_t size, TCPFrameType type, int64_t image_number) {
TcpFrameHeader h{};
h.type = static_cast<uint16_t>(type);
h.payload_size = size;
h.image_number = image_number >= 0 ? static_cast<uint64_t>(image_number) : 0;
h.socket_number = c.socket_number;
h.run_number = run_number;
if (!SendAll(c, &h, sizeof(h)))
return false;
if (size > 0 && !SendAll(c, data, size))
return false;
return true;
}
bool TCPStreamPusher::ReadExact(Connection& c, void* buf, size_t len) {
auto* p = static_cast<uint8_t*>(buf);
size_t got = 0;
while (got < len) {
if (!c.active)
return false;
const int local_fd = c.fd.load();
if (local_fd < 0)
return false;
pollfd pfd{};
pfd.fd = local_fd;
pfd.events = POLLIN;
const int prc = poll(&pfd, 1, 100);
if (prc == 0)
continue;
if (prc < 0) {
if (errno == EINTR)
continue;
return false;
}
if ((pfd.revents & (POLLHUP | POLLNVAL)) != 0 && !(pfd.revents & POLLIN))
return false;
if ((pfd.revents & POLLIN) == 0)
continue;
ssize_t rc = ::recv(local_fd, p + got, len - got, 0);
if (rc == 0)
return false;
if (rc < 0) {
if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK)
continue;
return false;
}
got += static_cast<size_t>(rc);
}
return true;
}
bool TCPStreamPusher::ReadExactPersistent(Connection& c, void* buf, size_t len) {
auto* p = static_cast<uint8_t*>(buf);
size_t got = 0;
while (got < len) {
if (!c.connected || c.broken)
return false;
const int local_fd = c.fd.load();
if (local_fd < 0)
return false;
pollfd pfd{};
pfd.fd = local_fd;
pfd.events = POLLIN;
const int prc = poll(&pfd, 1, 500);
if (prc == 0)
continue;
if (prc < 0) {
if (errno == EINTR)
continue;
return false;
}
// Only bail if the peer hung up / the fd is invalid and there is no readable
// data left to drain.
if ((pfd.revents & (POLLHUP | POLLNVAL)) != 0 && !(pfd.revents & POLLIN))
return false;
if ((pfd.revents & POLLIN) == 0)
continue;
ssize_t rc = ::recv(local_fd, p + got, len - got, 0);
if (rc == 0)
return false; // Peer closed connection
if (rc < 0) {
if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK)
continue;
return false;
}
got += static_cast<size_t>(rc);
}
return true;
}
void TCPStreamPusher::WriterThread(Connection* c) {
// Drain the queue until the end sentinel. We deliberately do NOT stop on c->active
// alone: StopDataCollectionThreads() clears active and then pushes the sentinel, and
// every DATA frame enqueued before that sentinel must still be sent. Otherwise END
// (sent by EndDataCollection after stopping this thread) could overtake images that
// are merely waiting in the queue, and the writer would finalize before receiving them.
while (true) {
auto e = c->queue.GetBlocking();
if (e.end)
break;
if (!e.z)
continue;
if (c->broken) {
e.z->release();
continue;
}
bool ok;
{
std::unique_lock ul(c->send_mutex);
ok = SendFrame(*c,
static_cast<const uint8_t*>(e.z->GetImage()),
e.z->GetImageSize(),
TCPFrameType::DATA,
e.z->GetImageNumber());
}
// Synchronous send: once send() returns, the payload has been copied into the
// kernel socket buffer, so the image-buffer slot can be returned immediately.
e.z->release();
if (!ok) {
c->broken = true;
logger.Error("TCP send failed on socket " + std::to_string(c->socket_number));
}
}
}
void TCPStreamPusher::PersistentAckThread(Connection* c) {
while (c->connected && !c->broken) {
TcpFrameHeader h{};
if (!ReadExactPersistent(*c, &h, sizeof(h))) {
if (c->connected && !c->broken) {
c->broken = true;
logger.Info("Persistent connection lost on socket " + std::to_string(c->socket_number));
}
break;
}
if (h.magic != JFJOCH_TCP_MAGIC || h.version != JFJOCH_TCP_VERSION) {
c->broken = true;
logger.Error("Invalid frame on persistent connection, socket " + std::to_string(c->socket_number));
break;
}
if (h.payload_size > JFJOCH_TCP_MAX_PAYLOAD_SIZE) {
c->broken = true;
logger.Error("Oversized payload on persistent connection, socket " + std::to_string(c->socket_number));
break;
}
// Any well-formed frame from the peer is a sign of life.
c->last_peer_activity_ns.store(SteadyNowNs(), std::memory_order_relaxed);
const auto frame_type = static_cast<TCPFrameType>(h.type);
// Keepalive pong from the writer
if (frame_type == TCPFrameType::KEEPALIVE) {
c->last_keepalive_recv = std::chrono::steady_clock::now();
if (h.payload_size > 0) {
std::vector<uint8_t> discard(h.payload_size);
ReadExactPersistent(*c, discard.data(), discard.size());
}
continue;
}
// Busy/backpressure heartbeat: the writer is alive but stalled. The liveness
// timestamp update above is all we need; just record reported FIFO occupancy.
if (frame_type == TCPFrameType::BUSY) {
c->last_ack_fifo_occupancy.store(h.ack_fifo_occupancy, std::memory_order_relaxed);
if (h.payload_size > 0) {
std::vector<uint8_t> discard(h.payload_size);
ReadExactPersistent(*c, discard.data(), discard.size());
}
continue;
}
if (frame_type != TCPFrameType::ACK) {
c->broken = true;
logger.Error("Unexpected frame type " + std::to_string(h.type) + " on socket " + std::to_string(c->socket_number));
break;
}
c->last_ack_fifo_occupancy.store(h.ack_fifo_occupancy, std::memory_order_relaxed);
// ACK frame — forward to data-collection ack logic
std::string error_text;
if (h.payload_size > 0) {
error_text.resize(h.payload_size);
if (!ReadExactPersistent(*c, error_text.data(), error_text.size())) {
c->broken = true;
break;
}
}
const auto ack_for = static_cast<TCPFrameType>(h.ack_for);
const bool ok = (h.flags & TCP_ACK_FLAG_OK) != 0;
const bool fatal = (h.flags & TCP_ACK_FLAG_FATAL) != 0;
const auto code = static_cast<TCPAckCode>(h.ack_code);
// Validate run number: discard stale ACKs from a previous run on a persistent connection
if (h.run_number != run_number) {
logger.Warning("Discarding ACK with stale run_number " + std::to_string(h.run_number)
+ " (expected " + std::to_string(run_number.load()) + ") on socket "
+ std::to_string(c->socket_number));
continue;
}
{
std::unique_lock ul(c->ack_mutex);
c->last_ack_code = code;
if (!error_text.empty())
c->last_ack_error = error_text;
if (ack_for == TCPFrameType::PREFLIGHT) {
c->preflight_ack_received = true;
c->preflight_ack_ok = ok;
if (!ok && error_text.empty())
c->last_ack_error = "PREFLIGHT rejected";
} else if (ack_for == TCPFrameType::START) {
c->start_ack_received = true;
c->start_ack_ok = ok;
if (!ok && error_text.empty())
c->last_ack_error = "START rejected";
} else if (ack_for == TCPFrameType::END) {
c->end_ack_received = true;
c->end_ack_ok = ok;
if (!ok && error_text.empty())
c->last_ack_error = "END rejected";
} else if (ack_for == TCPFrameType::CANCEL) {
c->cancel_ack_received = true;
c->cancel_ack_ok = ok;
if (!ok && error_text.empty())
c->last_ack_error = "CANCEL rejected";
} else if (ack_for == TCPFrameType::DATA) {
c->data_acked_total.fetch_add(1, std::memory_order_relaxed);
total_data_acked_total.fetch_add(1, std::memory_order_relaxed);
if (ok && !fatal) {
c->data_acked_ok.fetch_add(1, std::memory_order_relaxed);
total_data_acked_ok.fetch_add(1, std::memory_order_relaxed);
} else {
c->data_acked_bad.fetch_add(1, std::memory_order_relaxed);
total_data_acked_bad.fetch_add(1, std::memory_order_relaxed);
c->data_ack_error_reported = true;
if (!error_text.empty()) {
c->data_ack_error_text = error_text;
} else if (c->data_ack_error_text.empty()) {
c->data_ack_error_text = "DATA ACK failed";
}
}
}
}
c->ack_cv.notify_all();
}
// If the connection broke, also wake up anyone waiting for an ACK
c->ack_cv.notify_all();
}
void TCPStreamPusher::AcceptorThread() {
uint32_t next_socket_number = 0;
while (acceptor_running) {
int lfd = listen_fd.load();
if (lfd < 0)
break;
int new_fd = AcceptOne(lfd, std::chrono::milliseconds(500));
if (new_fd < 0)
continue;
std::vector<std::shared_ptr<Connection>> dead;
bool rejected = false;
size_t total = 0;
{
std::lock_guard lg(connections_mutex);
dead = DetachDeadConnections();
if (connections.size() >= max_connections) {
rejected = true;
} else {
SetupNewConnection(new_fd, next_socket_number++);
total = connections.size();
}
}
// Outside the lock: this joins writer threads that may be mid-send.
CloseDeadConnections(dead);
if (rejected) {
logger.Warning("Max connections (" + std::to_string(max_connections) +
") reached, rejecting new connection");
shutdown(new_fd, SHUT_RDWR);
close(new_fd);
continue;
}
logger.Info("Accepted writer connection (socket_number=" + std::to_string(next_socket_number - 1) +
", total=" + std::to_string(total) + ")");
}
}
void TCPStreamPusher::SetupNewConnection(int new_fd, uint32_t socket_number) {
auto c = std::make_shared<Connection>(send_queue_size);
c->socket_number = socket_number;
c->fd.store(new_fd);
int one = 1;
setsockopt(new_fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));
// Enable OS-level TCP keep-alive
setsockopt(new_fd, SOL_SOCKET, SO_KEEPALIVE, &one, sizeof(one));
int idle = 30; // longer than app-level keepalive to avoid interference
int intvl = 10;
int cnt = 3;
setsockopt(new_fd, IPPROTO_TCP, TCP_KEEPIDLE, &idle, sizeof(idle));
setsockopt(new_fd, IPPROTO_TCP, TCP_KEEPINTVL, &intvl, sizeof(intvl));
setsockopt(new_fd, IPPROTO_TCP, TCP_KEEPCNT, &cnt, sizeof(cnt));
if (send_buffer_size)
setsockopt(new_fd, SOL_SOCKET, SO_SNDBUF, &send_buffer_size.value(), sizeof(int32_t));
c->connected = true;
c->broken = false;
auto now = std::chrono::steady_clock::now();
c->last_keepalive_sent = now;
c->last_keepalive_recv = now;
c->last_peer_activity_ns.store(SteadyNowNs(), std::memory_order_relaxed);
auto* raw = c.get();
c->persistent_ack_future = std::async(std::launch::async, &TCPStreamPusher::PersistentAckThread, this, raw);
connections.emplace_back(std::move(c));
}
std::vector<std::shared_ptr<TCPStreamPusher::Connection>> TCPStreamPusher::DetachDeadConnections() {
// Must be called with connections_mutex held. Only unlinks them from the pool - which is quick -
// and hands them back for the caller to close once the mutex is released. Tearing them down here
// would join a writer thread that can be sitting in a send for up to max_backpressure_timeout,
// with connections_mutex held the whole time.
std::vector<std::shared_ptr<Connection>> dead;
auto it = connections.begin();
while (it != connections.end()) {
auto c = *it;
if (c->broken || !c->connected || !IsConnectionAlive(*c)) {
c->connected = false;
c->broken = true;
dead.push_back(c);
it = connections.erase(it);
} else {
++it;
}
}
return dead;
}
void TCPStreamPusher::CloseDeadConnections(const std::vector<std::shared_ptr<Connection>> &dead) {
// Blocking: joins the writer thread, which may be inside a send. Never call with
// connections_mutex held. These are out of the pool, but a control-plane call that copied the
// pool before the erase still holds a shared_ptr to them, so the joins below take the
// connection's own teardown lock rather than assuming exclusive access.
for (const auto &c : dead) {
StopDataCollectionThreads(*c);
CloseFd(c->fd);
JoinPersistentAck(*c);
logger.Info("Removed dead connection (socket_number=" + std::to_string(c->socket_number) + ")");
}
}
void TCPStreamPusher::KeepaliveThread() {
while (acceptor_running) {
// Sleep in small increments so we can exit promptly
for (int i = 0; i < 50 && acceptor_running; ++i)
std::this_thread::sleep_for(std::chrono::milliseconds(100));
if (!acceptor_running)
break;
// During data collection, the data flow itself serves as heartbeat
if (data_collection_active)
continue;
// Copy the pool out and send with the mutex released. SendFrame blocks until the peer
// liveness or backpressure timeout, and holding connections_mutex across that stalls
// SendImage and every /statistics poll behind an idle-time keepalive.
std::vector<std::shared_ptr<Connection>> local_connections;
{
std::lock_guard lg(connections_mutex);
local_connections = connections;
}
for (auto& cptr : local_connections) {
auto& c = *cptr;
if (c.broken || !c.connected)
continue;
std::unique_lock ul(c.send_mutex);
if (!SendFrame(c, nullptr, 0, TCPFrameType::KEEPALIVE, -1)) {
logger.Warning("Keepalive send failed on socket " + std::to_string(c.socket_number));
c.broken = true;
} else {
c.last_keepalive_sent = std::chrono::steady_clock::now();
}
}
std::vector<std::shared_ptr<Connection>> dead;
{
std::lock_guard lg(connections_mutex);
dead = DetachDeadConnections();
}
CloseDeadConnections(dead);
}
}
size_t TCPStreamPusher::GetConnectedWriters() const {
std::lock_guard lg(connections_mutex);
size_t count = 0;
for (const auto& c : connections) {
if (c->connected && !c->broken)
++count;
}
return count;
}
void TCPStreamPusher::StartDataCollectionThreads(Connection& c) {
{
std::unique_lock ul(c.ack_mutex);
c.preflight_ack_received = false;
c.preflight_ack_ok = false;
c.start_ack_received = false;
c.start_ack_ok = false;
c.end_ack_received = false;
c.end_ack_ok = false;
c.cancel_ack_received = false;
c.cancel_ack_ok = false;
c.last_ack_error.clear();
c.last_ack_code = TCPAckCode::None;
c.data_ack_error_reported = false;
c.data_ack_error_text.clear();
}
c.data_acked_ok.store(0, std::memory_order_relaxed);
c.data_acked_bad.store(0, std::memory_order_relaxed);
c.data_acked_total.store(0, std::memory_order_relaxed);
c.active = true;
c.writer_future = std::async(std::launch::async, &TCPStreamPusher::WriterThread, this, &c);
}
void TCPStreamPusher::StopDataCollectionThreads(Connection& c) {
// One thread at a time per connection: the acceptor can be reaping this connection while the
// control plane is starting or ending a run on it, and both end in writer_future.get().
std::lock_guard teardown(c.teardown_mutex);
if (!c.active)
return;
c.active = false;
// 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);
// active==false makes WriterThread exit without draining, so return any
// queued-but-unsent slots to the image buffer pool ourselves.
ImagePusherQueueElement e{};
while (c.queue.Get(e)) {
if (e.z)
e.z->release();
}
(void)c.queue.Put({.end = true});
}
c.ack_cv.notify_all();
if (c.writer_future.valid())
c.writer_future.get();
}
bool TCPStreamPusher::WaitForAck(Connection& c, TCPFrameType ack_for, std::chrono::milliseconds timeout, std::string* error_text) {
std::unique_lock ul(c.ack_mutex);
const bool ok = c.ack_cv.wait_for(ul, timeout, [&] {
if (ack_for == TCPFrameType::PREFLIGHT)
return c.preflight_ack_received || c.broken.load();
if (ack_for == TCPFrameType::START)
return c.start_ack_received || c.broken.load();
if (ack_for == TCPFrameType::END)
return c.end_ack_received || c.broken.load();
if (ack_for == TCPFrameType::CANCEL)
return c.cancel_ack_received || c.broken.load();
return false;
});
if (!ok) {
if (error_text) *error_text = "ACK timeout";
return false;
}
if (c.broken) {
if (error_text) *error_text = c.last_ack_error.empty() ? "Socket broken" : c.last_ack_error;
return false;
}
bool ack_ok = false;
if (ack_for == TCPFrameType::PREFLIGHT) ack_ok = c.preflight_ack_ok;
if (ack_for == TCPFrameType::START) ack_ok = c.start_ack_ok;
if (ack_for == TCPFrameType::END) ack_ok = c.end_ack_ok;
if (ack_for == TCPFrameType::CANCEL) ack_ok = c.cancel_ack_ok;
if (!ack_ok && error_text)
*error_text = c.last_ack_error.empty() ? "ACK rejected" : c.last_ack_error;
return ack_ok;
}
bool TCPStreamPusher::WaitForEndAck(Connection& c, std::chrono::milliseconds liveness_timeout, std::string* error_text) {
std::unique_lock ul(c.ack_mutex);
// After END is sent the writer may still be flushing a backlog to a slow filesystem.
// It keeps proving it is alive via BUSY heartbeats and DATA ACKs (both refresh
// last_peer_activity_ns), so we wait through arbitrarily long stalls and only give up
// if the peer falls completely silent for the whole liveness window.
while (!c.end_ack_received && !c.broken) {
c.ack_cv.wait_for(ul, liveness_timeout, [&] {
return c.end_ack_received || c.broken.load();
});
if (c.end_ack_received || c.broken)
break;
const int64_t silent_ns = SteadyNowNs() - c.last_peer_activity_ns.load(std::memory_order_relaxed);
if (silent_ns > std::chrono::duration_cast<std::chrono::nanoseconds>(liveness_timeout).count()) {
if (error_text) *error_text = "END ACK timeout (writer silent)";
return false;
}
}
if (c.broken) {
if (error_text) *error_text = c.last_ack_error.empty() ? "Socket broken" : c.last_ack_error;
return false;
}
if (!c.end_ack_ok && error_text)
*error_text = c.last_ack_error.empty() ? "END ACK rejected" : c.last_ack_error;
return c.end_ack_ok;
}
void TCPStreamPusher::Preflight(StartMessage& message) {
std::vector<std::shared_ptr<Connection>> local_connections;
{
std::lock_guard lg(connections_mutex);
for (const auto& c : connections) {
if (c->connected && !c->broken)
local_connections.push_back(c);
}
}
if (local_connections.empty())
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"No writers connected to " + endpoint);
// Set before the first frame goes out: PersistentAckThread discards any ACK whose run_number
// does not match, so an ACK for this run would be thrown away if run_number still named the
// previous one. StartDataCollection sets it again to the same value a moment later.
run_number = message.run_number;
for (size_t i = 0; i < local_connections.size(); i++) {
auto& c = *local_connections[i];
// Same assignment StartDataCollection makes, and it has to be the same one: only the
// writer holding the master file checks the output files, so a different index 0 here
// than there would check the run on the wrong writer.
message.socket_number = static_cast<int64_t>(c.socket_number);
message.write_master_file = (i == 0);
{
std::unique_lock ul(c.ack_mutex);
c.preflight_ack_received = false;
c.preflight_ack_ok = false;
c.last_ack_error.clear();
}
serializer.SerializeSequenceStart(message);
{
std::unique_lock ul(c.send_mutex);
if (!SendFrame(c, serialization_buffer.data(), serializer.GetBufferSize(), TCPFrameType::PREFLIGHT, -1))
throw JFJochException(JFJochExceptionCategory::FileWriteError,
"Timeout/failure sending PREFLIGHT on socket "
+ std::to_string(c.socket_number));
}
std::string ack_err;
if (!WaitForAck(c, TCPFrameType::PREFLIGHT, std::chrono::seconds(5), &ack_err))
throw JFJochException(JFJochExceptionCategory::FileWriteError,
"Pre-flight check failed on socket " + std::to_string(c.socket_number)
+ ": " + ack_err);
}
}
void TCPStreamPusher::StartDataCollection(StartMessage& message) {
if (message.images_per_file < 1)
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Images per file cannot be zero or negative");
images_per_file = message.images_per_file;
run_number = message.run_number;
run_name = message.run_name;
transmission_error = false;
total_data_acked_ok.store(0, std::memory_order_relaxed);
total_data_acked_bad.store(0, std::memory_order_relaxed);
total_data_acked_total.store(0, std::memory_order_relaxed);
// Stopping a writer thread joins it, and closing a detached connection joins it too, so both
// happen with connections_mutex released; the mutex is only taken to read and write the pool.
std::vector<std::shared_ptr<Connection>> current;
{
std::lock_guard lg(connections_mutex);
current = connections;
}
for (auto& c : current)
StopDataCollectionThreads(*c);
std::vector<std::shared_ptr<Connection>> dead;
{
std::lock_guard lg(connections_mutex);
dead = DetachDeadConnections();
}
CloseDeadConnections(dead);
std::vector<std::shared_ptr<Connection>> local_connections;
{
std::lock_guard lg(connections_mutex);
if (connections.empty())
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No writers connected to " + endpoint);
session_connections = connections;
calibration_connection = session_connections[0];
local_connections = session_connections;
}
logger.Info("Starting data collection with " + std::to_string(local_connections.size()) + " connected writers");
data_collection_active = true;
for (auto& c : local_connections)
StartDataCollectionThreads(*c);
std::vector<bool> started(local_connections.size(), false);
auto rollback_cancel = [&]() {
for (size_t i = 0; i < local_connections.size(); i++) {
auto& c = *local_connections[i];
if (!started[i] || c.broken)
continue;
std::unique_lock ul(c.send_mutex);
(void)SendFrame(c, nullptr, 0, TCPFrameType::CANCEL, -1);
std::string cancel_ack_err;
(void)WaitForAck(c, TCPFrameType::CANCEL, std::chrono::milliseconds(500), &cancel_ack_err);
}
for (auto& c : local_connections)
StopDataCollectionThreads(*c);
{
std::lock_guard lg(connections_mutex);
session_connections.clear();
calibration_connection.reset();
}
data_collection_active = false;
};
for (size_t i = 0; i < local_connections.size(); i++) {
auto& c = *local_connections[i];
message.socket_number = static_cast<int64_t>(c.socket_number);
message.write_master_file = (i == 0);
serializer.SerializeSequenceStart(message);
{
std::unique_lock ul(c.send_mutex);
if (!SendFrame(c, serialization_buffer.data(), serializer.GetBufferSize(), TCPFrameType::START, -1)) {
rollback_cancel();
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"Timeout/failure sending START on socket " + std::to_string(c.socket_number));
}
}
std::string ack_err;
if (!WaitForAck(c, TCPFrameType::START, std::chrono::seconds(5), &ack_err)) {
rollback_cancel();
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"START ACK failed on socket " + std::to_string(c.socket_number) + ": " + ack_err);
}
started[i] = true;
}
}
bool TCPStreamPusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number) {
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;
std::unique_lock ul(c.send_mutex);
return SendFrame(c, image_data, image_size, TCPFrameType::DATA, image_number);
}
bool 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.
std::shared_ptr<Connection> target;
{
std::lock_guard lg(connections_mutex);
const auto& use = (!session_connections.empty() ? session_connections : connections);
if (use.empty()) {
z.release();
return false;
}
auto idx = static_cast<size_t>((z.GetImageNumber() / images_per_file) % static_cast<int64_t>(use.size()));
target = use[idx];
}
if (!target) {
z.release();
return false;
}
// A full queue means the writer is applying backpressure (e.g. briefly stalled on a
// freshly-created file at startup). Keep waiting as long as it proves it is alive —
// BUSY heartbeats / ACKs refresh last_peer_activity_ns from a thread independent of
// its stalled write path — so a slow-but-healthy writer is held back, not dropped.
// Only a peer that has gone completely silent for the liveness window is declared dead;
// a writer that heartbeats but never drains is caught by SendAll's max_backpressure cap.
while (!target->broken && target->active) {
if (target->queue.PutTimeout(ImagePusherQueueElement{
.z = &z,
.end = false
}, std::chrono::milliseconds(50)))
return true;
const int64_t silent_ns = SteadyNowNs() - target->last_peer_activity_ns.load(std::memory_order_relaxed);
if (silent_ns > std::chrono::duration_cast<std::chrono::nanoseconds>(peer_liveness_timeout).count()) {
logger.Warning("No liveness from writer on socket " + std::to_string(target->socket_number) +
" for " + std::to_string(silent_ns / 1000000) + " ms while enqueuing; marking broken");
target->broken = true;
break;
}
}
z.release();
return false;
}
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);
local_connections = (!session_connections.empty() ? session_connections : connections);
}
for (auto& cptr : local_connections) {
auto& c = *cptr;
// Flush all queued DATA before sending END: stopping the writer thread drains the
// send queue and joins it, so every image is on the wire before END goes out. This
// prevents END from overtaking images still waiting in the queue when the writer
// thread has fallen behind (e.g. on a loaded machine), which would make the remote
// writer finalize early and silently drop the trailing images.
StopDataCollectionThreads(c);
if (c.broken) {
ret = false;
continue;
}
{
std::unique_lock ul(c.send_mutex);
if (!SendFrame(c, serialization_buffer.data(), serializer.GetBufferSize(), TCPFrameType::END, -1)) {
ret = false;
continue;
}
}
std::string ack_err;
if (!WaitForEndAck(c, peer_liveness_timeout, &ack_err))
ret = false;
}
{
std::lock_guard lg(connections_mutex);
session_connections.clear();
calibration_connection.reset();
}
data_collection_active = false;
transmission_error = !ret;
return ret;
}
bool TCPStreamPusher::SendCalibration(const CompressedImage& message) {
std::shared_ptr<Connection> target;
{
std::lock_guard lg(connections_mutex);
if (!data_collection_active) {
logger.Error("Refusing to send TCP calibration: no active run");
return false;
}
if (!calibration_connection) {
logger.Error("Refusing to send TCP calibration: calibration target is not set");
return false;
}
target = calibration_connection;
}
if (!target) {
logger.Error("Refusing to send TCP calibration: calibration target is null");
return false;
}
if (target->broken) {
logger.Error("Refusing to send TCP calibration: target socket {} is marked broken", target->socket_number);
return false;
}
if (!target->connected) {
logger.Error("Refusing to send TCP calibration: target socket {} is not connected", target->socket_number);
return false;
}
if (!target->active) {
logger.Error("Refusing to send TCP calibration: target socket {} is not active in current session", target->socket_number);
return false;
}
serializer.SerializeCalibration(message);
std::unique_lock ul(target->send_mutex);
if (target->broken) {
logger.Error("Refusing to send TCP calibration: target socket {} became broken before send", target->socket_number);
return false;
}
if (!target->connected) {
logger.Error("Refusing to send TCP calibration: target socket {} disconnected before send", target->socket_number);
return false;
}
if (!target->active) {
logger.Error("Refusing to send TCP calibration: target socket {} became inactive before send", target->socket_number);
return false;
}
if (!IsConnectionAlive(*target)) {
logger.Error("Refusing to send TCP calibration: target socket {} is not alive", target->socket_number);
target->broken = true;
return false;
}
const bool ok = SendFrame(*target,
serialization_buffer.data(),
serializer.GetBufferSize(),
TCPFrameType::CALIBRATION,
-1);
if (!ok) {
logger.Error("Failed to send TCP calibration on socket {}", target->socket_number);
target->broken = true;
return false;
}
return true;
}
std::string TCPStreamPusher::Finalize() {
std::string ret;
if (transmission_error)
ret += "Timeout sending images (e.g., writer disabled during data collection);";
std::lock_guard lg(connections_mutex);
for (size_t i = 0; i < connections.size(); i++) {
auto& c = *connections[i];
{
std::unique_lock ul(c.ack_mutex);
if (c.data_ack_error_reported && !c.data_ack_error_text.empty()) {
ret += "Writer " + std::to_string(c.socket_number) + ": " + c.data_ack_error_text + ";";
} else if (!c.last_ack_error.empty()) {
ret += "Writer " + std::to_string(c.socket_number) + ": " + c.last_ack_error + ";";
}
}
}
return ret;
}
std::string TCPStreamPusher::PrintSetup() const {
return "TCPStreamPusher: endpoint=" + endpoint +
" max_connections=" + std::to_string(max_connections) +
" connected=" + std::to_string(GetConnectedWriters());
}
std::optional<uint64_t> TCPStreamPusher::GetImagesWritten() const {
return total_data_acked_ok.load(std::memory_order_relaxed);
}
std::optional<uint64_t> TCPStreamPusher::GetImagesWriteError() const {
return total_data_acked_bad.load(std::memory_order_relaxed);
}
std::vector<int64_t> TCPStreamPusher::GetWriterFifoUtilization() const {
std::vector<int64_t> ret;
std::lock_guard lg(connections_mutex);
ret.reserve(connections.size());
for (const auto& c : connections) {
ret.push_back(c->last_ack_fifo_occupancy.load(std::memory_order_relaxed));
}
return ret;
}