cobalt/src/server.zig

const std = @import("std");
const protocol = @import("protocol.zig");
const dbmod = @import("db.zig");
const metrics = @import("metrics.zig");
const querymod = @import("query.zig");
const exprmod = @import("expr.zig");
const authmod = @import("auth.zig");
const sessionmod = @import("session.zig");
const ttlmod = @import("ttl.zig");
const collmod = @import("collections.zig");
const collttlmod = @import("collttl.zig");
const rowmod = @import("row.zig");
const txnmod = @import("txn.zig");
const txntablemod = @import("txn_table.zig");
const joinmod = @import("join.zig");
const aggregatemod = @import("aggregate.zig");
const windowmod = @import("window.zig");
const dumpmod = @import("dump.zig");
const pubsubmod = @import("pubsub.zig");

const Allocator = std.mem.Allocator;
const Db = dbmod.Db;

/// Which role a given opcode needs at minimum when auth is enabled. `ping`
/// and `auth` itself are handled before this is ever consulted (see
/// `dispatch`) - every other read op needs `readonly`, every mutation
/// needs `readwrite`, and `user_admin` (remote user management: create,
/// change password, change role, delete, list) needs `admin` - the one
/// opcode reserved for administrative server operations. Bootstrapping the
/// very first admin account still has to happen CLI-only over a local Db
/// handle (see `main.zig`'s `user add` and `serve --auth`'s refusal to
/// start without one), since a fresh server has no admin to authenticate
/// this opcode with in the first place.
///
/// `ttl_op` bundles four actions (see `protocol.TtlAction`) behind one
/// opcode, unlike every other entry here which is a single action - so it
/// is the one case that needs to peek at the payload's first byte (the
/// action tag) rather than deciding purely from `op`. `query` is read-only;
/// `expire`/`persist`/`sweep` all mutate the store and need `readwrite`. A
/// payload too short to carry an action byte falls back to `readwrite` (the
/// stricter requirement) - the actual decode in `handleTtlOp` will reject it
/// as malformed either way, so this only affects which error status a
/// bad frame gets before it is ever decoded.
fn requiredRole(op: protocol.Op, payload: []const u8) authmod.Role {
    return switch (op) {
        .ping, .auth, .auth_token => .readonly,
        .get, .scan, .prefix, .stats, .table_select, .table_join, .table_aggregate, .table_window, .table_explain, .table_export, .metrics => .readonly,
        .set, .del, .batch, .flush, .compact, .table_update, .table_delete, .table_import => .readwrite,
        .user_admin => .admin,
        .ttl_op => if (payload.len > 0 and payload[0] == @intFromEnum(protocol.TtlAction.query))
            .readonly
        else
            .readwrite,
        // Same action-byte peek as `.ttl_op` above, and for the same
        // reason: `coll_ttl_op` bundles the same four actions, just
        // against a whole collection instead of one flat key.
        .coll_ttl_op => if (payload.len > 0 and payload[0] == @intFromEnum(protocol.TtlAction.query))
            .readonly
        else
            .readwrite,
        // `coll_op` bundles thirty collection actions behind one opcode
        // (see `protocol.CollAction`) - same action-byte peek as `ttl_op`/
        // `coll_ttl_op`, but the read/write split is looked up from
        // `CollAction.isReadOnly` instead of a single fixed action value,
        // since unlike the four TTL actions this opcode's actions are not
        // all one enum-comparison away from each other.
        .coll_op => if (payload.len > 0)
            (if ((protocol.CollAction.fromByte(payload[0]) catch null)) |a|
                (if (a.isReadOnly()) authmod.Role.readonly else .readwrite)
            else
                .readwrite)
        else
            .readwrite,
        // `txn_op` is always a write (every entry is a put or a delete),
        // unlike `ttl_op`/`coll_ttl_op`/`coll_op` above it has no read-only
        // action to peek for, so it needs no payload inspection at all.
        .txn_op => .readwrite,
        // Pub/sub never touches `Db` at all (see `pubsub.zig`'s `Hub` doc
        // comment) - all three are gated `readonly` purely so a server with
        // auth enabled still requires SOME authenticated connection before
        // any of them works, matching every other opcode's default
        // posture, rather than because any of them reads storage.
        // `unsubscribe` in particular is always sent from a DIFFERENT
        // connection than the one that subscribed (see
        // `protocol.encodeUnsubscribe`'s doc comment) - it needs no
        // stronger role than `subscribe`/`publish` just because it arrives
        // over its own connection.
        .subscribe, .publish, .unsubscribe => .readonly,
    };
}

const PerConnLimitConfig = struct {
    capacity: f64,
    refillPerSecond: f64,
};

/// A TCP server that exposes a single `Db` over the binary protocol in
/// `protocol.zig`. One accept loop thread hands each connection off to
/// its own detached worker thread; every worker serialises its access to
/// `db` through `mutex` since `Db` itself keeps no internal locking.
pub const Server = struct {
    allocator: Allocator,
    db: *Db,
    mutex: std.Thread.Mutex = .{},
    net: std.net.Server,
    acceptThread: ?std.Thread = null,
    active: std.atomic.Value(usize) = std.atomic.Value(usize).init(0),
    stopping: std.atomic.Value(bool) = std.atomic.Value(bool).init(false),
    limiter: ?metrics.TokenBucket = null,
    perConnLimit: ?PerConnLimitConfig = null,
    requestsTotal: metrics.MetricCounter = .{},
    rateLimitedTotal: metrics.MetricCounter = .{},
    connRateLimitedTotal: metrics.MetricCounter = .{},
    authStore: ?authmod.Store = null,
    unauthorizedTotal: metrics.MetricCounter = .{},
    forbiddenTotal: metrics.MetricCounter = .{},
    sessions: ?*sessionmod.SessionStore = null,
    sessionTtlNanos: i128 = 0,
    sessionExpiredTotal: metrics.MetricCounter = .{},
    sweepIntervalNanos: u64 = 0,
    sweeperThread: ?std.Thread = null,
    sessionsSweptTotal: metrics.MetricCounter = .{},
    ttlSweepIntervalNanos: u64 = 0,
    ttlSweeperThread: ?std.Thread = null,
    ttlSweptTotal: metrics.MetricCounter = .{},
    collTtlSweepIntervalNanos: u64 = 0,
    collTtlSweeperThread: ?std.Thread = null,
    collTtlSweptTotal: metrics.MetricCounter = .{},
    /// Channel registry for `subscribe`/`publish` - see `pubsub.zig`'s doc
    /// comment on `Hub` for why this is entirely separate from `mutex`/`db`
    /// above (a live subscriber must never be able to block a store read or
    /// write, or vice versa).
    pubsub: pubsubmod.Hub = undefined,
    messagesPublishedTotal: metrics.MetricCounter = .{},

    pub fn init(allocator: Allocator, db: *Db, address: std.net.Address) !Server {
        const net = try address.listen(.{ .reuse_address = true });
        return .{ .allocator = allocator, .db = db, .net = net, .pubsub = pubsubmod.Hub.init(allocator) };
    }

    pub fn boundAddress(self: *Server) std.net.Address {
        return self.net.listen_address;
    }

    /// Number of connections currently being served. Exposed for tests
    /// that need to wait for a handler to finish without racing on
    /// thread teardown.
    pub fn activeConnections(self: *Server) usize {
        return self.active.load(.acquire);
    }

    /// Turns on a token-bucket rate limit shared across every connection:
    /// each dispatched request consumes one token before it is allowed to
    /// touch `db.mutex` at all, so an overload is rejected cheaply instead
    /// of queuing behind real storage work. Disabled (unlimited) by
    /// default; call once before `start()`.
    pub fn enableRateLimit(self: *Server, capacity: f64, refillPerSecond: f64) void {
        self.limiter = metrics.TokenBucket.init(capacity, refillPerSecond, @intCast(std.time.nanoTimestamp()));
    }

    /// Gives every accepted connection its own private token bucket,
    /// independent of every other connection's budget and of the shared
    /// `limiter` above. A fresh bucket is created per connection in
    /// `handleConnection` (not shared state, so no locking is needed beyond
    /// what TokenBucket already does internally) - one noisy client cannot
    /// starve another client's share of the server, which a single global
    /// bucket cannot express. Both limits can be enabled together: a
    /// request must clear the per-connection budget before it is even
    /// checked against the global one. Disabled by default; call once
    /// before `start()`.
    pub fn enablePerConnectionRateLimit(self: *Server, capacity: f64, refillPerSecond: f64) void {
        self.perConnLimit = .{ .capacity = capacity, .refillPerSecond = refillPerSecond };
    }

    /// Turns on username/password authentication and role enforcement for
    /// every connection: after this call, every opcode except `ping` and
    /// `auth` itself requires a prior successful `auth` on that connection
    /// (`unauthorized` otherwise), and the authenticated role must meet
    /// `requiredRole(op)` (`forbidden` otherwise). `store` is not owned by
    /// the server - the caller keeps managing the underlying `Db` user
    /// records (creating the first admin, etc.) through `store` directly.
    /// Disabled (no auth required, matching every prior version of this
    /// server) by default; call once before `start()`.
    pub fn enableAuth(self: *Server, store: authmod.Store) void {
        self.authStore = store;
    }

    /// Layers expiring session tokens on top of password auth (requires
    /// `enableAuth` to also be called - sessions are only ever minted from
    /// a successful password check, never a substitute for one). Once
    /// enabled: a successful `auth` response carries a fresh token in its
    /// payload (empty payload otherwise, unchanged from before this
    /// existed) valid for `ttlNanos`; a NEW connection can skip the
    /// password entirely by sending that token as `auth_token` instead;
    /// and - the actual behavior change from plain password auth - a
    /// connection's authenticated role stops being honored the moment its
    /// session's expiry passes, even if the TCP connection itself has
    /// stayed open the whole time (`dispatch` re-checks the wall clock on
    /// every request, not just at `auth` time). `store` is not owned by
    /// the server, exactly like `enableAuth`'s `authmod.Store` - the
    /// caller keeps it alive and can inspect/revoke sessions directly.
    pub fn enableSessions(self: *Server, store: *sessionmod.SessionStore, ttlNanos: i128) void {
        self.sessions = store;
        self.sessionTtlNanos = ttlNanos;
    }

    /// Runs `SessionStore.sweepExpired` on a background thread every
    /// `intervalNanos`, starting once `start()` is called. Without this, a
    /// session nobody ever presents again (the client that held it simply
    /// vanished) sits in `sessions` forever - `validate()` only ever cleans
    /// up a token someone actually tries to use. Requires `enableSessions`
    /// to have been called first (checked in `start()`, not here, since
    /// callers are free to call these two in either order before `start`).
    pub fn enableSessionSweep(self: *Server, intervalNanos: u64) void {
        self.sweepIntervalNanos = intervalNanos;
    }

    /// Runs `ttl.sweepExpired` on a background thread every `intervalNanos`,
    /// starting once `start()` is called - independent of `enableAuth`/
    /// `enableSessions`/`enableSessionSweep` above, since key TTLs (set via
    /// the `ttl_op` opcode or the CLI's local `put --ttl`/`expire`) are a
    /// property of the `Db` itself, not of authentication. Without this, a
    /// key whose TTL elapsed but that nobody ever reads again (so
    /// `ttl.getLive`'s lazy expiry never fires) sits in the store forever -
    /// exactly the same gap `enableSessionSweep` closes for session tokens.
    /// Call once before `start()`.
    pub fn enableTtlSweep(self: *Server, intervalNanos: u64) void {
        self.ttlSweepIntervalNanos = intervalNanos;
    }

    /// Runs `collttl.sweepExpired` on a background thread every
    /// `intervalNanos`, starting once `start()` is called - the whole-
    /// collection counterpart of `enableTtlSweep` above, independent of it
    /// (a server can enable one, both, or neither). Without this, a
    /// collection whose TTL elapsed but that nobody ever touches again (so
    /// `coll_ttl_op`'s lazy `expireIfDue` never fires, matching the CLI's
    /// local `coll` dispatcher) sits in the store forever. Call once before
    /// `start()`.
    pub fn enableCollTtlSweep(self: *Server, intervalNanos: u64) void {
        self.collTtlSweepIntervalNanos = intervalNanos;
    }

    /// Registers this server's request counters into `reg` under fixed
    /// names, so they show up alongside any other Prometheus metrics an
    /// embedding application already exports.
    pub fn registerMetrics(self: *Server, reg: *metrics.Registry) !void {
        try reg.registerCounter("cobalt_server_requests_total", &self.requestsTotal);
        try reg.registerCounter("cobalt_server_rate_limited_total", &self.rateLimitedTotal);
        try reg.registerCounter("cobalt_server_conn_rate_limited_total", &self.connRateLimitedTotal);
        try reg.registerCounter("cobalt_server_unauthorized_total", &self.unauthorizedTotal);
        try reg.registerCounter("cobalt_server_forbidden_total", &self.forbiddenTotal);
        try reg.registerCounter("cobalt_server_session_expired_total", &self.sessionExpiredTotal);
        try reg.registerCounter("cobalt_server_sessions_swept_total", &self.sessionsSweptTotal);
        try reg.registerCounter("cobalt_server_ttl_swept_total", &self.ttlSweptTotal);
        try reg.registerCounter("cobalt_server_coll_ttl_swept_total", &self.collTtlSweptTotal);
        try reg.registerCounter("cobalt_server_messages_published_total", &self.messagesPublishedTotal);
    }

    pub fn start(self: *Server) !void {
        self.acceptThread = try std.Thread.spawn(.{}, acceptLoop, .{self});
        if (self.sessions != null and self.sweepIntervalNanos > 0) {
            self.sweeperThread = try std.Thread.spawn(.{}, sweepLoop, .{self});
        }
        if (self.ttlSweepIntervalNanos > 0) {
            self.ttlSweeperThread = try std.Thread.spawn(.{}, ttlSweepLoop, .{self});
        }
        if (self.collTtlSweepIntervalNanos > 0) {
            self.collTtlSweeperThread = try std.Thread.spawn(.{}, collTtlSweepLoop, .{self});
        }
    }

    /// Blocks until the accept loop exits (only happens via `stop()` or
    /// an unrecoverable accept error). Used by long-running foreground
    /// servers such as the CLI's `serve` command.
    pub fn wait(self: *Server) void {
        if (self.acceptThread) |t| t.join();
    }

    /// Background loop started by `start()` when `enableSessionSweep` set a
    /// nonzero interval and sessions are enabled. Sleeps the full interval
    /// between sweeps (a sweep is pure housekeeping, never load-bearing for
    /// correctness - `validate()` already refuses an expired token on its
    /// own), and checks `stopping` both before and after each sleep so
    /// `stop()` never has to wait more than one interval to join it.
    fn sweepLoop(self: *Server) void {
        const store = self.sessions orelse return;
        while (!self.stopping.load(.acquire)) {
            std.time.sleep(self.sweepIntervalNanos);
            if (self.stopping.load(.acquire)) return;
            const dropped = store.sweepExpired(std.time.nanoTimestamp());
            if (dropped > 0) self.sessionsSweptTotal.incrBy(@intCast(dropped));
        }
    }

    /// Background loop started by `start()` when `enableTtlSweep` set a
    /// nonzero interval. Unlike `sweepLoop` above, `ttl.sweepExpired` reads
    /// and mutates the same `Db` every request handler touches, so this
    /// loop must take `self.mutex` around the call exactly like `dispatch`
    /// does for every opcode - `Db` itself keeps no internal locking (see
    /// the doc comment on `Server`). A failed sweep (e.g. a transient I/O
    /// error) is skipped rather than crashing the server; the next tick
    /// tries again.
    fn ttlSweepLoop(self: *Server) void {
        while (!self.stopping.load(.acquire)) {
            std.time.sleep(self.ttlSweepIntervalNanos);
            if (self.stopping.load(.acquire)) return;
            self.mutex.lock();
            const reaped = ttlmod.sweepExpired(self.db, self.allocator, std.time.nanoTimestamp()) catch {
                self.mutex.unlock();
                continue;
            };
            self.mutex.unlock();
            if (reaped > 0) self.ttlSweptTotal.incrBy(@intCast(reaped));
        }
    }