From a283b01556fec1721e6f5b6dd108f4c540ace3ce Mon Sep 17 00:00:00 2001 From: DualFroz Date: Thu, 10 Sep 2026 21:54:04 +0200 Subject: [PATCH] feat: add collTtlSweepLoop and acceptLoop to server --- src/server.zig | 156 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 156 insertions(+) diff --git a/src/server.zig b/src/server.zig index 9e677cb..982c7d6 100644 --- a/src/server.zig +++ b/src/server.zig @@ -314,3 +314,159 @@ pub const Server = struct { } } + /// Background loop started by `start()` when `enableCollTtlSweep` set a + /// nonzero interval - the whole-collection counterpart of `ttlSweepLoop` + /// above, same shape (own thread, takes `self.mutex` around the actual + /// `Db` access, skips a failed sweep rather than crashing). + /// `collttl.sweepExpired` walks every kind's marker namespace in one + /// prefix scan, so one thread covers Counter/Hash/Set/List/ZSet at once + /// rather than needing one sweeper per kind. + fn collTtlSweepLoop(self: *Server) void { + while (!self.stopping.load(.acquire)) { + std.time.sleep(self.collTtlSweepIntervalNanos); + if (self.stopping.load(.acquire)) return; + self.mutex.lock(); + const reaped = collttlmod.sweepExpired(self.db, self.allocator, std.time.nanoTimestamp()) catch { + self.mutex.unlock(); + continue; + }; + self.mutex.unlock(); + if (reaped > 0) self.collTtlSweptTotal.incrBy(@intCast(reaped)); + } + } + + fn acceptLoop(self: *Server) void { + while (true) { + const conn = self.net.accept() catch return; + if (self.stopping.load(.acquire)) { + conn.stream.close(); + return; + } + _ = self.active.fetchAdd(1, .acq_rel); + const t = std.Thread.spawn(.{}, handleConnection, .{ self, conn }) catch { + _ = self.active.fetchSub(1, .acq_rel); + conn.stream.close(); + continue; + }; + t.detach(); + } + } + + fn handleConnection(self: *Server, conn: std.net.Server.Connection) void { + defer conn.stream.close(); + defer _ = self.active.fetchSub(1, .acq_rel); + var stream = conn.stream; + + // A private bucket for this one connection, if enabled. It lives on + // this worker's own stack and is only ever touched by this thread, + // so unlike `self.limiter` it needs no cross-thread synchronization + // even though TokenBucket carries its own mutex regardless. + var connLimiter: ?metrics.TokenBucket = if (self.perConnLimit) |cfg| + metrics.TokenBucket.init(cfg.capacity, cfg.refillPerSecond, @intCast(std.time.nanoTimestamp())) + else + null; + + // Which role (if any) this one connection has authenticated as. + // Lives on this worker's stack, exactly like `connLimiter` above - + // a fresh, private slot per connection, never shared. + var connRole: ?authmod.Role = null; + + // When sessions are enabled (see `enableSessions`), the nanosecond + // timestamp at which `connRole` above stops being honored - set + // whenever `auth`/`auth_token` succeeds, re-checked on every + // request in `dispatch`. Stays null (never expires on its own) for + // a server with auth but no session store, matching the original + // "authenticated for as long as the connection is open" behavior. + var connExpiresAtNanos: ?i128 = null; + + while (true) { + var frame = protocol.readFrame(self.allocator, stream.reader()) catch return; + defer frame.deinit(); + const limPtr: ?*metrics.TokenBucket = if (connLimiter) |*cl| cl else null; + self.dispatch(&frame, stream, stream.writer(), limPtr, &connRole, &connExpiresAtNanos) catch return; + } + } + + fn writeErr(writer: anytype, status: protocol.Status, message: []const u8) !void { + try protocol.writeFrame(writer, @intFromEnum(status), message); + } + + fn dispatch(self: *Server, frame: *protocol.Frame, stream: std.net.Stream, writer: anytype, connLimiter: ?*metrics.TokenBucket, connRole: *?authmod.Role, connExpiresAtNanos: *?i128) !void { + self.requestsTotal.incr(); + if (connLimiter) |cl| { + if (!cl.tryAcquireNow(1)) { + self.connRateLimitedTotal.incr(); + return writeErr(writer, .rate_limited, "per-connection rate limit exceeded"); + } + } + if (self.limiter) |*lim| { + if (!lim.tryAcquireNow(1)) { + self.rateLimitedTotal.incr(); + return writeErr(writer, .rate_limited, "rate limit exceeded"); + } + } + + const op = protocol.Op.fromByte(frame.tag) catch { + return writeErr(writer, .bad_request, "unknown opcode"); + }; + + if (self.authStore) |store| { + if (op == .auth) return self.handleAuth(store, frame.payload, writer, connRole, connExpiresAtNanos); + if (op == .auth_token) return self.handleAuthToken(frame.payload, writer, connRole, connExpiresAtNanos); + if (op != .ping) { + const role = connRole.* orelse { + self.unauthorizedTotal.incr(); + return writeErr(writer, .unauthorized, "authentication required"); + }; + // A role survives only as long as its session does, when + // sessions are enabled: a session that expired since the + // last request on this same connection is treated exactly + // like never having authenticated at all, forcing a fresh + // `auth`/`auth_token` before anything else proceeds. + if (connExpiresAtNanos.*) |expiresAt| { + if (std.time.nanoTimestamp() >= expiresAt) { + connRole.* = null; + connExpiresAtNanos.* = null; + self.sessionExpiredTotal.incr(); + self.unauthorizedTotal.incr(); + return writeErr(writer, .unauthorized, "session expired"); + } + } + if (!role.atLeast(requiredRole(op, frame.payload))) { + self.forbiddenTotal.incr(); + return writeErr(writer, .forbidden, "insufficient role for this operation"); + } + } + } + + // `subscribe`/`publish` never touch `db` - handled here, before + // `self.mutex` is ever taken, so neither one can be blocked behind + // (or block) a store read/write. `subscribe` in particular can hold + // this stack frame open for the entire remaining lifetime of the + // connection (see `handleSubscribe`'s doc comment), which would + // starve every other connection forever if it happened under + // `self.mutex`. + if (op == .subscribe) return self.handleSubscribe(frame.payload, stream); + if (op == .publish) return self.handlePublish(frame.payload, writer); + if (op == .unsubscribe) return self.handleUnsubscribe(frame.payload, writer); + + self.mutex.lock(); + defer self.mutex.unlock(); + + switch (op) { + .ping => try protocol.writeFrame(writer, @intFromEnum(protocol.Status.ok), "PONG"), + + // Only reachable when `self.authStore == null` - a server with + // auth enabled answers `auth` requests in `handleAuth` above, + // before this switch is ever entered. A server with auth + // disabled has nothing to check credentials against, so it + // always accepts, matching every other opcode's "auth is + // simply not enforced" behavior on this server. + .auth => try protocol.writeFrame(writer, @intFromEnum(protocol.Status.ok), ""), + + // Only reachable when `self.authStore == null`, for the exact + // same reason as `.auth` above: a server without auth enabled + // has no notion of a session either, so it accepts any token + // unconditionally rather than ever consulting a session store. + .auth_token => try protocol.writeFrame(writer, @intFromEnum(protocol.Status.ok), ""), +