nxdns/src/upstream/health.zig

//! Per-endpoint health and exponential backoff. Pure: timestamps and jitter
//! arrive as parameters, so the pool owns the clock and the RNG and this file
//! is testable without a backend.
//!
//! Only peer faults reach this file. `transport.group` decides that; a local
//! resource error or a cancellation must never be recorded, or a full disk
//! would take every upstream out of service.
//!
//! Out-of-order completions are the normal case, not an edge case: two
//! concurrent exchanges against the same endpoint finish in either order, so
//! `at` moves backwards between calls routinely. Every timestamp field
//! therefore updates through `@max` on `.nanoseconds`, and `backoff_until` is
//! only ever extended. A late-arriving success must not shorten a backoff that
//! a later failure already set.
//!
//! Changes to state that describe the endpoint *now* are therefore
//! conditional, while the accumulated counters are not. Both record functions
//! split the same way:
//!
//!   * Data is data. `total_successes`, `total_failures`, the window and the
//!     `@max` of `last_success_at` / `last_error_at` always update, however old
//!     `at` is.
//!   * `consecutive_failures`, `backoff_until` and `last_error_buf` describe
//!     the present, so a newer recorded outcome overrules a stale call.
//!
//! `recordSuccess` clears `consecutive_failures` and `backoff_until` only when
//! its `at` is the newest outcome the state has seen, that is `at >=
//! max(last_success_at, last_error_at)` before the update. A stale success
//! leaves the backoff of a newer failure alone.
//!
//! `recordFailure` is the mirror image. See the rule at that function.
//!
//! `State` carries no lock. The pool owns the mutex.

const std = @import("std");

pub const Config = struct {
    /// Consecutive peer faults before the endpoint is put in backoff.
    failure_threshold: u8 = 2,
    base_backoff_ms: u32 = 500,
    max_backoff_ms: u32 = 60_000,
};

/// Rolling success-rate window, in samples. Equal to the bit width of
/// `State.window`.
pub const window_len = 32;

/// Bytes kept of an `@errorName`, truncated to fit.
pub const error_name_capacity = 48;

/// The shift is capped so `base_backoff_ms << shift` cannot run away; by then
/// `max_backoff_ms` has clamped the result many doublings ago.
const max_shift = 20;

pub const State = struct {
    consecutive_failures: u32,
    total_successes: u64,
    total_failures: u64,
    last_success_at: ?std.Io.Timestamp,
    last_error_at: ?std.Io.Timestamp,
    last_error_buf: [error_name_capacity]u8,
    /// Length of the `@errorName` held in `last_error_buf`, truncated to fit.
    last_error_len: u8,
    backoff_until: ?std.Io.Timestamp,
    /// Bitset, 1 = success, LSB = most recent.
    window: u32,
    window_filled: u8,

    pub const init: State = .{
        .consecutive_failures = 0,
        .total_successes = 0,
        .total_failures = 0,
        .last_success_at = null,
        .last_error_at = null,
        .last_error_buf = @splat(0),
        .last_error_len = 0,
        .backoff_until = null,
        .window = 0,
        .window_filled = 0,
    };

    pub fn recordSuccess(self: *State, at: std.Io.Timestamp) void {
        if (self.isNewestOutcome(at)) {
            self.consecutive_failures = 0;
            self.backoff_until = null;
        }
        self.total_successes += 1;
        self.push(1);
        self.last_success_at = later(self.last_success_at, at);
    }

    /// True when no recorded outcome is newer than `at`. Both timestamp fields
    /// update through `@max`, so their maximum is the newest outcome the state
    /// has seen and no separate field is needed.
    fn isNewestOutcome(self: *const State, at: std.Io.Timestamp) bool {
        return !newerThan(self.last_success_at, at) and !newerThan(self.last_error_at, at);
    }

    /// `err_name` is `@errorName` of a PeerFault member. `rand` supplies
    /// jitter; the caller owns the RNG so this stays pure and the test is
    /// deterministic.
    ///
    /// A stale failure, that is one whose `at` is older than an outcome already
    /// recorded, is held to the mirror image of the stale-success rule. Three
    /// cases, decided on the state *before* this call updates it:
    ///
    ///   1. A newer success exists. The endpoint answered after this failure,
    ///      so this failure cannot make it "consecutively failing" now: leave
    ///      `consecutive_failures` and `backoff_until` alone. A run of failures
    ///      that a later success ended is over.
    ///   2. No newer success, but a newer failure exists. The run of failures
    ///      is unbroken and this call is part of it, so the count and the
    ///      backoff update as usual; order inside a run does not matter. The
    ///      newer failure already wrote `last_error_buf`, so its text stays.
    ///   3. `at` is the newest outcome. Everything updates.
    ///
    /// `backoff_until` only ever extends through `@max`, so case 2 can lengthen
    /// a backoff but never shorten one.
    pub fn recordFailure(
        self: *State,
        at: std.Io.Timestamp,
        err_name: []const u8,
        cfg: Config,
        rand: u32,
    ) void {
        const newer_success = newerThan(self.last_success_at, at);
        const newer_failure = newerThan(self.last_error_at, at);

        if (!newer_failure) {
            const copied = @min(err_name.len, self.last_error_buf.len);
            @memcpy(self.last_error_buf[0..copied], err_name[0..copied]);
            self.last_error_len = @intCast(copied);
        }

        self.total_failures += 1;
        self.push(0);
        self.last_error_at = later(self.last_error_at, at);

        if (newer_success) return;

        self.consecutive_failures +|= 1;
        if (self.consecutive_failures < cfg.failure_threshold) return;

        const delay_ms = backoffDelayMs(self.consecutive_failures, cfg);
        const half = delay_ms / 2;
        const jittered = half + rand % (half + 1);
        const deadline: std.Io.Timestamp = .{
            .nanoseconds = at.nanoseconds + @as(i96, jittered) * std.time.ns_per_ms,
        };
        self.backoff_until = later(self.backoff_until, deadline);
    }

    pub fn available(self: *const State, now: std.Io.Timestamp) bool {
        const until = self.backoff_until orelse return true;
        return now.nanoseconds > until.nanoseconds;
    }

    /// Over the filled part of the window; 1.0 when the window is empty, so a
    /// fresh endpoint is not reported as failing.
    pub fn successRate(self: *const State) f32 {
        if (self.window_filled == 0) return 1.0;
        const filled: u6 = @intCast(@min(self.window_filled, window_len));
        const mask: u32 = if (filled == window_len)
            std.math.maxInt(u32)
        else
            (@as(u32, 1) << @intCast(filled)) - 1;
        const successes = @popCount(self.window & mask);
        return @as(f32, @floatFromInt(successes)) / @as(f32, @floatFromInt(filled));
    }

    pub fn lastError(self: *const State) []const u8 {
        return self.last_error_buf[0..self.last_error_len];
    }

    fn push(self: *State, bit: u1) void {
        self.window = (self.window << 1) | bit;
        self.window_filled = @min(self.window_filled + 1, window_len);
    }
};

/// True when `existing` is set and strictly newer than `at`. Equal timestamps
/// are not stale: a call at the timestamp of the newest outcome still counts as
/// current.
fn newerThan(existing: ?std.Io.Timestamp, at: std.Io.Timestamp) bool {
    const previous = existing orelse return false;
    return previous.nanoseconds > at.nanoseconds;
}

fn later(existing: ?std.Io.Timestamp, candidate: std.Io.Timestamp) std.Io.Timestamp {
    const previous = existing orelse return candidate;
    return .{ .nanoseconds = @max(previous.nanoseconds, candidate.nanoseconds) };
}

/// `min(max_backoff_ms, base_backoff_ms << shift)` with
/// `shift = min(consecutive_failures - failure_threshold, max_shift)`. The
/// arithmetic runs in u64 so a large `base_backoff_ms` cannot wrap before the
/// clamp applies.
fn backoffDelayMs(consecutive_failures: u32, cfg: Config) u32 {
    const over = consecutive_failures - cfg.failure_threshold;
    const shift: u6 = @intCast(@min(over, max_shift));
    const shifted = @as(u64, cfg.base_backoff_ms) << shift;
    return @intCast(@min(@as(u64, cfg.max_backoff_ms), shifted));
}

const testing = std.testing;

fn ts(nanoseconds: i96) std.Io.Timestamp {
    return .{ .nanoseconds = nanoseconds };
}

fn ms(count: i96) i96 {
    return count * std.time.ns_per_ms;
}

test "a failure below the threshold leaves the endpoint available" {
    const cfg: Config = .{};
    var state: State = .init;
    state.recordFailure(ts(0), "Timeout", cfg, 0);

    try testing.expectEqual(@as(u32, 1), state.consecutive_failures);
    try testing.expectEqual(@as(u64, 1), state.total_failures);
    try testing.expectEqual(@as(?std.Io.Timestamp, null), state.backoff_until);
    try testing.expect(state.available(ts(0)));
}

test "reaching the threshold puts the endpoint in backoff until the deadline" {
    const cfg: Config = .{};
    var state: State = .init;
    state.recordFailure(ts(0), "Timeout", cfg, 0);
    state.recordFailure(ts(0), "Timeout", cfg, 0);

    // Two failures, threshold 2, shift 0: delay 500 ms, jitter 0 => 250 ms.
    const until = state.backoff_until.?;
    try testing.expectEqual(ms(250), until.nanoseconds);
    try testing.expect(!state.available(ts(0)));
    try testing.expect(!state.available(until));
    try testing.expect(state.available(ts(until.nanoseconds + 1)));
}

test "consecutive failures grow the delay and saturate at max_backoff_ms" {
    const cfg: Config = .{};
    var state: State = .init;

    var previous: i96 = -1;
    var i: usize = 0;
    while (i < 40) : (i += 1) {
        state.recordFailure(ts(0), "Timeout", cfg, 0);
        if (state.backoff_until) |until| {
            try testing.expect(until.nanoseconds >= previous);
            previous = until.nanoseconds;
        }
    }

    // Jitter 0 halves the delay, and the delay itself is clamped.
    try testing.expectEqual(ms(cfg.max_backoff_ms / 2), state.backoff_until.?.nanoseconds);
    try testing.expectEqual(@as(u32, 500), backoffDelayMs(2, cfg));
    try testing.expectEqual(@as(u32, 1000), backoffDelayMs(3, cfg));
    try testing.expectEqual(@as(u32, 2000), backoffDelayMs(4, cfg));
    try testing.expectEqual(cfg.max_backoff_ms, backoffDelayMs(40, cfg));
}

test "a success resets the consecutive count, the window and the backoff" {
    const cfg: Config = .{};
    var state: State = .init;
    state.recordFailure(ts(0), "Timeout", cfg, 0);
    state.recordFailure(ts(0), "Timeout", cfg, 0);
    try testing.expect(state.backoff_until != null);

    state.recordSuccess(ts(ms(1)));
    try testing.expectEqual(@as(u32, 0), state.consecutive_failures);
    try testing.expectEqual(@as(?std.Io.Timestamp, null), state.backoff_until);
    try testing.expectEqual(@as(u64, 1), state.total_successes);
    try testing.expectEqual(@as(u32, 1), state.window & 1);
    try testing.expect(state.available(ts(0)));
}

test "jitter stays inside half the delay and the whole delay" {
    const cfg: Config = .{};
    const delay = backoffDelayMs(2, cfg);

    for ([_]u32{ 0, std.math.maxInt(u32), 1, 12345 }) |rand| {
        var state: State = .init;
        state.recordFailure(ts(0), "Timeout", cfg, rand);
        state.recordFailure(ts(0), "Timeout", cfg, rand);
        const offset = state.backoff_until.?.nanoseconds;
        try testing.expect(offset >= ms(delay / 2));
        try testing.expect(offset <= ms(delay));
    }
}

test "an out-of-order success does not move last_success_at backwards" {
    var state: State = .init;
    state.recordSuccess(ts(100));
    state.recordSuccess(ts(50));
    try testing.expectEqual(@as(i96, 100), state.last_success_at.?.nanoseconds);
}

test "an out-of-order failure does not shorten the backoff" {
    const cfg: Config = .{};
    var state: State = .init;
    state.recordFailure(ts(ms(100)), "Timeout", cfg, 0);
    state.recordFailure(ts(ms(100)), "Timeout", cfg, 0);
    const until = state.backoff_until.?.nanoseconds;

    state.recordFailure(ts(ms(50)), "Timeout", cfg, 0);
    try testing.expect(state.backoff_until.?.nanoseconds >= until);
    try testing.expectEqual(@as(i96, ms(100)), state.last_error_at.?.nanoseconds);
}

test "an out-of-order success does not clear the backoff of a newer failure" {
    const cfg: Config = .{};
    var state: State = .init;
    state.recordFailure(ts(ms(100)), "Timeout", cfg, 0);
    state.recordFailure(ts(ms(100)), "Timeout", cfg, 0);
    const until = state.backoff_until.?.nanoseconds;

    state.recordSuccess(ts(ms(50)));
    try testing.expectEqual(@as(i96, until), state.backoff_until.?.nanoseconds);
    try testing.expectEqual(@as(u32, 2), state.consecutive_failures);
    try testing.expectEqual(@as(u64, 1), state.total_successes);
    try testing.expectEqual(@as(u32, 1), state.window & 1);
    try testing.expectEqual(@as(i96, ms(50)), state.last_success_at.?.nanoseconds);
}

test "a stale failure behind a newer success does not raise the consecutive count" {
    const cfg: Config = .{};
    var state: State = .init;
    state.recordFailure(ts(ms(10)), "Timeout", cfg, 0);
    state.recordSuccess(ts(ms(100)));
    state.recordFailure(ts(ms(20)), "Timeout", cfg, 0);
    state.recordFailure(ts(ms(30)), "Timeout", cfg, 0);

    try testing.expectEqual(@as(u32, 0), state.consecutive_failures);
    try testing.expectEqual(@as(?std.Io.Timestamp, null), state.backoff_until);
    try testing.expect(state.available(ts(ms(31))));
}

test "a stale failure behind a newer success still counts into the totals" {
    const cfg: Config = .{};
    var state: State = .init;
    state.recordSuccess(ts(ms(100)));
    state.recordFailure(ts(ms(20)), "Timeout", cfg, 0);

    try testing.expectEqual(@as(u64, 1), state.total_failures);
    try testing.expectEqual(@as(u8, 2), state.window_filled);
    try testing.expectEqual(@as(u32, 0), state.window & 1);
    try testing.expectEqual(@as(i96, ms(20)), state.last_error_at.?.nanoseconds);
    try testing.expectEqualStrings("Timeout", state.lastError());
}

test "a stale failure with no newer success still counts as consecutive" {
    const cfg: Config = .{};
    var state: State = .init;
    state.recordFailure(ts(ms(100)), "Timeout", cfg, 0);
    state.recordFailure(ts(ms(50)), "Timeout", cfg, 0);

    try testing.expectEqual(@as(u32, 2), state.consecutive_failures);
    try testing.expect(state.backoff_until != null);
    try testing.expectEqual(@as(i96, ms(100)), state.last_error_at.?.nanoseconds);
}

test "a stale failure does not overwrite the error of a newer failure" {
    const cfg: Config = .{};
    var state: State = .init;
    state.recordFailure(ts(ms(100)), "Timeout", cfg, 0);
    state.recordFailure(ts(ms(50)), "ConnectFailed", cfg, 0);

    try testing.expectEqualStrings("Timeout", state.lastError());
}

test "the newest failure records its own error and extends the backoff" {
    const cfg: Config = .{};
    var state: State = .init;
    state.recordFailure(ts(ms(50)), "Timeout", cfg, 0);
    state.recordFailure(ts(ms(100)), "ConnectFailed", cfg, 0);

    try testing.expectEqualStrings("ConnectFailed", state.lastError());
    try testing.expectEqual(@as(u32, 2), state.consecutive_failures);
    try testing.expectEqual(ms(100) + ms(250), state.backoff_until.?.nanoseconds);
}

test "the newest success clears the backoff of an older failure" {
    const cfg: Config = .{};
    var state: State = .init;
    state.recordFailure(ts(ms(100)), "Timeout", cfg, 0);
    state.recordFailure(ts(ms(100)), "Timeout", cfg, 0);
    try testing.expect(state.backoff_until != null);

    state.recordSuccess(ts(ms(101)));
    try testing.expectEqual(@as(?std.Io.Timestamp, null), state.backoff_until);
    try testing.expectEqual(@as(u32, 0), state.consecutive_failures);
}

test "a success at the timestamp of the newest failure clears the backoff" {
    const cfg: Config = .{};
    var state: State = .init;
    state.recordFailure(ts(ms(100)), "Timeout", cfg, 0);
    state.recordFailure(ts(ms(100)), "Timeout", cfg, 0);

    state.recordSuccess(ts(ms(100)));
    try testing.expectEqual(@as(?std.Io.Timestamp, null), state.backoff_until);
    try testing.expectEqual(@as(u32, 0), state.consecutive_failures);
}

test "successRate over a half-success window is 0.5" {
    const cfg: Config = .{};
    var state: State = .init;
    try testing.expectEqual(@as(f32, 1.0), state.successRate());

    var i: usize = 0;
    while (i < 4) : (i += 1) {
        state.recordSuccess(ts(0));
        state.recordFailure(ts(0), "Timeout", cfg, 0);
    }
    try testing.expectEqual(@as(u8, 8), state.window_filled);
    try testing.expectEqual(@as(f32, 0.5), state.successRate());
}

test "successRate counts only the filled part of the window" {
    var state: State = .init;
    state.recordSuccess(ts(0));
    try testing.expectEqual(@as(f32, 1.0), state.successRate());

    var i: usize = 0;
    while (i < window_len * 2) : (i += 1) state.recordSuccess(ts(0));
    try testing.expectEqual(@as(u8, window_len), state.window_filled);
    try testing.expectEqual(@as(f32, 1.0), state.successRate());
}

test "lastError returns the last recorded name, truncated not overflowed" {
    const cfg: Config = .{};
    var state: State = .init;
    try testing.expectEqualStrings("", state.lastError());

    state.recordFailure(ts(0), "ConnectFailed", cfg, 0);
    try testing.expectEqualStrings("ConnectFailed", state.lastError());

    state.recordFailure(ts(0), "Timeout", cfg, 0);
    try testing.expectEqualStrings("Timeout", state.lastError());

    const long = "A" ** 200;
    state.recordFailure(ts(0), long, cfg, 0);
    try testing.expectEqual(@as(usize, 48), state.lastError().len);
    try testing.expectEqualStrings(long[0..48], state.lastError());
}