Skip to content

tcp balancer logic #38

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Rustamchuk merged 2 commits into from
Jul 16, 2026
Merged

tcp balancer logic #38

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b524b18
tcp balancer logic
Jul 16, 2026
4781835
tcp balancer logic
Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 69 additions & 31 deletions statshouse.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,7 @@ class TransportUDPBase {
MAX_DATAGRAM_SIZE = 65507, // https://stackoverflow.com/questions/42609561/udp-maximum-packet-size/42610200
MAX_KEYS = 47,
MAX_FULL_KEY_SIZE = 1024, // roughly metric plus all tags
MAX_HOST_TAG_LEN = 128,
};

// Metric name builder. Use by calling transport.metric("metric").tag("tag1").tag("tag2").write_count(1);
Expand All @@ -197,9 +198,6 @@ class TransportUDPBase {
}
// Assigns arbitrary tag value by key name
MetricBuilder & tag(string_view key, string_view str) {
if (key.size() == 2 && key.data()[0] == '_' && key.data()[1] == 'h') {
host_set = true;
}
auto begin = buffer + buffer_pos;
auto end = buffer + MAX_FULL_KEY_SIZE;
if (STATSHOUSE_UNLIKELY(!(begin = pack_string(begin, end, key)))) { did_not_fit = true; return *this; }
Expand Down Expand Up @@ -262,7 +260,6 @@ class TransportUDPBase {
}
TransportUDPBase *transport;
bool env_set = false; // so current default_env will be added in pack_header even if set later.
bool host_set = false; // if explicit _h is set, do not add default hostname tag.
bool did_not_fit = false; // we do not want any exceptions in writing metrics
size_t next_tag = 1;
size_t tags_count = 0;
Expand All @@ -276,16 +273,10 @@ class TransportUDPBase {
if (::gethostname(host_buf, sizeof(host_buf) - 1) == 0) {
host_buf[sizeof(host_buf) - 1] = '\0';
host_tag.assign(host_buf);
if (host_tag.size() > MAX_HOST_TAG_LEN) {
host_tag.resize(MAX_HOST_TAG_LEN);
}
}
if (host_tag.empty()) {
return;
}
host_tag_tl_pair.resize(4 + host_tag.size() + 4); // key name, host, padding
auto begin = &host_tag_tl_pair[0];
auto end = begin + host_tag_tl_pair.size();
begin = pack32(begin, end, 2 | (uint32_t('_') << 8) | (uint32_t('h') << 16));
begin = pack_string(begin, end, host_tag);
host_tag_tl_pair.resize(begin - &host_tag_tl_pair[0]);
}
virtual ~TransportUDPBase() = default;

Expand Down Expand Up @@ -401,6 +392,7 @@ class TransportUDPBase {
protected:
Stats stats; // allow implementations to update stats directly for simplicity
mutable mutex mu;
std::string host_tag;
bool write_count_locked(const MetricBuilder &metric, double count, uint32_t tsUnixSec = 0) {
return write_count_impl(metric, count, tsUnixSec);
}
Expand Down Expand Up @@ -430,8 +422,6 @@ class TransportUDPBase {
std::string default_env;
std::string default_env_tl_pair; // "0", default_env in TL format to optimized writing
std::string app_name;
std::string host_tag;
std::string host_tag_tl_pair; // "_h", hostname in TL format;

size_t max_payload_size = DEFAULT_DATAGRAM_SIZE;
bool immediate_flush = false;
Expand Down Expand Up @@ -616,19 +606,13 @@ class TransportUDPBase {
if (STATSHOUSE_UNLIKELY(!enoughSpace(begin, end, metric.buffer_pos))) { return nullptr; }
std::memcpy(begin, metric.buffer, metric.buffer_pos);
const bool add_env = !metric.env_set && !default_env_tl_pair.empty();
const bool add_host = !metric.host_set && !host_tag_tl_pair.empty();
put32(begin + metric.tags_count_pos(), metric.tags_count + int(add_env) + int(add_host));
put32(begin + metric.tags_count_pos(), metric.tags_count + int(add_env));
begin += metric.buffer_pos;
if (add_env) {
if (STATSHOUSE_UNLIKELY(!enoughSpace(begin, end, default_env_tl_pair.size()))) { return nullptr; }
std::memcpy(begin, default_env_tl_pair.data(), default_env_tl_pair.size());
begin += default_env_tl_pair.size();
}
if (add_host) {
if (STATSHOUSE_UNLIKELY(!enoughSpace(begin, end, host_tag_tl_pair.size()))) { return nullptr; }
std::memcpy(begin, host_tag_tl_pair.data(), host_tag_tl_pair.size());
begin += host_tag_tl_pair.size();
}
if (fields_mask & TL_STATSHOUSE_METRIC_COUNTER_FIELDS_MASK) {
if (STATSHOUSE_UNLIKELY(!(begin = pack64(begin, end, doubleBits(counter))))) { return nullptr;}
}
Expand Down Expand Up @@ -901,7 +885,8 @@ class TransportTCP : public TransportUDPBase {
TransportTCP(const std::string &host, int port, const std::string &resolve_targets = "")
: host_{host}
, resolve_targets_{resolve_targets}
, port_{port} {
, port_{port}
, handshake_(build_tcp_handshake_v2(host_tag)) {
(void)refresh_pools();
primary_worker_.th = std::thread(&TransportTCP::run_worker, this, &primary_worker_);
secondary_worker_.th = std::thread(&TransportTCP::run_worker, this, &secondary_worker_);
Expand Down Expand Up @@ -938,7 +923,10 @@ class TransportTCP : public TransportUDPBase {
}

private:
enum { tcpConnBucketCount = 512 };
enum {
tcpConnBucketCount = 512,
stuckReconDelaySec = 30,
};
struct resolved_addr {
::sockaddr_storage storage{};
socklen_t len = 0;
Expand All @@ -952,6 +940,8 @@ class TransportTCP : public TransportUDPBase {
int sock = -1;
std::thread th;
std::atomic<size_t> would_block_bytes{0};
std::atomic_bool force_reconnect{false};
std::chrono::steady_clock::time_point last_stuck_recon{std::chrono::steady_clock::now()};
};
struct send_result {
bool ok = false;
Expand All @@ -962,6 +952,7 @@ class TransportTCP : public TransportUDPBase {
std::string host_;
std::string resolve_targets_;
int port_ = 0;
std::string handshake_; // precomputed reconnect key (statshousev2 + host)
worker primary_worker_;
worker secondary_worker_;
std::mutex route_mu_;
Expand Down Expand Up @@ -996,6 +987,11 @@ class TransportTCP : public TransportUDPBase {
return flush_packet_result::ok;
}
if (enqueue(*secondary, payload)) {
primary->force_reconnect.store(true, std::memory_order_relaxed);
{
std::lock_guard<std::mutex> lock{primary->mu};
primary->cv.notify_one();
}
std::lock_guard<std::mutex> lock{route_mu_};
std::swap(primary_route_, secondary_route_);
return flush_packet_result::ok;
Expand All @@ -1017,21 +1013,47 @@ class TransportTCP : public TransportUDPBase {
void run_worker(worker *w) {
for (;;) {
std::vector<char> payload;
bool have_payload = false;
bool need_recon = false;
{
std::unique_lock<std::mutex> lock{w->mu};
w->cv.wait(lock, [this, w]() { return stop_.load() || !w->queue.empty(); });
if (w->queue.empty() && stop_.load()) {
w->cv.wait(lock, [this, w]() {
return stop_.load() || !w->queue.empty() || w->force_reconnect.load(std::memory_order_relaxed);
});
if (w->force_reconnect.exchange(false, std::memory_order_relaxed)) {
need_recon = true;
}
if (!w->queue.empty()) {
payload = std::move(w->queue.front());
w->queue.pop_front();
have_payload = true;
} else if (stop_.load()) {
break;
}
payload = std::move(w->queue.front());
w->queue.pop_front();
}
send_payload_with_reconnect(*w, payload);
if (need_recon) {
maybe_stuck_reconnect(*w);
}
if (have_payload) {
send_payload_with_reconnect(*w, payload);
}
}
}

void maybe_stuck_reconnect(worker &w) {
const auto now = std::chrono::steady_clock::now();
if (now < w.last_stuck_recon + std::chrono::seconds(stuckReconDelaySec)) {
return;
}
close_socket(w.sock);
w.last_stuck_recon = now;
}

void send_payload_with_reconnect(worker &w, const std::vector<char> &payload) {
for (;;) {
if (w.force_reconnect.exchange(false, std::memory_order_relaxed)) {
maybe_stuck_reconnect(w);
}
if (w.sock < 0 && !reconnect(w)) {
if (stop_.load()) {
return;
Expand Down Expand Up @@ -1107,8 +1129,7 @@ class TransportTCP : public TransportUDPBase {
set_errno_error(err, "statshouse::TransportTCP connect() failed");
return false;
}
static const char kHeader[] = "statshousev1";
const send_result header_write = send_all(sock, kHeader, sizeof(kHeader) - 1);
const send_result header_write = send_all(sock, handshake_.data(), handshake_.size());
if (!header_write.ok) {
int err = header_write.err;
::close(sock);
Expand Down Expand Up @@ -1241,10 +1262,27 @@ class TransportTCP : public TransportUDPBase {
kb.tag("1", "5"); // lang: cpp
kb.tag("2", "1"); // kind: would block
kb.tag("3", app_name_locked());
if (!host_tag.empty()) {
kb.tag("_h", host_tag);
}
const double lost = static_cast<double>(n);
(void)write_values_locked(kb, &lost, 1, 0, 0);
}

static std::string build_tcp_handshake_v2(const std::string &host_tag) {
std::string buf;
buf.reserve(11 + 1 + 4 + host_tag.size());
buf.append("statshousev");
buf.push_back('2');
const uint32_t n = static_cast<uint32_t>(host_tag.size());
buf.push_back(static_cast<char>(n));
buf.push_back(static_cast<char>(n >> 8));
buf.push_back(static_cast<char>(n >> 16));
buf.push_back(static_cast<char>(n >> 24));
buf.append(host_tag);
return buf;
}

void set_errno_error(int err, const char *msg) {
stats.last_error.clear();
stats.last_error.append(msg);
Expand Down