Merge pull request #1 from dualfroz/feat/pending-work-2026-09
feat: land pending September work
This commit is contained in:
32 files changed
+40566
-125
No files matched your search
+1338
File diff suppressed because it is too large.
Load diff
+450
@@ -0,0 +1,450 @@
|
||||
const std = @import("std");
|
||||
const dbmod = @import("db.zig");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
const Db = dbmod.Db;
|
||||
const Hmac256 = std.crypto.auth.hmac.sha2.HmacSha256;
|
||||
|
||||
/// Username/password authentication and coarse-grained role authorization,
|
||||
/// layered over the same Db as tables/indexes/collections under its own
|
||||
/// namespace tag 'A' (row.zig/index.zig use 'R'/'I'/'M', collections.zig
|
||||
/// uses 'X', so no keyspace ever collides):
|
||||
/// "A" "U\0" username -> UserRecord{role, rounds, salt[16], hash[32]}
|
||||
/// Passwords are never stored: only a per-user salted PBKDF2-HMAC-SHA256
|
||||
/// hash, using `std.crypto.pwhash.pbkdf2` - the same primitive the standard
|
||||
/// library itself ships for exactly this purpose, not a hand-rolled one.
|
||||
pub const nsTag: u8 = 'A';
|
||||
const userPrefixBytes = [_]u8{ nsTag, 'U', 0 };
|
||||
const saltLen = 16;
|
||||
const hashLen = 32;
|
||||
const recordLen = 1 + 4 + saltLen + hashLen;
|
||||
|
||||
pub const AuthError = error{
|
||||
UserNotFound,
|
||||
UserExists,
|
||||
InvalidCredentials,
|
||||
Malformed,
|
||||
};
|
||||
|
||||
/// Roles form a total order, not a set of independent flags: readwrite
|
||||
/// implies readonly, admin implies both. `atLeast` is the only check any
|
||||
/// caller (server.zig's dispatch) ever needs.
|
||||
pub const Role = enum(u8) {
|
||||
readonly = 0,
|
||||
readwrite = 1,
|
||||
admin = 2,
|
||||
|
||||
pub fn fromByte(b: u8) !Role {
|
||||
return std.meta.intToEnum(Role, b) catch error.Malformed;
|
||||
}
|
||||
|
||||
pub fn parse(text: []const u8) !Role {
|
||||
if (std.mem.eql(u8, text, "readonly")) return .readonly;
|
||||
if (std.mem.eql(u8, text, "readwrite")) return .readwrite;
|
||||
if (std.mem.eql(u8, text, "admin")) return .admin;
|
||||
return error.Malformed;
|
||||
}
|
||||
|
||||
pub fn name(self: Role) []const u8 {
|
||||
return switch (self) {
|
||||
.readonly => "readonly",
|
||||
.readwrite => "readwrite",
|
||||
.admin => "admin",
|
||||
};
|
||||
}
|
||||
|
||||
pub fn atLeast(self: Role, min: Role) bool {
|
||||
return @intFromEnum(self) >= @intFromEnum(min);
|
||||
}
|
||||
};
|
||||
|
||||
/// Recommended real-world cost factor for interactive login (OWASP's floor
|
||||
/// for PBKDF2-HMAC-SHA256 as of 2023). A constant rather than a value
|
||||
/// baked into `hashPassword` so tests can pass a tiny round count and stay
|
||||
/// fast, while `main.zig` and any other real caller default to this.
|
||||
pub const recommendedRounds: u32 = 100_000;
|
||||
|
||||
fn userKey(allocator: Allocator, username: []const u8) ![]u8 {
|
||||
const out = try allocator.alloc(u8, userPrefixBytes.len + username.len);
|
||||
@memcpy(out[0..userPrefixBytes.len], &userPrefixBytes);
|
||||
@memcpy(out[userPrefixBytes.len..], username);
|
||||
return out;
|
||||
}
|
||||
|
||||
const UserRecord = struct {
|
||||
role: Role,
|
||||
rounds: u32,
|
||||
salt: [saltLen]u8,
|
||||
hash: [hashLen]u8,
|
||||
|
||||
fn encode(self: UserRecord) [recordLen]u8 {
|
||||
var out: [recordLen]u8 = undefined;
|
||||
out[0] = @intFromEnum(self.role);
|
||||
std.mem.writeInt(u32, out[1..5], self.rounds, .little);
|
||||
@memcpy(out[5 .. 5 + saltLen], &self.salt);
|
||||
@memcpy(out[5 + saltLen .. recordLen], &self.hash);
|
||||
return out;
|
||||
}
|
||||
|
||||
fn decode(data: []const u8) !UserRecord {
|
||||
if (data.len != recordLen) return error.Malformed;
|
||||
var rec: UserRecord = undefined;
|
||||
rec.role = try Role.fromByte(data[0]);
|
||||
rec.rounds = std.mem.readInt(u32, data[1..5], .little);
|
||||
@memcpy(&rec.salt, data[5 .. 5 + saltLen]);
|
||||
@memcpy(&rec.hash, data[5 + saltLen .. recordLen]);
|
||||
return rec;
|
||||
}
|
||||
};
|
||||
|
||||
fn randomSalt() [saltLen]u8 {
|
||||
var salt: [saltLen]u8 = undefined;
|
||||
std.crypto.random.bytes(&salt);
|
||||
return salt;
|
||||
}
|
||||
|
||||
fn hashPassword(password: []const u8, salt: [saltLen]u8, rounds: u32) ![hashLen]u8 {
|
||||
var out: [hashLen]u8 = undefined;
|
||||
try std.crypto.pwhash.pbkdf2(&out, password, &salt, rounds, Hmac256);
|
||||
return out;
|
||||
}
|
||||
|
||||
/// Fixed salt/rounds used to cost a lookup against a username that does not
|
||||
/// exist, so `verify` spends roughly the same CPU whether the account is
|
||||
/// missing or the password was merely wrong - closing the most obvious
|
||||
/// timing side channel for username enumeration over the network (the
|
||||
/// key comparison itself still only runs on the "user exists" path, since
|
||||
/// there is nothing real to compare against otherwise).
|
||||
const dummySalt = [_]u8{0} ** saltLen;
|
||||
const dummyRounds: u32 = 4096;
|
||||
|
||||
pub const UserInfo = struct {
|
||||
username: []u8,
|
||||
role: Role,
|
||||
|
||||
pub fn deinit(self: UserInfo, allocator: Allocator) void {
|
||||
allocator.free(self.username);
|
||||
}
|
||||
};
|
||||
|
||||
pub fn freeUserList(allocator: Allocator, users: []UserInfo) void {
|
||||
for (users) |u| u.deinit(allocator);
|
||||
allocator.free(users);
|
||||
}
|
||||
|
||||
/// Username/password/role management layered over a `*Db`. Every method
|
||||
/// touches only keys under `userPrefixBytes`; unrelated data in the same
|
||||
/// Db (tables, collections, raw KV) is never read or written. Not
|
||||
/// synchronized on its own, same as every other module built over Db -
|
||||
/// server.zig serializes access with its own connection-handling mutex.
|
||||
pub const Store = struct {
|
||||
db: *Db,
|
||||
|
||||
pub fn init(db: *Db) Store {
|
||||
return .{ .db = db };
|
||||
}
|
||||
|
||||
pub fn createUser(self: Store, allocator: Allocator, username: []const u8, password: []const u8, role: Role, rounds: u32) !void {
|
||||
const key = try userKey(allocator, username);
|
||||
defer allocator.free(key);
|
||||
if (try self.db.get(allocator, key)) |existing| {
|
||||
allocator.free(existing);
|
||||
return error.UserExists;
|
||||
}
|
||||
const salt = randomSalt();
|
||||
const record = UserRecord{
|
||||
.role = role,
|
||||
.rounds = rounds,
|
||||
.salt = salt,
|
||||
.hash = try hashPassword(password, salt, rounds),
|
||||
};
|
||||
const encoded = record.encode();
|
||||
try self.db.put(key, &encoded);
|
||||
}
|
||||
|
||||
fn mustGet(self: Store, allocator: Allocator, key: []const u8) !UserRecord {
|
||||
const raw = try self.db.get(allocator, key) orelse return error.UserNotFound;
|
||||
defer allocator.free(raw);
|
||||
return UserRecord.decode(raw);
|
||||
}
|
||||
|
||||
pub fn setPassword(self: Store, allocator: Allocator, username: []const u8, newPassword: []const u8, rounds: u32) !void {
|
||||
const key = try userKey(allocator, username);
|
||||
defer allocator.free(key);
|
||||
var record = try self.mustGet(allocator, key);
|
||||
record.rounds = rounds;
|
||||
record.salt = randomSalt();
|
||||
record.hash = try hashPassword(newPassword, record.salt, rounds);
|
||||
const encoded = record.encode();
|
||||
try self.db.put(key, &encoded);
|
||||
}
|
||||
|
||||
pub fn setRole(self: Store, allocator: Allocator, username: []const u8, role: Role) !void {
|
||||
const key = try userKey(allocator, username);
|
||||
defer allocator.free(key);
|
||||
var record = try self.mustGet(allocator, key);
|
||||
record.role = role;
|
||||
const encoded = record.encode();
|
||||
try self.db.put(key, &encoded);
|
||||
}
|
||||
|
||||
pub fn deleteUser(self: Store, allocator: Allocator, username: []const u8) !void {
|
||||
const key = try userKey(allocator, username);
|
||||
defer allocator.free(key);
|
||||
_ = try self.mustGet(allocator, key);
|
||||
try self.db.delete(key);
|
||||
}
|
||||
|
||||
/// Verifies `password` for `username`, returning the stored role on
|
||||
/// success. See `dummySalt`/`dummyRounds` for why a missing user still
|
||||
/// pays a full PBKDF2 pass before failing. The final comparison uses
|
||||
/// `std.crypto.utils.timingSafeEql` rather than `std.mem.eql`, since an
|
||||
/// early-exit byte comparison on the hash would leak how many leading
|
||||
/// bytes matched to a network attacker timing repeated attempts.
|
||||
pub fn verify(self: Store, allocator: Allocator, username: []const u8, password: []const u8) !Role {
|
||||
const key = try userKey(allocator, username);
|
||||
defer allocator.free(key);
|
||||
const raw = try self.db.get(allocator, key) orelse {
|
||||
_ = hashPassword(password, dummySalt, dummyRounds) catch {};
|
||||
return error.InvalidCredentials;
|
||||
};
|
||||
defer allocator.free(raw);
|
||||
const record = try UserRecord.decode(raw);
|
||||
const candidate = try hashPassword(password, record.salt, record.rounds);
|
||||
if (!std.crypto.utils.timingSafeEql([hashLen]u8, candidate, record.hash)) return error.InvalidCredentials;
|
||||
return record.role;
|
||||
}
|
||||
|
||||
/// Whether at least one admin account exists - `main.zig` uses this to
|
||||
/// refuse turning on server-side auth enforcement until an operator has
|
||||
/// created a way back in, so enabling `--auth` can never lock everyone
|
||||
/// out of a fresh database.
|
||||
pub fn hasAnyAdmin(self: Store, allocator: Allocator) !bool {
|
||||
var it = try self.db.scanPrefix(allocator, &userPrefixBytes);
|
||||
defer it.deinit();
|
||||
while (try it.next()) |entry| {
|
||||
const record = UserRecord.decode(entry.value) catch continue;
|
||||
if (record.role == .admin) return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
/// Returns every user with their role, independently owned - free with
|
||||
/// `freeUserList`. Order follows key order (lexicographic by username).
|
||||
pub fn listUsers(self: Store, allocator: Allocator) ![]UserInfo {
|
||||
var it = try self.db.scanPrefix(allocator, &userPrefixBytes);
|
||||
defer it.deinit();
|
||||
var out = std.ArrayList(UserInfo).init(allocator);
|
||||
errdefer {
|
||||
for (out.items) |u| u.deinit(allocator);
|
||||
out.deinit();
|
||||
}
|
||||
while (try it.next()) |entry| {
|
||||
const record = UserRecord.decode(entry.value) catch continue;
|
||||
const uname = entry.key[userPrefixBytes.len..];
|
||||
try out.append(.{ .username = try allocator.dupe(u8, uname), .role = record.role });
|
||||
}
|
||||
return out.toOwnedSlice();
|
||||
}
|
||||
};
|
||||
|
||||
fn testDb(allocator: Allocator, dir: std.fs.Dir) !Db {
|
||||
const path = try dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(path);
|
||||
return Db.open(allocator, path, .{});
|
||||
}
|
||||
|
||||
test "role ordering is a total order across readonly/readwrite/admin" {
|
||||
try std.testing.expect(Role.readwrite.atLeast(.readonly));
|
||||
try std.testing.expect(Role.admin.atLeast(.readonly));
|
||||
try std.testing.expect(Role.admin.atLeast(.readwrite));
|
||||
try std.testing.expect(Role.admin.atLeast(.admin));
|
||||
try std.testing.expect(!Role.readonly.atLeast(.readwrite));
|
||||
try std.testing.expect(!Role.readwrite.atLeast(.admin));
|
||||
}
|
||||
|
||||
test "Role.parse rejects an unknown role name" {
|
||||
try std.testing.expectEqual(Role.admin, try Role.parse("admin"));
|
||||
try std.testing.expectError(error.Malformed, Role.parse("superuser"));
|
||||
}
|
||||
|
||||
test "createUser then verify accepts the correct password" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
|
||||
const store = Store.init(&db);
|
||||
try store.createUser(allocator, "alice", "hunter2", .readwrite, 16);
|
||||
|
||||
const role = try store.verify(allocator, "alice", "hunter2");
|
||||
try std.testing.expectEqual(Role.readwrite, role);
|
||||
}
|
||||
|
||||
test "verify rejects a wrong password without revealing which part was wrong" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
|
||||
const store = Store.init(&db);
|
||||
try store.createUser(allocator, "alice", "hunter2", .readonly, 16);
|
||||
|
||||
try std.testing.expectError(error.InvalidCredentials, store.verify(allocator, "alice", "wrong"));
|
||||
}
|
||||
|
||||
test "verify rejects an unknown username the same way as a wrong password" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
|
||||
const store = Store.init(&db);
|
||||
try std.testing.expectError(error.InvalidCredentials, store.verify(allocator, "ghost", "anything"));
|
||||
}
|
||||
|
||||
test "createUser rejects a duplicate username" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
|
||||
const store = Store.init(&db);
|
||||
try store.createUser(allocator, "alice", "hunter2", .readonly, 16);
|
||||
try std.testing.expectError(error.UserExists, store.createUser(allocator, "alice", "different", .admin, 16));
|
||||
}
|
||||
|
||||
test "setPassword replaces the password and invalidates the old one" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
|
||||
const store = Store.init(&db);
|
||||
try store.createUser(allocator, "alice", "old-pass", .readonly, 16);
|
||||
try store.setPassword(allocator, "alice", "new-pass", 16);
|
||||
|
||||
try std.testing.expectError(error.InvalidCredentials, store.verify(allocator, "alice", "old-pass"));
|
||||
try std.testing.expectEqual(Role.readonly, try store.verify(allocator, "alice", "new-pass"));
|
||||
}
|
||||
|
||||
test "setPassword on a missing user reports UserNotFound" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
|
||||
const store = Store.init(&db);
|
||||
try std.testing.expectError(error.UserNotFound, store.setPassword(allocator, "ghost", "x", 16));
|
||||
}
|
||||
|
||||
test "setRole changes role without disturbing the password" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
|
||||
const store = Store.init(&db);
|
||||
try store.createUser(allocator, "alice", "hunter2", .readonly, 16);
|
||||
try store.setRole(allocator, "alice", .admin);
|
||||
|
||||
try std.testing.expectEqual(Role.admin, try store.verify(allocator, "alice", "hunter2"));
|
||||
}
|
||||
|
||||
test "deleteUser removes the account so verify and mustGet both fail" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
|
||||
const store = Store.init(&db);
|
||||
try store.createUser(allocator, "alice", "hunter2", .readonly, 16);
|
||||
try store.deleteUser(allocator, "alice");
|
||||
|
||||
try std.testing.expectError(error.InvalidCredentials, store.verify(allocator, "alice", "hunter2"));
|
||||
try std.testing.expectError(error.UserNotFound, store.deleteUser(allocator, "alice"));
|
||||
}
|
||||
|
||||
test "hasAnyAdmin reflects whether any created user carries the admin role" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
|
||||
const store = Store.init(&db);
|
||||
try std.testing.expect(!(try store.hasAnyAdmin(allocator)));
|
||||
|
||||
try store.createUser(allocator, "alice", "hunter2", .readwrite, 16);
|
||||
try std.testing.expect(!(try store.hasAnyAdmin(allocator)));
|
||||
|
||||
try store.createUser(allocator, "root", "toor", .admin, 16);
|
||||
try std.testing.expect(try store.hasAnyAdmin(allocator));
|
||||
}
|
||||
|
||||
test "listUsers returns every account with its role" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
|
||||
const store = Store.init(&db);
|
||||
try store.createUser(allocator, "alice", "a-pass", .readonly, 16);
|
||||
try store.createUser(allocator, "bob", "b-pass", .readwrite, 16);
|
||||
try store.createUser(allocator, "root", "r-pass", .admin, 16);
|
||||
|
||||
const users = try store.listUsers(allocator);
|
||||
defer freeUserList(allocator, users);
|
||||
try std.testing.expectEqual(@as(usize, 3), users.len);
|
||||
|
||||
var sawAdmin = false;
|
||||
for (users) |u| {
|
||||
if (std.mem.eql(u8, u.username, "root")) {
|
||||
try std.testing.expectEqual(Role.admin, u.role);
|
||||
sawAdmin = true;
|
||||
}
|
||||
}
|
||||
try std.testing.expect(sawAdmin);
|
||||
}
|
||||
|
||||
test "users survive a Db close and reopen" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
{
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
const store = Store.init(&db);
|
||||
try store.createUser(allocator, "alice", "hunter2", .admin, 16);
|
||||
}
|
||||
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
const store = Store.init(&db);
|
||||
try std.testing.expectEqual(Role.admin, try store.verify(allocator, "alice", "hunter2"));
|
||||
}
|
||||
|
||||
test "a malformed stored record is reported as Malformed, not a crash" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try testDb(allocator, tmp.dir);
|
||||
defer db.close();
|
||||
|
||||
const key = try userKey(allocator, "alice");
|
||||
defer allocator.free(key);
|
||||
try db.put(key, "too short to be a real record");
|
||||
|
||||
const store = Store.init(&db);
|
||||
try std.testing.expectError(error.Malformed, store.verify(allocator, "alice", "anything"));
|
||||
}
|
||||
+527
@@ -0,0 +1,527 @@
|
||||
const std = @import("std");
|
||||
const db_mod = @import("db.zig");
|
||||
const encode = @import("encode.zig");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
const Db = db_mod.Db;
|
||||
const Snapshot = db_mod.Snapshot;
|
||||
const WriteBatch = db_mod.WriteBatch;
|
||||
|
||||
// Portable logical backup format for cobalt. Dumps a consistent, point-in-time
|
||||
// view of the live keyspace (as returned by Snapshot.scan - tombstones already
|
||||
// resolved) into a self-describing file that can be replayed into any Db.
|
||||
//
|
||||
// Layout:
|
||||
// header: magic("CBLB" 4B) | version(u32) | seq(u64)
|
||||
// entry*: keyLen(u32) | valueLen(u32) | key | value | entryCrc(u32)
|
||||
// entryCrc = crc32(keyLen le bytes || valueLen le bytes || key || value)
|
||||
// footer: magic("CBLE" 4B) | entryCount(u64) | fileCrc(u32)
|
||||
// fileCrc = crc32 of every byte written before the footer magic
|
||||
// (header + all entries), rolling-accumulated while writing.
|
||||
//
|
||||
// Restore validates the whole file (magics, per-entry crc, trailing count,
|
||||
// whole-file crc) before touching the target Db, so a corrupt or truncated
|
||||
// backup never applies a partial write.
|
||||
|
||||
const headerMagic = "CBLB";
|
||||
const footerMagic = "CBLE";
|
||||
const formatVersion: u32 = 1;
|
||||
|
||||
pub const BackupError = error{
|
||||
BadMagic,
|
||||
UnsupportedVersion,
|
||||
Truncated,
|
||||
Corrupt,
|
||||
CountMismatch,
|
||||
};
|
||||
|
||||
pub const BackupInfo = struct {
|
||||
seq: u64,
|
||||
entryCount: u64,
|
||||
};
|
||||
|
||||
const RollingCrc = struct {
|
||||
hasher: std.hash.Crc32,
|
||||
|
||||
fn init() RollingCrc {
|
||||
return .{ .hasher = std.hash.Crc32.init() };
|
||||
}
|
||||
|
||||
fn feed(self: *RollingCrc, bytes: []const u8) void {
|
||||
self.hasher.update(bytes);
|
||||
}
|
||||
|
||||
fn value(self: RollingCrc) u32 {
|
||||
var copy = self.hasher;
|
||||
return copy.final();
|
||||
}
|
||||
};
|
||||
|
||||
fn writeAllCrc(writer: anytype, roll: *RollingCrc, bytes: []const u8) !void {
|
||||
try writer.writeAll(bytes);
|
||||
roll.feed(bytes);
|
||||
}
|
||||
|
||||
fn putU32Crc(writer: anytype, roll: *RollingCrc, v: u32) !void {
|
||||
var b: [4]u8 = undefined;
|
||||
std.mem.writeInt(u32, &b, v, encode.endian);
|
||||
try writeAllCrc(writer, roll, &b);
|
||||
}
|
||||
|
||||
fn putU64Crc(writer: anytype, roll: *RollingCrc, v: u64) !void {
|
||||
var b: [8]u8 = undefined;
|
||||
std.mem.writeInt(u64, &b, v, encode.endian);
|
||||
try writeAllCrc(writer, roll, &b);
|
||||
}
|
||||
|
||||
/// Writes a full, consistent backup of `snapshot` to `path` (created/truncated).
|
||||
/// Returns the number of entries written.
|
||||
pub fn dumpToFile(allocator: Allocator, snapshot: Snapshot, path: []const u8) !u64 {
|
||||
const file = try std.fs.cwd().createFile(path, .{ .truncate = true });
|
||||
defer file.close();
|
||||
|
||||
var bufferedWriter = std.io.bufferedWriter(file.writer());
|
||||
const writer = bufferedWriter.writer();
|
||||
|
||||
var roll = RollingCrc.init();
|
||||
try writeAllCrc(writer, &roll, headerMagic);
|
||||
try putU32Crc(writer, &roll, formatVersion);
|
||||
try putU64Crc(writer, &roll, snapshot.seq);
|
||||
|
||||
var iter = try snapshot.scan(allocator, null, null);
|
||||
defer iter.deinit();
|
||||
|
||||
var count: u64 = 0;
|
||||
var entryScratch = std.ArrayList(u8).init(allocator);
|
||||
defer entryScratch.deinit();
|
||||
|
||||
while (try iter.next()) |entry| {
|
||||
entryScratch.clearRetainingCapacity();
|
||||
var lenBuf: [8]u8 = undefined;
|
||||
std.mem.writeInt(u32, lenBuf[0..4], @intCast(entry.key.len), encode.endian);
|
||||
std.mem.writeInt(u32, lenBuf[4..8], @intCast(entry.value.len), encode.endian);
|
||||
try entryScratch.appendSlice(&lenBuf);
|
||||
try entryScratch.appendSlice(entry.key);
|
||||
try entryScratch.appendSlice(entry.value);
|
||||
const entryCrc = std.hash.Crc32.hash(entryScratch.items);
|
||||
|
||||
try writeAllCrc(writer, &roll, entryScratch.items);
|
||||
try putU32Crc(writer, &roll, entryCrc);
|
||||
count += 1;
|
||||
}
|
||||
|
||||
try writer.writeAll(footerMagic);
|
||||
var countBuf: [8]u8 = undefined;
|
||||
std.mem.writeInt(u64, &countBuf, count, encode.endian);
|
||||
try writer.writeAll(&countBuf);
|
||||
|
||||
const fileCrc = roll.value();
|
||||
var crcBuf: [4]u8 = undefined;
|
||||
std.mem.writeInt(u32, &crcBuf, fileCrc, encode.endian);
|
||||
try writer.writeAll(&crcBuf);
|
||||
|
||||
try bufferedWriter.flush();
|
||||
return count;
|
||||
}
|
||||
|
||||
const ParsedEntry = struct {
|
||||
key: []const u8,
|
||||
value: []const u8,
|
||||
};
|
||||
|
||||
/// One pass over a fully-loaded backup buffer: validates structure end to end
|
||||
/// (magics, per-entry crc, declared count, whole-file crc) without allocating
|
||||
/// per-entry copies. Returns the entries slice (positions into `data`) plus
|
||||
/// backup metadata, or a BackupError describing the first problem found.
|
||||
const ParsedBackup = struct {
|
||||
seq: u64,
|
||||
entries: []ParsedEntry,
|
||||
};
|
||||
|
||||
fn parse(allocator: Allocator, data: []const u8) !ParsedBackup {
|
||||
if (data.len < headerMagic.len + 4 + 8) return BackupError.Truncated;
|
||||
if (!std.mem.eql(u8, data[0..4], headerMagic)) return BackupError.BadMagic;
|
||||
|
||||
const version = encode.readU32(data, 4);
|
||||
if (version != formatVersion) return BackupError.UnsupportedVersion;
|
||||
const seq = encode.readU64(data, 8);
|
||||
|
||||
var pos: usize = 16;
|
||||
var entries = std.ArrayList(ParsedEntry).init(allocator);
|
||||
errdefer entries.deinit();
|
||||
|
||||
while (true) {
|
||||
if (pos + footerMagic.len <= data.len and std.mem.eql(u8, data[pos .. pos + 4], footerMagic)) {
|
||||
break;
|
||||
}
|
||||
if (pos + 8 > data.len) return BackupError.Truncated;
|
||||
const keyLen = encode.readU32(data, pos);
|
||||
const valueLen = encode.readU32(data, pos + 4);
|
||||
const entryStart = pos;
|
||||
const keyStart = pos + 8;
|
||||
const valueStart = keyStart + keyLen;
|
||||
const valueEnd = valueStart + valueLen;
|
||||
if (valueEnd + 4 > data.len) return BackupError.Truncated;
|
||||
|
||||
const entryBytes = data[entryStart..valueEnd];
|
||||
const storedCrc = encode.readU32(data, valueEnd);
|
||||
const actualCrc = std.hash.Crc32.hash(entryBytes);
|
||||
if (storedCrc != actualCrc) return BackupError.Corrupt;
|
||||
|
||||
try entries.append(.{ .key = data[keyStart..valueStart], .value = data[valueStart..valueEnd] });
|
||||
pos = valueEnd + 4;
|
||||
}
|
||||
|
||||
if (pos + 4 + 8 + 4 > data.len) return BackupError.Truncated;
|
||||
const declaredCount = encode.readU64(data, pos + 4);
|
||||
if (declaredCount != entries.items.len) return BackupError.CountMismatch;
|
||||
|
||||
const storedFileCrc = encode.readU32(data, pos + 4 + 8);
|
||||
const actualFileCrc = std.hash.Crc32.hash(data[0..pos]);
|
||||
if (storedFileCrc != actualFileCrc) return BackupError.Corrupt;
|
||||
|
||||
return .{ .seq = seq, .entries = try entries.toOwnedSlice() };
|
||||
}
|
||||
|
||||
/// Reads and structurally validates a backup file without applying it,
|
||||
/// returning its metadata (seq, entry count). Useful to verify a backup
|
||||
/// before trusting it for disaster recovery.
|
||||
pub fn inspectFile(allocator: Allocator, path: []const u8) !BackupInfo {
|
||||
const data = try std.fs.cwd().readFileAlloc(allocator, path, 1 << 34);
|
||||
defer allocator.free(data);
|
||||
const parsed = try parse(allocator, data);
|
||||
defer allocator.free(parsed.entries);
|
||||
return .{ .seq = parsed.seq, .entryCount = parsed.entries.len };
|
||||
}
|
||||
|
||||
/// Restores every entry from the backup at `path` into `db`, in chunks of
|
||||
/// `chunkSize` applied atomically via WriteBatch. The whole file is validated
|
||||
/// (magics, per-entry and whole-file crc, declared count) before any chunk is
|
||||
/// applied, so a corrupt backup fails without mutating the target database.
|
||||
/// Existing keys in `db` outside the backup are left untouched; keys present
|
||||
/// in both are overwritten with the backup's value.
|
||||
pub fn restoreFromFile(allocator: Allocator, db: *Db, path: []const u8, chunkSize: usize) !BackupInfo {
|
||||
const data = try std.fs.cwd().readFileAlloc(allocator, path, 1 << 34);
|
||||
defer allocator.free(data);
|
||||
|
||||
const parsed = try parse(allocator, data);
|
||||
defer allocator.free(parsed.entries);
|
||||
|
||||
var i: usize = 0;
|
||||
while (i < parsed.entries.len) {
|
||||
const end = @min(i + chunkSize, parsed.entries.len);
|
||||
var batch = WriteBatch.init(allocator);
|
||||
defer batch.deinit();
|
||||
for (parsed.entries[i..end]) |entry| {
|
||||
try batch.put(entry.key, entry.value);
|
||||
}
|
||||
try db.apply(&batch);
|
||||
i = end;
|
||||
}
|
||||
|
||||
return .{ .seq = parsed.seq, .entryCount = parsed.entries.len };
|
||||
}
|
||||
|
||||
test "backup round trip: dump then restore into a fresh db reproduces every key" {
|
||||
const allocator = std.testing.allocator;
|
||||
const srcDir = "test-backup-src-1";
|
||||
const dstDir = "test-backup-dst-1";
|
||||
const backupPath = "test-backup-1.cblb";
|
||||
std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
std.fs.cwd().deleteTree(dstDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(dstDir) catch {};
|
||||
defer std.fs.cwd().deleteFile(backupPath) catch {};
|
||||
|
||||
var src = try Db.open(allocator, srcDir, .{});
|
||||
defer src.close();
|
||||
|
||||
var keyBuf: [16]u8 = undefined;
|
||||
var valBuf: [16]u8 = undefined;
|
||||
var i: usize = 0;
|
||||
while (i < 200) : (i += 1) {
|
||||
const key = try std.fmt.bufPrint(&keyBuf, "k{d:0>6}", .{i});
|
||||
const value = try std.fmt.bufPrint(&valBuf, "v{d:0>6}", .{i});
|
||||
try src.put(key, value);
|
||||
}
|
||||
// overwrite one, delete another: backup must reflect resolved state.
|
||||
try src.put("k000005", "updated");
|
||||
try src.delete("k000010");
|
||||
|
||||
const snap = try src.snapshot();
|
||||
const count = try dumpToFile(allocator, snap, backupPath);
|
||||
try std.testing.expectEqual(@as(u64, 199), count);
|
||||
|
||||
var dst = try Db.open(allocator, dstDir, .{});
|
||||
defer dst.close();
|
||||
const info = try restoreFromFile(allocator, &dst, backupPath, 32);
|
||||
try std.testing.expectEqual(@as(u64, 199), info.entryCount);
|
||||
|
||||
i = 0;
|
||||
while (i < 200) : (i += 1) {
|
||||
const key = try std.fmt.bufPrint(&keyBuf, "k{d:0>6}", .{i});
|
||||
const got = try dst.get(allocator, key);
|
||||
if (i == 10) {
|
||||
try std.testing.expect(got == null);
|
||||
continue;
|
||||
}
|
||||
defer allocator.free(got.?);
|
||||
if (i == 5) {
|
||||
try std.testing.expectEqualStrings("updated", got.?);
|
||||
} else {
|
||||
const expected = try std.fmt.bufPrint(&valBuf, "v{d:0>6}", .{i});
|
||||
try std.testing.expectEqualStrings(expected, got.?);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
test "backup of empty db round trips to zero entries" {
|
||||
const allocator = std.testing.allocator;
|
||||
const srcDir = "test-backup-src-empty";
|
||||
const dstDir = "test-backup-dst-empty";
|
||||
const backupPath = "test-backup-empty.cblb";
|
||||
std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
std.fs.cwd().deleteTree(dstDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(dstDir) catch {};
|
||||
defer std.fs.cwd().deleteFile(backupPath) catch {};
|
||||
|
||||
var src = try Db.open(allocator, srcDir, .{});
|
||||
defer src.close();
|
||||
const snap = try src.snapshot();
|
||||
const count = try dumpToFile(allocator, snap, backupPath);
|
||||
try std.testing.expectEqual(@as(u64, 0), count);
|
||||
|
||||
const info = try inspectFile(allocator, backupPath);
|
||||
try std.testing.expectEqual(@as(u64, 0), info.entryCount);
|
||||
|
||||
var dst = try Db.open(allocator, dstDir, .{});
|
||||
defer dst.close();
|
||||
_ = try restoreFromFile(allocator, &dst, backupPath, 32);
|
||||
var iter = try dst.scan(allocator, null, null);
|
||||
defer iter.deinit();
|
||||
try std.testing.expect((try iter.next()) == null);
|
||||
}
|
||||
|
||||
test "backup is point-in-time: writes after snapshot are excluded" {
|
||||
const allocator = std.testing.allocator;
|
||||
const srcDir = "test-backup-src-pit";
|
||||
const backupPath = "test-backup-pit.cblb";
|
||||
std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
defer std.fs.cwd().deleteFile(backupPath) catch {};
|
||||
|
||||
var src = try Db.open(allocator, srcDir, .{});
|
||||
defer src.close();
|
||||
try src.put("before", "1");
|
||||
const snap = try src.snapshot();
|
||||
try src.put("after", "2");
|
||||
|
||||
const count = try dumpToFile(allocator, snap, backupPath);
|
||||
try std.testing.expectEqual(@as(u64, 1), count);
|
||||
|
||||
const info = try inspectFile(allocator, backupPath);
|
||||
try std.testing.expectEqual(snap.seq, info.seq);
|
||||
}
|
||||
|
||||
test "restore into db with pre-existing disjoint keys merges without touching them" {
|
||||
const allocator = std.testing.allocator;
|
||||
const srcDir = "test-backup-src-merge";
|
||||
const dstDir = "test-backup-dst-merge";
|
||||
const backupPath = "test-backup-merge.cblb";
|
||||
std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
std.fs.cwd().deleteTree(dstDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(dstDir) catch {};
|
||||
defer std.fs.cwd().deleteFile(backupPath) catch {};
|
||||
|
||||
var src = try Db.open(allocator, srcDir, .{});
|
||||
defer src.close();
|
||||
try src.put("from-backup", "yes");
|
||||
const count = try dumpToFile(allocator, try src.snapshot(), backupPath);
|
||||
try std.testing.expectEqual(@as(u64, 1), count);
|
||||
|
||||
var dst = try Db.open(allocator, dstDir, .{});
|
||||
defer dst.close();
|
||||
try dst.put("already-here", "kept");
|
||||
|
||||
_ = try restoreFromFile(allocator, &dst, backupPath, 32);
|
||||
|
||||
const kept = try dst.get(allocator, "already-here");
|
||||
defer allocator.free(kept.?);
|
||||
try std.testing.expectEqualStrings("kept", kept.?);
|
||||
|
||||
const restored = try dst.get(allocator, "from-backup");
|
||||
defer allocator.free(restored.?);
|
||||
try std.testing.expectEqualStrings("yes", restored.?);
|
||||
}
|
||||
|
||||
test "restore chunking applies every entry across multiple write batches" {
|
||||
const allocator = std.testing.allocator;
|
||||
const srcDir = "test-backup-src-chunks";
|
||||
const dstDir = "test-backup-dst-chunks";
|
||||
const backupPath = "test-backup-chunks.cblb";
|
||||
std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
std.fs.cwd().deleteTree(dstDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(dstDir) catch {};
|
||||
defer std.fs.cwd().deleteFile(backupPath) catch {};
|
||||
|
||||
var src = try Db.open(allocator, srcDir, .{});
|
||||
defer src.close();
|
||||
var keyBuf: [16]u8 = undefined;
|
||||
var i: usize = 0;
|
||||
while (i < 97) : (i += 1) {
|
||||
const key = try std.fmt.bufPrint(&keyBuf, "c{d:0>6}", .{i});
|
||||
try src.put(key, "x");
|
||||
}
|
||||
_ = try dumpToFile(allocator, try src.snapshot(), backupPath);
|
||||
|
||||
var dst = try Db.open(allocator, dstDir, .{});
|
||||
defer dst.close();
|
||||
const info = try restoreFromFile(allocator, &dst, backupPath, 10);
|
||||
try std.testing.expectEqual(@as(u64, 97), info.entryCount);
|
||||
|
||||
var found: usize = 0;
|
||||
var iter = try dst.scan(allocator, null, null);
|
||||
defer iter.deinit();
|
||||
while (try iter.next()) |_| found += 1;
|
||||
try std.testing.expectEqual(@as(usize, 97), found);
|
||||
}
|
||||
|
||||
test "corrupted entry payload is rejected before any restore write happens" {
|
||||
const allocator = std.testing.allocator;
|
||||
const srcDir = "test-backup-src-corrupt";
|
||||
const dstDir = "test-backup-dst-corrupt";
|
||||
const backupPath = "test-backup-corrupt.cblb";
|
||||
std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
std.fs.cwd().deleteTree(dstDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(dstDir) catch {};
|
||||
defer std.fs.cwd().deleteFile(backupPath) catch {};
|
||||
|
||||
var src = try Db.open(allocator, srcDir, .{});
|
||||
defer src.close();
|
||||
try src.put("intact", "value-one");
|
||||
try src.put("victim", "value-two");
|
||||
_ = try dumpToFile(allocator, try src.snapshot(), backupPath);
|
||||
|
||||
{
|
||||
const file = try std.fs.cwd().openFile(backupPath, .{ .mode = .read_write });
|
||||
defer file.close();
|
||||
// flip a byte squarely inside the entry region (well past the 16 byte header).
|
||||
try file.seekTo(24);
|
||||
var byte: [1]u8 = undefined;
|
||||
_ = try file.read(&byte);
|
||||
try file.seekTo(24);
|
||||
try file.writeAll(&.{byte[0] ^ 0xFF});
|
||||
}
|
||||
|
||||
var dst = try Db.open(allocator, dstDir, .{});
|
||||
defer dst.close();
|
||||
const result = restoreFromFile(allocator, &dst, backupPath, 32);
|
||||
try std.testing.expectError(BackupError.Corrupt, result);
|
||||
|
||||
// nothing from the corrupt backup should have been applied: validation
|
||||
// happens over the whole in-memory buffer before the first WriteBatch.
|
||||
var iter = try dst.scan(allocator, null, null);
|
||||
defer iter.deinit();
|
||||
try std.testing.expect((try iter.next()) == null);
|
||||
}
|
||||
|
||||
test "truncated file is rejected as Truncated, not a crash" {
|
||||
const allocator = std.testing.allocator;
|
||||
const srcDir = "test-backup-src-trunc";
|
||||
const dstDir = "test-backup-dst-trunc";
|
||||
const backupPath = "test-backup-trunc.cblb";
|
||||
std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
std.fs.cwd().deleteTree(dstDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(dstDir) catch {};
|
||||
defer std.fs.cwd().deleteFile(backupPath) catch {};
|
||||
|
||||
var src = try Db.open(allocator, srcDir, .{});
|
||||
defer src.close();
|
||||
try src.put("a", "b");
|
||||
_ = try dumpToFile(allocator, try src.snapshot(), backupPath);
|
||||
|
||||
{
|
||||
const data = try std.fs.cwd().readFileAlloc(allocator, backupPath, 1 << 20);
|
||||
defer allocator.free(data);
|
||||
const truncated = data[0 .. data.len - 5];
|
||||
const file = try std.fs.cwd().createFile(backupPath, .{ .truncate = true });
|
||||
defer file.close();
|
||||
try file.writeAll(truncated);
|
||||
}
|
||||
|
||||
var dst = try Db.open(allocator, dstDir, .{});
|
||||
defer dst.close();
|
||||
const result = restoreFromFile(allocator, &dst, backupPath, 32);
|
||||
try std.testing.expectError(BackupError.Truncated, result);
|
||||
}
|
||||
|
||||
test "bad magic header is rejected without reading the rest of the file" {
|
||||
const allocator = std.testing.allocator;
|
||||
const backupPath = "test-backup-badmagic.cblb";
|
||||
defer std.fs.cwd().deleteFile(backupPath) catch {};
|
||||
|
||||
const file = try std.fs.cwd().createFile(backupPath, .{ .truncate = true });
|
||||
{
|
||||
defer file.close();
|
||||
try file.writeAll("XXXX" ++ [_]u8{0} ** 12);
|
||||
}
|
||||
|
||||
const result = inspectFile(allocator, backupPath);
|
||||
try std.testing.expectError(BackupError.BadMagic, result);
|
||||
}
|
||||
|
||||
test "unsupported version is rejected" {
|
||||
const allocator = std.testing.allocator;
|
||||
const backupPath = "test-backup-badversion.cblb";
|
||||
defer std.fs.cwd().deleteFile(backupPath) catch {};
|
||||
|
||||
const file = try std.fs.cwd().createFile(backupPath, .{ .truncate = true });
|
||||
{
|
||||
defer file.close();
|
||||
var buf: [16]u8 = undefined;
|
||||
@memcpy(buf[0..4], headerMagic);
|
||||
std.mem.writeInt(u32, buf[4..8], 99, encode.endian);
|
||||
std.mem.writeInt(u64, buf[8..16], 0, encode.endian);
|
||||
try file.writeAll(&buf);
|
||||
}
|
||||
|
||||
const result = inspectFile(allocator, backupPath);
|
||||
try std.testing.expectError(BackupError.UnsupportedVersion, result);
|
||||
}
|
||||
|
||||
test "declared entry count mismatch is caught even if per-entry crcs are valid" {
|
||||
const allocator = std.testing.allocator;
|
||||
const srcDir = "test-backup-src-countmismatch";
|
||||
const backupPath = "test-backup-countmismatch.cblb";
|
||||
std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
defer std.fs.cwd().deleteTree(srcDir) catch {};
|
||||
defer std.fs.cwd().deleteFile(backupPath) catch {};
|
||||
|
||||
var src = try Db.open(allocator, srcDir, .{});
|
||||
defer src.close();
|
||||
try src.put("a", "1");
|
||||
try src.put("b", "2");
|
||||
_ = try dumpToFile(allocator, try src.snapshot(), backupPath);
|
||||
|
||||
{
|
||||
const data = try std.fs.cwd().readFileAlloc(allocator, backupPath, 1 << 20);
|
||||
defer allocator.free(data);
|
||||
// footer is the last 4(magic)+8(count)+4(crc)=16 bytes; corrupt count only,
|
||||
// which necessarily also breaks the whole-file crc, but count is checked first.
|
||||
const countPos = data.len - 12;
|
||||
var mutated = try allocator.dupe(u8, data);
|
||||
defer allocator.free(mutated);
|
||||
std.mem.writeInt(u64, mutated[countPos..][0..8], 999, encode.endian);
|
||||
const file = try std.fs.cwd().createFile(backupPath, .{ .truncate = true });
|
||||
defer file.close();
|
||||
try file.writeAll(mutated);
|
||||
}
|
||||
|
||||
const result = inspectFile(allocator, backupPath);
|
||||
try std.testing.expectError(BackupError.CountMismatch, result);
|
||||
}
|
||||
+703
@@ -0,0 +1,703 @@
|
||||
const std = @import("std");
|
||||
const protocol = @import("protocol.zig");
|
||||
const dbmod = @import("db.zig");
|
||||
const rowmod = @import("row.zig");
|
||||
const authmod = @import("auth.zig");
|
||||
const ttlmod = @import("ttl.zig");
|
||||
const collmod = @import("collections.zig");
|
||||
const aggregatemod = @import("aggregate.zig");
|
||||
const windowmod = @import("window.zig");
|
||||
const querymod = @import("query.zig");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
|
||||
pub const ClientError = error{
|
||||
ServerBadRequest,
|
||||
ServerInternal,
|
||||
ServerConflict,
|
||||
ServerRateLimited,
|
||||
ServerUnauthorized,
|
||||
ServerForbidden,
|
||||
ServerUserNotFound,
|
||||
ServerKeyNotFound,
|
||||
ServerCollectionNotFound,
|
||||
UnexpectedStatus,
|
||||
};
|
||||
|
||||
/// Result of a remote `table export`: the row count `dump.zig`'s exporter
|
||||
/// reported plus the rendered bytes, independently allocated so they
|
||||
/// outlive the response frame they arrived in.
|
||||
pub const RemoteExport = struct {
|
||||
rowCount: u64,
|
||||
content: []u8,
|
||||
|
||||
pub fn deinit(self: RemoteExport, allocator: Allocator) void {
|
||||
allocator.free(self.content);
|
||||
}
|
||||
};
|
||||
|
||||
fn statusToError(status: protocol.Status) ClientError {
|
||||
return switch (status) {
|
||||
.bad_request => error.ServerBadRequest,
|
||||
.internal => error.ServerInternal,
|
||||
.conflict => error.ServerConflict,
|
||||
.rate_limited => error.ServerRateLimited,
|
||||
.unauthorized => error.ServerUnauthorized,
|
||||
.forbidden => error.ServerForbidden,
|
||||
else => error.UnexpectedStatus,
|
||||
};
|
||||
}
|
||||
|
||||
/// A blocking client for the wire protocol in `protocol.zig`. Each call
|
||||
/// writes one request frame and reads exactly one response frame, so a
|
||||
/// `Client` must not be shared across threads without external locking.
|
||||
pub const Client = struct {
|
||||
allocator: Allocator,
|
||||
stream: std.net.Stream,
|
||||
/// The session token from the last successful `authenticate` call, if
|
||||
/// the server had `enableSessions` on (see server.zig) - null on a
|
||||
/// server without sessions, or before any successful `authenticate`.
|
||||
/// Owned by the `Client`, freed on `close()` and on every subsequent
|
||||
/// `authenticate`/`authenticateWithToken` call.
|
||||
token: ?[]u8 = null,
|
||||
|
||||
pub fn connect(allocator: Allocator, address: std.net.Address) !Client {
|
||||
const stream = try std.net.tcpConnectToAddress(address);
|
||||
return .{ .allocator = allocator, .stream = stream };
|
||||
}
|
||||
|
||||
pub fn close(self: *Client) void {
|
||||
if (self.token) |t| self.allocator.free(t);
|
||||
self.stream.close();
|
||||
}
|
||||
|
||||
fn setToken(self: *Client, payload: []const u8) !void {
|
||||
if (self.token) |old| self.allocator.free(old);
|
||||
self.token = if (payload.len > 0) try self.allocator.dupe(u8, payload) else null;
|
||||
}
|
||||
|
||||
/// The session token captured by the last successful `authenticate`,
|
||||
/// if any (see `token` above). Callers that want to reconnect later
|
||||
/// without a password should dupe this before the `Client` closes.
|
||||
pub fn sessionToken(self: *Client) ?[]const u8 {
|
||||
return self.token;
|
||||
}
|
||||
|
||||
fn roundTrip(self: *Client, op: protocol.Op, payload: []const u8) !protocol.Frame {
|
||||
try protocol.writeFrame(self.stream.writer(), @intFromEnum(op), payload);
|
||||
return protocol.readFrame(self.allocator, self.stream.reader());
|
||||
}
|
||||
|
||||
/// Authenticates this connection against a server that has auth
|
||||
/// enabled (see `server.zig`'s `enableAuth`). Every request on this
|
||||
/// stream after a successful call is dispatched under the role
|
||||
/// returned by the server's `auth.Store.verify`; a server without auth
|
||||
/// enabled always answers `ok` regardless of the credentials given. If
|
||||
/// the server also has sessions enabled (`enableSessions`), the
|
||||
/// response carries a fresh token captured into `sessionToken()` -
|
||||
/// otherwise `sessionToken()` reads null afterwards, same as before
|
||||
/// this call was ever made.
|
||||
pub fn authenticate(self: *Client, username: []const u8, password: []const u8) !void {
|
||||
const payload = try protocol.encodeAuth(self.allocator, username, password);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.auth, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
try self.setToken(frame.payload);
|
||||
}
|
||||
|
||||
/// Authenticates this connection using a token obtained from another
|
||||
/// connection's `sessionToken()`, instead of a username/password pair
|
||||
/// (see server.zig's `auth_token` opcode) - skips the password
|
||||
/// entirely and resumes the same role until the session expires. A
|
||||
/// server without sessions enabled accepts any token unconditionally,
|
||||
/// matching `authenticate`'s behavior against a server without auth.
|
||||
pub fn authenticateWithToken(self: *Client, token: []const u8) !void {
|
||||
const payload = try protocol.encodeAuthToken(self.allocator, token);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.auth_token, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
pub fn ping(self: *Client) !void {
|
||||
var frame = try self.roundTrip(.ping, "");
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
pub fn get(self: *Client, allocator: Allocator, key: []const u8) !?[]u8 {
|
||||
var frame = try self.roundTrip(.get, key);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
return switch (status) {
|
||||
.ok => try allocator.dupe(u8, frame.payload),
|
||||
.not_found => null,
|
||||
else => statusToError(status),
|
||||
};
|
||||
}
|
||||
|
||||
pub fn set(self: *Client, key: []const u8, value: []const u8) !void {
|
||||
const payload = try protocol.encodeSet(self.allocator, key, value);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.set, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
pub fn delete(self: *Client, key: []const u8) !void {
|
||||
var frame = try self.roundTrip(.del, key);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
pub fn scan(self: *Client, allocator: Allocator, args: protocol.ScanArgs) ![]protocol.OwnedEntry {
|
||||
const payload = try protocol.encodeScan(self.allocator, args);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.scan, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeEntries(allocator, frame.payload);
|
||||
}
|
||||
|
||||
pub fn prefix(self: *Client, allocator: Allocator, prefixBytes: []const u8, limit: u32) ![]protocol.OwnedEntry {
|
||||
const payload = try protocol.encodePrefix(self.allocator, prefixBytes, limit);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.prefix, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeEntries(allocator, frame.payload);
|
||||
}
|
||||
|
||||
pub fn stats(self: *Client) !dbmod.Stats {
|
||||
var frame = try self.roundTrip(.stats, "");
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeStats(frame.payload);
|
||||
}
|
||||
|
||||
/// Fetches the server's combined Prometheus text export (its own
|
||||
/// request/rate-limit/auth/session counters from `Server.registerMetrics`
|
||||
/// plus `Db.stats()`, in the one payload `Server.handleMetrics` builds) -
|
||||
/// the network counterpart of the CLI's local-only `stats --prom`.
|
||||
/// Returned text is independently owned; free with `allocator.free`.
|
||||
pub fn metricsText(self: *Client, allocator: Allocator) ![]u8 {
|
||||
var frame = try self.roundTrip(.metrics, "");
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return allocator.dupe(u8, frame.payload);
|
||||
}
|
||||
|
||||
pub fn flush(self: *Client) !void {
|
||||
var frame = try self.roundTrip(.flush, "");
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
pub fn compact(self: *Client) !void {
|
||||
var frame = try self.roundTrip(.compact, "");
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
pub fn batch(self: *Client, ops: []const protocol.BatchOp) !void {
|
||||
const payload = try protocol.encodeBatch(self.allocator, ops);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.batch, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
/// Remote `table explain`: the wire counterpart of the CLI's local
|
||||
/// `table explain`, reporting which strategy `query.Table.explain`
|
||||
/// says a matching `tableSelect` would take against the server's
|
||||
/// `*Db`, without running the query. `exprText` is the same `expr.zig`
|
||||
/// grammar as `tableSelect`'s (empty means no predicate, i.e. always a
|
||||
/// full scan). Returned `Plan` owns its `indexName` - free with
|
||||
/// `Plan.deinit`.
|
||||
pub fn tableExplain(self: *Client, allocator: Allocator, table: []const u8, exprText: []const u8) !querymod.Plan {
|
||||
const payload = try protocol.encodeTableExplain(self.allocator, table, exprText);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.table_explain, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodePlan(allocator, frame.payload);
|
||||
}
|
||||
|
||||
/// Remote `table update`: the wire counterpart of the CLI's local
|
||||
/// `table update`, applying `sets` to every row matching `exprText`
|
||||
/// (empty means every row, the same convention `tableSelect`/
|
||||
/// `tableExplain` use) against the server's `*Db` via
|
||||
/// `query.Table.updateWhere`. Returns the number of rows updated.
|
||||
pub fn tableUpdate(self: *Client, table: []const u8, sets: []const rowmod.Field, exprText: []const u8) !u64 {
|
||||
const payload = try protocol.encodeTableUpdate(self.allocator, table, sets, exprText);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.table_update, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeTtlSweepCount(frame.payload);
|
||||
}
|
||||
|
||||
/// Remote `table delete`: the wire counterpart of the CLI's local
|
||||
/// `table delete`, removing every row matching `exprText` (empty means
|
||||
/// every row, the same convention `tableUpdate`/`tableExplain` use)
|
||||
/// against the server's `*Db` via `query.Table.deleteWhere`. Returns the
|
||||
/// number of rows deleted.
|
||||
pub fn tableDelete(self: *Client, table: []const u8, exprText: []const u8) !u64 {
|
||||
const payload = try protocol.encodeTableDelete(self.allocator, table, exprText);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.table_delete, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeTtlSweepCount(frame.payload);
|
||||
}
|
||||
|
||||
/// Remote `table export`: the wire counterpart of the CLI's local
|
||||
/// `table export <table> csv|json <path>`, rendering the whole table on
|
||||
/// the server and returning both the row count and the rendered bytes
|
||||
/// in one round trip. `content` is a fresh allocation independent of
|
||||
/// the response frame - free it with `RemoteExport.deinit`.
|
||||
pub fn tableExport(self: *Client, allocator: Allocator, table: []const u8, format: protocol.DumpFormat) !RemoteExport {
|
||||
const payload = try protocol.encodeTableExport(self.allocator, table, format);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.table_export, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
const result = try protocol.decodeExportResult(frame.payload);
|
||||
const content = try allocator.dupe(u8, result.content);
|
||||
return .{ .rowCount = result.rowCount, .content = content };
|
||||
}
|
||||
|
||||
/// Remote `table import`: the wire counterpart of the CLI's local
|
||||
/// `table import <table> csv|json <path>`. `content` is data the
|
||||
/// caller already read into memory (e.g. from a local file) and sends
|
||||
/// to the server in one frame; the server runs it through `dump.zig`'s
|
||||
/// importer there. Returns the number of rows imported.
|
||||
pub fn tableImport(self: *Client, table: []const u8, format: protocol.DumpFormat, content: []const u8) !u64 {
|
||||
const payload = try protocol.encodeTableImport(self.allocator, table, format, content);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.table_import, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeTtlSweepCount(frame.payload);
|
||||
}
|
||||
|
||||
/// Remote `table select`: `exprText` is the same `expr.zig` predicate
|
||||
/// grammar the local CLI's `table select` accepts (empty selects every
|
||||
/// row). `offset`/`limit` page the result the same way
|
||||
/// `query.Table.selectPage` does locally (0 limit = unlimited), so a
|
||||
/// large remote result does not have to come back in a single frame.
|
||||
/// `orderBy`, when non-empty, sorts server-side via
|
||||
/// `Table.selectPageOrdered` before that page window applies (one or
|
||||
/// more columns applied left to right, ties on earlier columns broken
|
||||
/// by later ones) - the network counterpart of the local CLI's `table
|
||||
/// select --order-by`. Returned rows are independently owned - free
|
||||
/// with `protocol.freeRows`.
|
||||
pub fn tableSelect(self: *Client, allocator: Allocator, table: []const u8, exprText: []const u8, offset: u32, limit: u32, orderBy: []const querymod.OrderBy) ![]rowmod.Row {
|
||||
const payload = try protocol.encodeTableSelect(self.allocator, table, exprText, offset, limit, orderBy);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.table_select, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeRows(allocator, frame.payload);
|
||||
}
|
||||
|
||||
/// Remote `table join`: the wire counterpart of the CLI's local `table
|
||||
/// join [--left|--right|--full]`, running whichever of `join.zig`'s four
|
||||
/// equi-join variants `mode` names against the server's `*Db` instead of
|
||||
/// a locally opened one. `leftExprText`/`rightExprText` are the same
|
||||
/// `expr.zig` predicate grammar as `tableSelect`'s `exprText` (empty
|
||||
/// selects every row on that side), filtering each table before the
|
||||
/// join runs - the network counterpart of `innerJoin`'s
|
||||
/// `leftPred`/`rightPred` parameters. `offset`/`limit` window the
|
||||
/// already-joined result server-side (0/0 = unlimited), the same
|
||||
/// convention `tableSelect` uses, via `join.applyPage`. Returned pairs
|
||||
/// are independently owned - free with `protocol.freeJoinPairs`.
|
||||
pub fn tableJoin(
|
||||
self: *Client,
|
||||
allocator: Allocator,
|
||||
mode: protocol.JoinMode,
|
||||
leftTable: []const u8,
|
||||
leftField: []const u8,
|
||||
rightTable: []const u8,
|
||||
rightField: []const u8,
|
||||
offset: u32,
|
||||
limit: u32,
|
||||
leftExprText: []const u8,
|
||||
rightExprText: []const u8,
|
||||
) ![]protocol.JoinPair {
|
||||
const payload = try protocol.encodeTableJoin(self.allocator, mode, leftTable, leftField, rightTable, rightField, offset, limit, leftExprText, rightExprText);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.table_join, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeJoinedRows(allocator, frame.payload);
|
||||
}
|
||||
|
||||
/// Remote `table aggregate <table> * <specs> [expr...]` (no `GROUP BY`):
|
||||
/// the wire counterpart of the CLI's local `table aggregate` with a `*`
|
||||
/// group field, calling `aggregate.runGlobal` server-side over rows
|
||||
/// filtered by `exprText` (same `expr.zig` grammar, empty selects every
|
||||
/// row, as `tableSelect`). `havingText` (empty = no HAVING) filters the
|
||||
/// single result server-side via `aggregate.globalMatchesHaving` -
|
||||
/// returns zero values (not an error) when it excludes the only row.
|
||||
/// Otherwise always returns exactly `specs.len` values in the same order
|
||||
/// as `specs` - free with `protocol.freeAggValues`.
|
||||
pub fn tableAggregateGlobal(self: *Client, allocator: Allocator, table: []const u8, specs: []const aggregatemod.AggSpec, havingText: []const u8, exprText: []const u8) ![]rowmod.Value {
|
||||
const payload = try protocol.encodeTableAggregate(self.allocator, table, "", specs, 0, 0, &.{}, havingText, exprText);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.table_aggregate, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeAggGlobalResult(allocator, frame.payload);
|
||||
}
|
||||
|
||||
/// Remote `table aggregate <table> <groupField> <specs> [expr...]
|
||||
/// [having <expr>] [--limit N] [--offset N]`: the wire counterpart of
|
||||
/// the CLI's local `table aggregate` with a real GROUP BY field, calling
|
||||
/// `aggregate.run` server-side, then (when `havingText` is non-empty)
|
||||
/// `aggregate.filterHaving` to narrow the groups, then windowing what
|
||||
/// remains server-side with `aggregate.applyPageGroups` (0/0 =
|
||||
/// unlimited, the same convention `tableJoin`'s `offset`/`limit` use -
|
||||
/// HAVING always runs before the window, same as SQL). `orderBy`
|
||||
/// (possibly empty) sorts the HAVING-narrowed groups server-side via
|
||||
/// `aggregate.sortGroupsBy` before that window, the same
|
||||
/// HAVING-then-ORDER-BY-then-LIMIT clause order SQL evaluates - each
|
||||
/// `orderBy[i].field` must name either `groupField` or one of `specs`'
|
||||
/// result labels (see `aggregate.specLabel`). One `GroupRow` per
|
||||
/// distinct group value kept by HAVING and the window, in `orderBy`
|
||||
/// order when non-empty, otherwise first-seen order (not sorted), same
|
||||
/// as the local library call - free with `protocol.freeAggGroupRows`.
|
||||
pub fn tableAggregateGroups(self: *Client, allocator: Allocator, table: []const u8, groupField: []const u8, specs: []const aggregatemod.AggSpec, offset: u32, limit: u32, orderBy: []const querymod.OrderBy, havingText: []const u8, exprText: []const u8) ![]aggregatemod.GroupRow {
|
||||
const payload = try protocol.encodeTableAggregate(self.allocator, table, groupField, specs, offset, limit, orderBy, havingText, exprText);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.table_aggregate, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeAggGroupRows(allocator, frame.payload);
|
||||
}
|
||||
|
||||
/// Remote `table window <table> <kind[:field]> [expr...] [partition
|
||||
/// <field>[,field...]] [order <field[:desc]>[,field...]] [--limit N]
|
||||
/// [--offset N]`: the wire counterpart of the CLI's local `table
|
||||
/// window`, calling `window.compute` server-side over rows filtered by
|
||||
/// `exprText` (same `expr.zig` grammar, empty selects every row).
|
||||
/// Unlike `tableAggregateGroups`, this never collapses rows - every
|
||||
/// matching row keeps its own `protocol.WindowRow` - but `offset`/`limit`
|
||||
/// still window WHICH rows come back, walked in the same
|
||||
/// partition-then-order display sequence `window.orderedIndices`
|
||||
/// produces server-side (0/0 means unlimited, same convention as every
|
||||
/// other paginated opcode). Returns one `WindowRow` per surviving row, in
|
||||
/// that display order - free with `protocol.freeWindowRows`.
|
||||
pub fn tableWindow(self: *Client, allocator: Allocator, table: []const u8, specs: []const windowmod.WindowSpec, partitionBy: []const []const u8, orderBy: []const querymod.OrderBy, offset: u32, limit: u32, exprText: []const u8, windowFrame: windowmod.FrameMode) ![]protocol.WindowRow {
|
||||
const payload = try protocol.encodeTableWindow(self.allocator, table, specs, partitionBy, orderBy, offset, limit, exprText, windowFrame);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.table_window, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeWindowRows(allocator, frame.payload);
|
||||
}
|
||||
|
||||
/// Remote `subscribe <channel>[,channel...]`: registers this connection
|
||||
/// with the server's `pubsub.Hub` and blocks until the server's ack
|
||||
/// frame comes back, so a `publish` issued on a different connection
|
||||
/// right after this returns is guaranteed to reach this subscription -
|
||||
/// see `server.zig`'s `handleSubscribe` doc comment for why the ack
|
||||
/// exists at all. After this call, `self` must only ever be used with
|
||||
/// `nextMessage` - the server never answers another request opcode on
|
||||
/// this same connection again (a real, documented protocol
|
||||
/// simplification, not an oversight: see `handleSubscribe`'s doc
|
||||
/// comment on why a subscribed connection is receive-only). Returns the
|
||||
/// `pubsub.Hub`-assigned subscriber id, the only handle a later
|
||||
/// `unsubscribe` call - necessarily made through a DIFFERENT `Client`,
|
||||
/// since this one can never send another request - can use to name this
|
||||
/// subscription.
|
||||
pub fn subscribe(self: *Client, channels: []const []const u8) !u64 {
|
||||
const payload = try protocol.encodeSubscribe(self.allocator, channels);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.subscribe, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
const ack = try protocol.decodeSubscribeAck(frame.payload);
|
||||
return ack.subscriberId;
|
||||
}
|
||||
|
||||
/// Remote `unsubscribe <subscriber-id> <channel>[,channel...]`: the
|
||||
/// client-side counterpart of `server.zig`'s `unsubscribe` opcode. Sent
|
||||
/// over an ordinary connection (this one may be reused for other
|
||||
/// requests afterward, unlike the subscribe connection it targets) to
|
||||
/// drop a subset of `subscriberId`'s channels - or, if it names every
|
||||
/// channel that subscriber still has left, to end it outright (see
|
||||
/// `pubsub.Hub.unsubscribe`'s doc comment). Returns how many of
|
||||
/// `channels` were actually removed.
|
||||
pub fn unsubscribe(self: *Client, subscriberId: u64, channels: []const []const u8) !u32 {
|
||||
const payload = try protocol.encodeUnsubscribe(self.allocator, subscriberId, channels);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.unsubscribe, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeUnsubscribeResult(frame.payload);
|
||||
}
|
||||
|
||||
/// Reads the next pushed message on a connection that already called
|
||||
/// `subscribe`, or `null` once the server closes the connection (the
|
||||
/// same "no more frames coming" signal `roundTrip` treats as an error
|
||||
/// everywhere else, but is the routine end of a subscription here).
|
||||
/// Blocks indefinitely otherwise - there is no polling variant, exactly
|
||||
/// like `Subscriber.waitMessage` server-side. Returns an independently
|
||||
/// owned `protocol.PubSubMessage` - free with its own `deinit`.
|
||||
pub fn nextMessage(self: *Client, allocator: Allocator) !?protocol.PubSubMessage {
|
||||
var frame = protocol.readFrame(self.allocator, self.stream.reader()) catch |err| switch (err) {
|
||||
error.EndOfStream => return null,
|
||||
else => return err,
|
||||
};
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return try protocol.decodePubSubMessage(allocator, frame.payload);
|
||||
}
|
||||
|
||||
/// Remote `publish <channel> <message>`: delivers `message` to every
|
||||
/// subscriber currently registered on `channel` across every connection
|
||||
/// the server has, returning how many subscribers it reached (0 is a
|
||||
/// normal outcome for a channel nobody is listening to, not an error -
|
||||
/// see `pubsub.Hub.publish`'s doc comment). Any connection may publish,
|
||||
/// whether or not it has ever subscribed to anything itself.
|
||||
pub fn publish(self: *Client, channel: []const u8, message: []const u8) !u64 {
|
||||
const payload = try protocol.encodePublish(self.allocator, channel, message);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.publish, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodePublishResult(frame.payload);
|
||||
}
|
||||
|
||||
/// Remote `txn <script>`: applies every entry in `entries` as one atomic
|
||||
/// `Transaction` commit across however many tables they touch - the wire
|
||||
/// counterpart of the CLI's local `txn <script-file>` (main.zig's
|
||||
/// `runTxn`). Returns the number of operations committed, always
|
||||
/// `entries.len` on success since a partial commit is impossible; a
|
||||
/// conflicting transaction surfaces as `error.ServerConflict` with
|
||||
/// nothing applied, matching `error.Conflict` from a local `Transaction.
|
||||
/// commit()`.
|
||||
pub fn txnOp(self: *Client, entries: []const protocol.TxnOpEntry) !u32 {
|
||||
const payload = try protocol.encodeTxnOp(self.allocator, entries);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.txn_op, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeTxnResult(frame.payload);
|
||||
}
|
||||
|
||||
fn userAdminRoundTrip(self: *Client, action: protocol.UserAdminAction, username: []const u8, password: []const u8, role: authmod.Role, rounds: u32) !protocol.Frame {
|
||||
const payload = try protocol.encodeUserAdmin(self.allocator, action, username, password, role, rounds);
|
||||
defer self.allocator.free(payload);
|
||||
return self.roundTrip(.user_admin, payload);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `auth.Store.createUser`, gated server-side
|
||||
/// behind `admin` role whenever the server has auth enabled (see
|
||||
/// `server.zig`'s `requiredRole`). `rounds` of 0 asks the server to use
|
||||
/// its own `auth.recommendedRounds` default.
|
||||
pub fn adminCreateUser(self: *Client, username: []const u8, password: []const u8, role: authmod.Role, rounds: u32) !void {
|
||||
var frame = try self.userAdminRoundTrip(.create, username, password, role, rounds);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `auth.Store.setPassword`. Reports
|
||||
/// `error.ServerUserNotFound` (not the generic `UnexpectedStatus`) when
|
||||
/// the account does not exist, since that is a normal, expected outcome
|
||||
/// worth distinguishing from every other non-ok status.
|
||||
pub fn adminSetPassword(self: *Client, username: []const u8, password: []const u8, rounds: u32) !void {
|
||||
var frame = try self.userAdminRoundTrip(.set_password, username, password, .readonly, rounds);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status == .not_found) return error.ServerUserNotFound;
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `auth.Store.setRole`.
|
||||
pub fn adminSetRole(self: *Client, username: []const u8, role: authmod.Role) !void {
|
||||
var frame = try self.userAdminRoundTrip(.set_role, username, "", role, 0);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status == .not_found) return error.ServerUserNotFound;
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `auth.Store.deleteUser`.
|
||||
pub fn adminDeleteUser(self: *Client, username: []const u8) !void {
|
||||
var frame = try self.userAdminRoundTrip(.delete, username, "", .readonly, 0);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status == .not_found) return error.ServerUserNotFound;
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `auth.Store.listUsers`. Returned entries are
|
||||
/// independently owned - free with `authmod.freeUserList`.
|
||||
pub fn adminListUsers(self: *Client, allocator: Allocator) ![]authmod.UserInfo {
|
||||
var frame = try self.userAdminRoundTrip(.list, "", "", .readonly, 0);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeUserList(allocator, frame.payload);
|
||||
}
|
||||
|
||||
fn ttlOpRoundTrip(self: *Client, action: protocol.TtlAction, key: []const u8, seconds: u64) !protocol.Frame {
|
||||
const payload = try protocol.encodeTtlOp(self.allocator, action, key, seconds);
|
||||
defer self.allocator.free(payload);
|
||||
return self.roundTrip(.ttl_op, payload);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `ttl.setExpireSeconds`, gated server-side the
|
||||
/// same way the CLI's local `expire` is: `error.ServerKeyNotFound` (not
|
||||
/// the generic `UnexpectedStatus`) when `key` does not currently exist,
|
||||
/// since that is a normal, expected outcome worth distinguishing from
|
||||
/// every other non-ok status - matching `adminSetPassword`'s treatment
|
||||
/// of an unknown username.
|
||||
pub fn ttlExpire(self: *Client, key: []const u8, seconds: u64) !void {
|
||||
var frame = try self.ttlOpRoundTrip(.expire, key, seconds);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status == .not_found) return error.ServerKeyNotFound;
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `ttl.clearExpire`.
|
||||
pub fn ttlPersist(self: *Client, key: []const u8) !void {
|
||||
var frame = try self.ttlOpRoundTrip(.persist, key, 0);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `ttl.remaining`.
|
||||
pub fn ttlQuery(self: *Client, key: []const u8) !ttlmod.TtlResult {
|
||||
var frame = try self.ttlOpRoundTrip(.query, key, 0);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeTtlResult(frame.payload);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `ttl.sweepExpired`: reaps every due key on
|
||||
/// the server and reports how many were removed.
|
||||
pub fn ttlSweep(self: *Client) !u64 {
|
||||
var frame = try self.ttlOpRoundTrip(.sweep, "", 0);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeTtlSweepCount(frame.payload);
|
||||
}
|
||||
|
||||
fn collTtlOpRoundTrip(self: *Client, action: protocol.TtlAction, kind: collmod.Kind, name: []const u8, seconds: u64) !protocol.Frame {
|
||||
const payload = try protocol.encodeCollTtlOp(self.allocator, action, kind, name, seconds);
|
||||
defer self.allocator.free(payload);
|
||||
return self.roundTrip(.coll_ttl_op, payload);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `collttl.setExpireSeconds`, gated server-side
|
||||
/// the same way the CLI's local `coll <kind> expire` is:
|
||||
/// `error.ServerCollectionNotFound` (not the generic `UnexpectedStatus`)
|
||||
/// when the collection does not currently exist, matching `ttlExpire`'s
|
||||
/// treatment of an unknown flat key.
|
||||
pub fn collTtlExpire(self: *Client, kind: collmod.Kind, name: []const u8, seconds: u64) !void {
|
||||
var frame = try self.collTtlOpRoundTrip(.expire, kind, name, seconds);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status == .not_found) return error.ServerCollectionNotFound;
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `collttl.clearExpire`.
|
||||
pub fn collTtlPersist(self: *Client, kind: collmod.Kind, name: []const u8) !void {
|
||||
var frame = try self.collTtlOpRoundTrip(.persist, kind, name, 0);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `collttl.remaining`.
|
||||
pub fn collTtlQuery(self: *Client, kind: collmod.Kind, name: []const u8) !ttlmod.TtlResult {
|
||||
var frame = try self.collTtlOpRoundTrip(.query, kind, name, 0);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeTtlResult(frame.payload);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `collttl.sweepExpired`: reaps every due
|
||||
/// collection of every kind on the server and reports how many were
|
||||
/// removed. The `kind` argument to `collTtlOpRoundTrip` is a
|
||||
/// placeholder here (see `protocol.encodeCollTtlOp`'s doc comment) -
|
||||
/// `.counter` is chosen arbitrarily, same convention as `ttlSweep`
|
||||
/// passing an empty key.
|
||||
pub fn collTtlSweep(self: *Client) !u64 {
|
||||
var frame = try self.collTtlOpRoundTrip(.sweep, .counter, "", 0);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeTtlSweepCount(frame.payload);
|
||||
}
|
||||
|
||||
/// The wire counterpart of `main.zig`'s local `coll <kind> <op> ...`
|
||||
/// (backed by `collections.zig`) - one generic round trip for every
|
||||
/// `protocol.CollAction`, mirroring that opcode's own "one opcode, many
|
||||
/// actions" shape rather than exposing thirty near-identical wrapper
|
||||
/// methods here. `args` are the same CLI-text-form string arguments
|
||||
/// (values/members/indices/scores) `runColl` already parses locally, so
|
||||
/// the server decodes exactly the same numbers with exactly the same
|
||||
/// `std.fmt.parseInt`/`parseFloat` calls the local command uses. Unlike
|
||||
/// `collTtlExpire`, a missing hash field/list element/zset member is not
|
||||
/// an error here - it comes back as `OwnedCollReply.nil`, exactly like
|
||||
/// the local CLI prints `(nil)` rather than failing. Returned reply is
|
||||
/// independently owned; free with `reply.deinit(allocator)`.
|
||||
pub fn collOp(self: *Client, allocator: Allocator, action: protocol.CollAction, name: []const u8, args: []const []const u8) !protocol.OwnedCollReply {
|
||||
const payload = try protocol.encodeCollOp(self.allocator, action, name, args);
|
||||
defer self.allocator.free(payload);
|
||||
var frame = try self.roundTrip(.coll_op, payload);
|
||||
defer frame.deinit();
|
||||
const status = try protocol.Status.fromByte(frame.tag);
|
||||
if (status != .ok) return statusToError(status);
|
||||
return protocol.decodeCollReply(allocator, frame.payload);
|
||||
}
|
||||
};
|
||||
|
||||
test "statusToError maps every non-ok status to a distinct error" {
|
||||
try std.testing.expectEqual(error.ServerBadRequest, statusToError(.bad_request));
|
||||
try std.testing.expectEqual(error.ServerInternal, statusToError(.internal));
|
||||
try std.testing.expectEqual(error.ServerConflict, statusToError(.conflict));
|
||||
try std.testing.expectEqual(error.ServerRateLimited, statusToError(.rate_limited));
|
||||
try std.testing.expectEqual(error.ServerUnauthorized, statusToError(.unauthorized));
|
||||
try std.testing.expectEqual(error.ServerForbidden, statusToError(.forbidden));
|
||||
try std.testing.expectEqual(error.UnexpectedStatus, statusToError(.not_found));
|
||||
}
|
||||
+1423
File diff suppressed because it is too large.
Load diff
+354
@@ -0,0 +1,354 @@
|
||||
const std = @import("std");
|
||||
const dbmod = @import("db.zig");
|
||||
const ttlmod = @import("ttl.zig");
|
||||
const collmod = @import("collections.zig");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
const Db = dbmod.Db;
|
||||
const Kind = collmod.Kind;
|
||||
pub const TtlResult = ttlmod.TtlResult;
|
||||
|
||||
/// Per-collection expiry markers, sibling to ttl.zig's per-key markers but
|
||||
/// for a different unit: a typed collection (Counter/Hash/Set/List/ZSet) is
|
||||
/// not one physical Db key, it is a whole namespace of them (every field of
|
||||
/// a hash, every element of a list, both rows of every zset entry), so
|
||||
/// "expiring" one here always means dropping the ENTIRE namespace via
|
||||
/// collections.drop, not a single db.delete.
|
||||
/// "Y" <kind tag byte> 0x00 <name> -> raw i128 (expiresAtNanos, big-endian)
|
||||
///
|
||||
/// This used to share ttl.zig's 'T' tag (with a leading 0x01 sub-byte to
|
||||
/// keep individual marker keys from colliding byte-for-byte). That was a
|
||||
/// real bug, not just an odd design: ttl.zig's own `sweepExpired` scans by
|
||||
/// the single-byte prefix "T", and every collttl marker IS a byte-for-byte
|
||||
/// extension of that same prefix, so the flat sweeper's scan matched them
|
||||
/// too. Both modules store the identical i128-expiresAt payload shape, so
|
||||
/// `readI128` on a collttl marker "succeeded" without error; if it looked
|
||||
/// due, the flat sweeper would `db.delete` a harmless nonexistent flat key
|
||||
/// and then reconstruct and delete the ORIGINAL collttl marker key via its
|
||||
/// own `clearExpire(key[1..])` - silently clearing the collection's expiry
|
||||
/// marker without ever calling `collections.dropCollection`, so the
|
||||
/// collection's actual data was never removed even though flat-key
|
||||
/// `ttlSweptTotal` had counted a phantom reap for it. Caught by a real
|
||||
/// test that runs both sweepers against the same Db at once (see
|
||||
/// server.zig's "flat ttl sweeper and the coll ttl sweeper enabled at once"
|
||||
/// test) - a manual `ttl expire-sweep`/`ttl_op sweep` run on a Db that also
|
||||
/// has due collection markers would trip the same bug even with only one
|
||||
/// sweeper. Fixed by giving collttl.zig its own top-level tag ('Y', not
|
||||
/// used by row.zig/index.zig/collections.zig/auth.zig/ttl.zig), so it is
|
||||
/// prefix-disjoint from every other namespace, not merely full-key-disjoint
|
||||
/// the way the shared-'T' scheme assumed.
|
||||
pub const nsTag: u8 = 'Y';
|
||||
const collSub: u8 = 0x01;
|
||||
const sep: u8 = 0x00;
|
||||
|
||||
fn markerKey(allocator: Allocator, k: Kind, name: []const u8) ![]u8 {
|
||||
const out = try allocator.alloc(u8, name.len + 3);
|
||||
out[0] = nsTag;
|
||||
out[1] = collSub;
|
||||
out[2] = k.tagByte();
|
||||
// name is appended without its own leading separator on purpose: the
|
||||
// fixed 3-byte header (nsTag, collSub, kindTag) already makes every
|
||||
// marker key length-disjoint per kind boundary at byte 2, so no two
|
||||
// different (kind, name) pairs can ever produce the same marker key.
|
||||
@memcpy(out[3..], name);
|
||||
return out;
|
||||
}
|
||||
|
||||
fn putI128(out: *std.ArrayList(u8), v: i128) !void {
|
||||
var b: [16]u8 = undefined;
|
||||
std.mem.writeInt(i128, &b, v, .big);
|
||||
try out.appendSlice(&b);
|
||||
}
|
||||
|
||||
fn readI128(data: []const u8) !i128 {
|
||||
if (data.len != 16) return error.Malformed;
|
||||
return std.mem.readInt(i128, data[0..16], .big);
|
||||
}
|
||||
|
||||
/// Sets (or replaces) an absolute expiry deadline on the collection `name`
|
||||
/// of kind `k`. Does not verify the collection currently exists - callers
|
||||
/// that need existence-checked semantics (the CLI `coll <kind> expire`
|
||||
/// command) check collections.exists first, same as ttl.zig's setExpireAt
|
||||
/// vs the CLI's flat-key `expire`.
|
||||
pub fn setExpireAt(db: *Db, allocator: Allocator, k: Kind, name: []const u8, expiresAtNanos: i128) !void {
|
||||
const mk = try markerKey(allocator, k, name);
|
||||
defer allocator.free(mk);
|
||||
var out = std.ArrayList(u8).init(allocator);
|
||||
defer out.deinit();
|
||||
try putI128(&out, expiresAtNanos);
|
||||
try db.put(mk, out.items);
|
||||
}
|
||||
|
||||
pub fn setExpireSeconds(db: *Db, allocator: Allocator, k: Kind, name: []const u8, nowNanos: i128, seconds: u64) !void {
|
||||
try setExpireAt(db, allocator, k, name, nowNanos + @as(i128, @intCast(seconds)) * ttlmod.secondNanos);
|
||||
}
|
||||
|
||||
pub fn getExpiresAt(db: *Db, allocator: Allocator, k: Kind, name: []const u8) !?i128 {
|
||||
const mk = try markerKey(allocator, k, name);
|
||||
defer allocator.free(mk);
|
||||
const raw = try db.get(allocator, mk) orelse return null;
|
||||
defer allocator.free(raw);
|
||||
return try readI128(raw);
|
||||
}
|
||||
|
||||
pub fn clearExpire(db: *Db, allocator: Allocator, k: Kind, name: []const u8) !void {
|
||||
const mk = try markerKey(allocator, k, name);
|
||||
defer allocator.free(mk);
|
||||
try db.delete(mk);
|
||||
}
|
||||
|
||||
/// Redis TTL(2)-shaped result, same three cases as ttl.zig's `remaining`:
|
||||
/// `.noKey` when the collection has no physical presence at all,
|
||||
/// `.noExpiry` when it exists but carries no marker, `.seconds` otherwise.
|
||||
pub fn remaining(db: *Db, allocator: Allocator, k: Kind, name: []const u8, nowNanos: i128) !TtlResult {
|
||||
if (!try collmod.collectionExists(db, allocator, k, name)) return .noKey;
|
||||
const expiresAt = try getExpiresAt(db, allocator, k, name) orelse return .noExpiry;
|
||||
if (nowNanos >= expiresAt) return .{ .seconds = 0 };
|
||||
return .{ .seconds = @intCast(@divFloor(expiresAt - nowNanos, ttlmod.secondNanos)) };
|
||||
}
|
||||
|
||||
/// Lazily enforces expiry on access: if the collection carries a due
|
||||
/// marker, drops every physical key it owns (collections.drop - a real
|
||||
/// multi-key delete through the normal WAL/tombstone path, not a pretend
|
||||
/// removal) plus the marker itself, and reports true. The CLI's `coll`
|
||||
/// dispatcher calls this before every non-TTL operation so a caller never
|
||||
/// observes a logically-expired collection just because no sweep has run.
|
||||
pub fn expireIfDue(db: *Db, allocator: Allocator, k: Kind, name: []const u8, nowNanos: i128) !bool {
|
||||
const expiresAt = try getExpiresAt(db, allocator, k, name) orelse return false;
|
||||
if (!ttlmod.isDue(expiresAt, nowNanos)) return false;
|
||||
_ = try collmod.dropCollection(db, allocator, k, name);
|
||||
try clearExpire(db, allocator, k, name);
|
||||
return true;
|
||||
}
|
||||
|
||||
/// Same lazy check as `expireIfDue`, applied to every name in `names`
|
||||
/// instead of a single one. Closes the gap documented at the `Set`
|
||||
/// multi-set operations' call sites (`intersect`/`unionOf`/`difference`):
|
||||
/// those accept more than one set name, but only the first one was ever
|
||||
/// checked on the read path (the CLI's blanket single-name check runs on
|
||||
/// `args[0]` before dispatch), leaving `args[1..]` reachable only through
|
||||
/// the periodic/manual sweep. Returns the count actually reaped, same
|
||||
/// counting convention as `sweepExpired`.
|
||||
pub fn expireIfDueMany(db: *Db, allocator: Allocator, k: Kind, names: []const []const u8, nowNanos: i128) !usize {
|
||||
var reaped: usize = 0;
|
||||
for (names) |name| {
|
||||
if (try expireIfDue(db, allocator, k, name, nowNanos)) reaped += 1;
|
||||
}
|
||||
return reaped;
|
||||
}
|
||||
|
||||
/// Active sweep across every kind at once: scans the whole collection-TTL
|
||||
/// marker namespace, drops every collection whose marker is due. Returns
|
||||
/// the number of collections reaped (not the number of physical keys - a
|
||||
/// dropped hash with 5 fields counts once here, same granularity as
|
||||
/// ttl.zig's sweepExpired counts flat keys).
|
||||
pub fn sweepExpired(db: *Db, allocator: Allocator, nowNanos: i128) !usize {
|
||||
const Due = struct { k: Kind, name: []u8 };
|
||||
var due = std.ArrayList(Due).init(allocator);
|
||||
defer {
|
||||
for (due.items) |d| allocator.free(d.name);
|
||||
due.deinit();
|
||||
}
|
||||
{
|
||||
const prefix = &[_]u8{ nsTag, collSub };
|
||||
var iter = try db.scanPrefix(allocator, prefix);
|
||||
defer iter.deinit();
|
||||
while (try iter.next()) |entry| {
|
||||
// entry.key = nsTag collSub kindTag name...
|
||||
if (entry.key.len < 3) continue;
|
||||
const k: Kind = Kind.fromTagByte(entry.key[2]) orelse continue;
|
||||
const expiresAt = try readI128(entry.value);
|
||||
if (ttlmod.isDue(expiresAt, nowNanos)) {
|
||||
try due.append(.{ .k = k, .name = try allocator.dupe(u8, entry.key[3..]) });
|
||||
}
|
||||
}
|
||||
}
|
||||
for (due.items) |d| {
|
||||
_ = try collmod.dropCollection(db, allocator, d.k, d.name);
|
||||
try clearExpire(db, allocator, d.k, d.name);
|
||||
}
|
||||
return due.items.len;
|
||||
}
|
||||
|
||||
fn openTemp(tmp: *std.testing.TmpDir, allocator: Allocator) !Db {
|
||||
const dbPath = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(dbPath);
|
||||
return Db.open(allocator, dbPath, .{});
|
||||
}
|
||||
|
||||
test "remaining reports noKey for a collection that was never created" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
try std.testing.expectEqual(TtlResult.noKey, try remaining(&db, allocator, .hash, "missing", 0));
|
||||
}
|
||||
|
||||
test "remaining reports noExpiry then counted seconds for a real hash" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
_ = try collmod.Hash.set(&db, allocator, "profile", "name", "ada");
|
||||
try std.testing.expectEqual(TtlResult.noExpiry, try remaining(&db, allocator, .hash, "profile", 0));
|
||||
|
||||
try setExpireSeconds(&db, allocator, .hash, "profile", 1000, 30);
|
||||
const r = try remaining(&db, allocator, .hash, "profile", 1000);
|
||||
try std.testing.expectEqual(@as(u64, 30), r.seconds);
|
||||
}
|
||||
|
||||
test "expireIfDue drops every field of a hash, not just one key" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
_ = try collmod.Hash.set(&db, allocator, "session:1", "user", "ada");
|
||||
_ = try collmod.Hash.set(&db, allocator, "session:1", "role", "admin");
|
||||
try setExpireSeconds(&db, allocator, .hash, "session:1", 0, 5);
|
||||
|
||||
try std.testing.expect(!try expireIfDue(&db, allocator, .hash, "session:1", 4 * ttlmod.secondNanos));
|
||||
try std.testing.expectEqual(@as(usize, 2), try collmod.Hash.len(&db, allocator, "session:1"));
|
||||
|
||||
try std.testing.expect(try expireIfDue(&db, allocator, .hash, "session:1", 5 * ttlmod.secondNanos));
|
||||
try std.testing.expectEqual(@as(usize, 0), try collmod.Hash.len(&db, allocator, "session:1"));
|
||||
try std.testing.expect(!try collmod.collectionExists(&db, allocator, .hash, "session:1"));
|
||||
try std.testing.expectEqual(@as(?i128, null), try getExpiresAt(&db, allocator, .hash, "session:1"));
|
||||
}
|
||||
|
||||
test "expireIfDueMany reaps only the due names among several, mirroring a Set.intersect call" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
_ = try collmod.Set.add(&db, allocator, "a", "x");
|
||||
_ = try collmod.Set.add(&db, allocator, "b", "x");
|
||||
_ = try collmod.Set.add(&db, allocator, "c", "x");
|
||||
try setExpireSeconds(&db, allocator, .set, "b", 0, 5);
|
||||
|
||||
// Mirrors `coll set inter a b c`: the CLI's blanket lazy check only
|
||||
// ever covered args[0] ("a"); this call is what now covers args[1..].
|
||||
const reaped = try expireIfDueMany(&db, allocator, .set, &.{ "b", "c" }, 5 * ttlmod.secondNanos);
|
||||
try std.testing.expectEqual(@as(usize, 1), reaped);
|
||||
|
||||
try std.testing.expect(!try collmod.collectionExists(&db, allocator, .set, "b"));
|
||||
try std.testing.expect(try collmod.collectionExists(&db, allocator, .set, "a"));
|
||||
try std.testing.expect(try collmod.collectionExists(&db, allocator, .set, "c"));
|
||||
|
||||
// Now Set.intersect(a, c) sees a live pair and finds the real member,
|
||||
// instead of silently intersecting against a phantom "b" that a
|
||||
// periodic sweep hadn't reaped yet.
|
||||
const inter = try collmod.Set.intersect(&db, allocator, &.{ "a", "c" });
|
||||
defer {
|
||||
for (inter) |m| allocator.free(m);
|
||||
allocator.free(inter);
|
||||
}
|
||||
try std.testing.expectEqual(@as(usize, 1), inter.len);
|
||||
try std.testing.expectEqualStrings("x", inter[0]);
|
||||
}
|
||||
|
||||
test "expireIfDueMany is a no-op count of zero when nothing is due" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
_ = try collmod.Set.add(&db, allocator, "a", "x");
|
||||
_ = try collmod.Set.add(&db, allocator, "b", "x");
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 0), try expireIfDueMany(&db, allocator, .set, &.{ "a", "b" }, 0));
|
||||
}
|
||||
|
||||
test "clearExpire makes a collection immortal again" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
try collmod.List.pushRight(&db, allocator, "queue", "job-1");
|
||||
try setExpireSeconds(&db, allocator, .list, "queue", 0, 1);
|
||||
try clearExpire(&db, allocator, .list, "queue");
|
||||
|
||||
try std.testing.expectEqual(TtlResult.noExpiry, try remaining(&db, allocator, .list, "queue", 100 * ttlmod.secondNanos));
|
||||
try std.testing.expect(!try expireIfDue(&db, allocator, .list, "queue", 100 * ttlmod.secondNanos));
|
||||
try std.testing.expectEqual(@as(usize, 1), try collmod.List.len(&db, allocator, "queue"));
|
||||
}
|
||||
|
||||
test "sweepExpired reaps due collections of every kind without touching live ones" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
_ = try collmod.Counter.incrBy(&db, allocator, "hits", 3);
|
||||
try setExpireSeconds(&db, allocator, .counter, "hits", 0, 1);
|
||||
|
||||
_ = try collmod.Hash.set(&db, allocator, "profile", "name", "ada");
|
||||
try setExpireSeconds(&db, allocator, .hash, "profile", 0, 100);
|
||||
|
||||
try collmod.ZSet.add(&db, allocator, "leaderboard", "alice", 10);
|
||||
try setExpireSeconds(&db, allocator, .zset, "leaderboard", 0, 1);
|
||||
|
||||
_ = try collmod.Set.add(&db, allocator, "tags", "zig");
|
||||
|
||||
const reaped = try sweepExpired(&db, allocator, 2 * ttlmod.secondNanos);
|
||||
try std.testing.expectEqual(@as(usize, 2), reaped);
|
||||
|
||||
try std.testing.expect(!try collmod.collectionExists(&db, allocator, .counter, "hits"));
|
||||
try std.testing.expect(!try collmod.collectionExists(&db, allocator, .zset, "leaderboard"));
|
||||
try std.testing.expect(try collmod.collectionExists(&db, allocator, .hash, "profile"));
|
||||
try std.testing.expect(try collmod.collectionExists(&db, allocator, .set, "tags"));
|
||||
|
||||
// Idempotent: a second sweep at the same instant finds nothing left.
|
||||
try std.testing.expectEqual(@as(usize, 0), try sweepExpired(&db, allocator, 2 * ttlmod.secondNanos));
|
||||
}
|
||||
|
||||
test "collection sweep survives a wal reopen: markers are durable" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
{
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
_ = try collmod.Hash.set(&db, allocator, "durable", "k", "v");
|
||||
try setExpireSeconds(&db, allocator, .hash, "durable", 0, 5);
|
||||
}
|
||||
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
const r = try remaining(&db, allocator, .hash, "durable", 1 * ttlmod.secondNanos);
|
||||
try std.testing.expectEqual(@as(u64, 4), r.seconds);
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 1), try sweepExpired(&db, allocator, 10 * ttlmod.secondNanos));
|
||||
try std.testing.expect(!try collmod.collectionExists(&db, allocator, .hash, "durable"));
|
||||
}
|
||||
|
||||
test "markers for the same name under different kinds never collide" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
_ = try collmod.Hash.set(&db, allocator, "shared", "f", "v");
|
||||
_ = try collmod.Set.add(&db, allocator, "shared", "m");
|
||||
|
||||
try setExpireSeconds(&db, allocator, .hash, "shared", 0, 5);
|
||||
// The set named "shared" must still read as noExpiry - the marker for
|
||||
// the hash "shared" must not leak across the kind boundary.
|
||||
try std.testing.expectEqual(TtlResult.noExpiry, try remaining(&db, allocator, .set, "shared", 0));
|
||||
|
||||
const hashResult = try remaining(&db, allocator, .hash, "shared", 0);
|
||||
try std.testing.expectEqual(@as(u64, 5), hashResult.seconds);
|
||||
}
|
||||
@@ -0,0 +1,437 @@
|
||||
const std = @import("std");
|
||||
const dbmod = @import("db.zig");
|
||||
const batchmod = @import("batch.zig");
|
||||
const txnmod = @import("txn.zig");
|
||||
const txntable = @import("txn_table.zig");
|
||||
const querymod = @import("query.zig");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
const Db = dbmod.Db;
|
||||
const Transaction = txnmod.Transaction;
|
||||
const TableTxn = txntable.TableTxn;
|
||||
const WriteBatch = batchmod.WriteBatch;
|
||||
|
||||
/// A thread-safe facade over `*Db`. `Db` itself keeps no internal locking -
|
||||
/// a fact documented (and worked around ad hoc) at every call site that
|
||||
/// already needs concurrent access from more than one OS thread:
|
||||
/// server.zig takes its own `mutex` before touching `db` on every request,
|
||||
/// and every session note about txn_table.zig/query.zig/collections.zig
|
||||
/// repeats the same known gap for the same reason. `GuardedDb` is that
|
||||
/// exact pattern promoted to one reusable, independently tested primitive
|
||||
/// instead of being reimplemented by hand wherever it is next needed.
|
||||
///
|
||||
/// This does not make `Db` itself thread-safe: it makes "every access to
|
||||
/// this particular `Db` goes through this particular lock" enforceable,
|
||||
/// which is the same guarantee `Server` already relies on. Two independent
|
||||
/// `GuardedDb` wrapping the same underlying `*Db` would NOT be safe - the
|
||||
/// mutex only serializes callers that agree to go through it.
|
||||
pub const GuardedDb = struct {
|
||||
db: *Db,
|
||||
mutex: std.Thread.Mutex = .{},
|
||||
|
||||
pub fn init(db: *Db) GuardedDb {
|
||||
return .{ .db = db };
|
||||
}
|
||||
|
||||
/// Exclusive handle returned by `acquire()`. Holds the lock until
|
||||
/// `release()` runs, so a caller that needs several `Db` calls to
|
||||
/// happen as one atomic unit (a read-modify-write, or several puts
|
||||
/// that must never be observed half-applied by another guarded
|
||||
/// reader) can acquire once and issue every call through `.db`
|
||||
/// without releasing in between - the same shape as
|
||||
/// `Server.dispatch` locking `self.mutex` once per request no matter
|
||||
/// how many `Db` calls that request needs.
|
||||
pub const Guard = struct {
|
||||
db: *Db,
|
||||
mutex: *std.Thread.Mutex,
|
||||
|
||||
pub fn release(self: Guard) void {
|
||||
self.mutex.unlock();
|
||||
}
|
||||
};
|
||||
|
||||
pub fn acquire(self: *GuardedDb) Guard {
|
||||
self.mutex.lock();
|
||||
return .{ .db = self.db, .mutex = &self.mutex };
|
||||
}
|
||||
|
||||
pub fn put(self: *GuardedDb, key: []const u8, value: []const u8) !void {
|
||||
const g = self.acquire();
|
||||
defer g.release();
|
||||
try g.db.put(key, value);
|
||||
}
|
||||
|
||||
pub fn delete(self: *GuardedDb, key: []const u8) !void {
|
||||
const g = self.acquire();
|
||||
defer g.release();
|
||||
try g.db.delete(key);
|
||||
}
|
||||
|
||||
pub fn get(self: *GuardedDb, allocator: Allocator, key: []const u8) !?[]u8 {
|
||||
const g = self.acquire();
|
||||
defer g.release();
|
||||
return g.db.get(allocator, key);
|
||||
}
|
||||
|
||||
pub fn apply(self: *GuardedDb, batch: *WriteBatch) !void {
|
||||
const g = self.acquire();
|
||||
defer g.release();
|
||||
try g.db.apply(batch);
|
||||
}
|
||||
|
||||
pub fn flush(self: *GuardedDb) !void {
|
||||
const g = self.acquire();
|
||||
defer g.release();
|
||||
try g.db.flush();
|
||||
}
|
||||
|
||||
pub fn compact(self: *GuardedDb) !void {
|
||||
const g = self.acquire();
|
||||
defer g.release();
|
||||
try g.db.compact();
|
||||
}
|
||||
|
||||
/// Atomic read-modify-write counter: decodes the current value at
|
||||
/// `key` as a little-endian i64 (0 if absent or the wrong width),
|
||||
/// adds `delta`, writes the new value back - the whole read, decide,
|
||||
/// write sequence runs under one lock acquisition, so no other
|
||||
/// guarded caller can observe or clobber the value partway through.
|
||||
/// A bare `db.get` followed later by a separate `db.put` cannot make
|
||||
/// that guarantee, which is exactly the lost-update hazard this
|
||||
/// exists to close (see the test below that proves it with real
|
||||
/// concurrent OS threads, not a simulated interleaving).
|
||||
pub fn incrBy(self: *GuardedDb, allocator: Allocator, key: []const u8, delta: i64) !i64 {
|
||||
const g = self.acquire();
|
||||
defer g.release();
|
||||
var current: i64 = 0;
|
||||
if (try g.db.get(allocator, key)) |raw| {
|
||||
defer allocator.free(raw);
|
||||
if (raw.len == 8) current = std.mem.readInt(i64, raw[0..8], .little);
|
||||
}
|
||||
const next = current + delta;
|
||||
var buf: [8]u8 = undefined;
|
||||
std.mem.writeInt(i64, &buf, next, .little);
|
||||
try g.db.put(key, &buf);
|
||||
return next;
|
||||
}
|
||||
};
|
||||
|
||||
fn openTemp(tmp: *std.testing.TmpDir, allocator: Allocator) !Db {
|
||||
const dbPath = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(dbPath);
|
||||
return Db.open(allocator, dbPath, .{});
|
||||
}
|
||||
|
||||
test "put/get/delete round-trip through the guard like a plain Db call" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
var guarded = GuardedDb.init(&db);
|
||||
|
||||
try guarded.put("k", "v1");
|
||||
{
|
||||
const got = try guarded.get(allocator, "k");
|
||||
defer if (got) |g| allocator.free(g);
|
||||
try std.testing.expectEqualStrings("v1", got.?);
|
||||
}
|
||||
|
||||
try guarded.delete("k");
|
||||
const afterDelete = try guarded.get(allocator, "k");
|
||||
try std.testing.expectEqual(@as(?[]u8, null), afterDelete);
|
||||
}
|
||||
|
||||
test "incrBy on a fresh key starts from zero and returns the running total" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
var guarded = GuardedDb.init(&db);
|
||||
|
||||
try std.testing.expectEqual(@as(i64, 5), try guarded.incrBy(allocator, "ctr", 5));
|
||||
try std.testing.expectEqual(@as(i64, 3), try guarded.incrBy(allocator, "ctr", -2));
|
||||
try std.testing.expectEqual(@as(i64, 103), try guarded.incrBy(allocator, "ctr", 100));
|
||||
}
|
||||
|
||||
const IncrWorker = struct {
|
||||
guarded: *GuardedDb,
|
||||
key: []const u8,
|
||||
iterations: usize,
|
||||
|
||||
fn run(self: IncrWorker) void {
|
||||
const allocator = std.heap.page_allocator;
|
||||
var i: usize = 0;
|
||||
while (i < self.iterations) : (i += 1) {
|
||||
_ = self.guarded.incrBy(allocator, self.key, 1) catch |err| {
|
||||
std.debug.panic("incrBy failed under real thread contention: {s}", .{@errorName(err)});
|
||||
};
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
test "incrBy survives real concurrent OS threads with zero lost updates" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
var guarded = GuardedDb.init(&db);
|
||||
|
||||
const threadCount = 6;
|
||||
const perThread = 300;
|
||||
var threads: [threadCount]std.Thread = undefined;
|
||||
for (&threads) |*t| {
|
||||
t.* = try std.Thread.spawn(.{}, IncrWorker.run, .{IncrWorker{
|
||||
.guarded = &guarded,
|
||||
.key = "shared-counter",
|
||||
.iterations = perThread,
|
||||
}});
|
||||
}
|
||||
for (threads) |t| t.join();
|
||||
|
||||
const raw = try guarded.get(allocator, "shared-counter");
|
||||
defer allocator.free(raw.?);
|
||||
const total = std.mem.readInt(i64, raw.?[0..8], .little);
|
||||
// A bare unsynchronized `db.get` + `db.put` read-modify-write would
|
||||
// routinely lose updates here (two threads read the same stale value,
|
||||
// both write back the same next value, one increment vanishes). The
|
||||
// mutex inside GuardedDb makes every incrBy fully serialized, so the
|
||||
// total below must land exactly on threadCount*perThread every single
|
||||
// run - not "usually", not "close to" - or the guard has a real bug.
|
||||
try std.testing.expectEqual(@as(i64, threadCount * perThread), total);
|
||||
}
|
||||
|
||||
const TransferWorker = struct {
|
||||
guarded: *GuardedDb,
|
||||
rounds: usize,
|
||||
|
||||
fn run(self: TransferWorker) void {
|
||||
const allocator = std.heap.page_allocator;
|
||||
var i: usize = 0;
|
||||
while (i < self.rounds) : (i += 1) {
|
||||
const g = self.guarded.acquire();
|
||||
defer g.release();
|
||||
const aRaw = g.db.get(allocator, "balanceA") catch @panic("get a failed");
|
||||
defer allocator.free(aRaw.?);
|
||||
const bRaw = g.db.get(allocator, "balanceB") catch @panic("get b failed");
|
||||
defer allocator.free(bRaw.?);
|
||||
const a = std.mem.readInt(i64, aRaw.?[0..8], .little);
|
||||
const b = std.mem.readInt(i64, bRaw.?[0..8], .little);
|
||||
if (a < 1) continue;
|
||||
var aBuf: [8]u8 = undefined;
|
||||
var bBuf: [8]u8 = undefined;
|
||||
std.mem.writeInt(i64, &aBuf, a - 1, .little);
|
||||
std.mem.writeInt(i64, &bBuf, b + 1, .little);
|
||||
g.db.put("balanceA", &aBuf) catch @panic("put a failed");
|
||||
g.db.put("balanceB", &bBuf) catch @panic("put b failed");
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const InvariantReader = struct {
|
||||
guarded: *GuardedDb,
|
||||
checks: usize,
|
||||
violations: *std.atomic.Value(usize),
|
||||
|
||||
fn run(self: InvariantReader) void {
|
||||
const allocator = std.heap.page_allocator;
|
||||
var i: usize = 0;
|
||||
while (i < self.checks) : (i += 1) {
|
||||
const g = self.guarded.acquire();
|
||||
const aRaw = g.db.get(allocator, "balanceA") catch @panic("get a failed");
|
||||
const bRaw = g.db.get(allocator, "balanceB") catch @panic("get b failed");
|
||||
const a = std.mem.readInt(i64, aRaw.?[0..8], .little);
|
||||
const b = std.mem.readInt(i64, bRaw.?[0..8], .little);
|
||||
g.release();
|
||||
allocator.free(aRaw.?);
|
||||
allocator.free(bRaw.?);
|
||||
if (a + b != 1000) _ = self.violations.fetchAdd(1, .monotonic);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
test "acquire lets a writer move value between two keys atomically under a concurrent reader" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
var guarded = GuardedDb.init(&db);
|
||||
|
||||
var initA: [8]u8 = undefined;
|
||||
var initB: [8]u8 = undefined;
|
||||
std.mem.writeInt(i64, &initA, 1000, .little);
|
||||
std.mem.writeInt(i64, &initB, 0, .little);
|
||||
try guarded.put("balanceA", &initA);
|
||||
try guarded.put("balanceB", &initB);
|
||||
|
||||
var violations = std.atomic.Value(usize).init(0);
|
||||
const writerThread = try std.Thread.spawn(.{}, TransferWorker.run, .{TransferWorker{
|
||||
.guarded = &guarded,
|
||||
.rounds = 1000,
|
||||
}});
|
||||
const readerThread = try std.Thread.spawn(.{}, InvariantReader.run, .{InvariantReader{
|
||||
.guarded = &guarded,
|
||||
.checks = 1000,
|
||||
.violations = &violations,
|
||||
}});
|
||||
writerThread.join();
|
||||
readerThread.join();
|
||||
|
||||
// Every single reader observation is required to see the invariant
|
||||
// balanceA + balanceB == 1000 hold, because each transfer's two puts
|
||||
// happen inside one `acquire()`/`release()` pair the reader's own
|
||||
// acquire() must wait out entirely - a real race would show up here
|
||||
// as a nonzero violation count, not as a crash, which is why this
|
||||
// asserts the counter is exactly zero instead of merely "did not hang".
|
||||
try std.testing.expectEqual(@as(usize, 0), violations.load(.monotonic));
|
||||
|
||||
const finalARaw = try guarded.get(allocator, "balanceA");
|
||||
defer allocator.free(finalARaw.?);
|
||||
const finalBRaw = try guarded.get(allocator, "balanceB");
|
||||
defer allocator.free(finalBRaw.?);
|
||||
const finalA = std.mem.readInt(i64, finalARaw.?[0..8], .little);
|
||||
const finalB = std.mem.readInt(i64, finalBRaw.?[0..8], .little);
|
||||
try std.testing.expectEqual(@as(i64, 1000), finalA + finalB);
|
||||
try std.testing.expectEqual(@as(i64, 0), finalA);
|
||||
try std.testing.expectEqual(@as(i64, 1000), finalB);
|
||||
}
|
||||
|
||||
test "apply is usable concurrently: disjoint batches from different threads never clobber each other" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
var guarded = GuardedDb.init(&db);
|
||||
|
||||
const ApplyWorker = struct {
|
||||
guarded: *GuardedDb,
|
||||
prefix: u8,
|
||||
count: usize,
|
||||
|
||||
fn run(self: @This()) void {
|
||||
var batch = WriteBatch.init(std.heap.page_allocator);
|
||||
defer batch.deinit();
|
||||
var i: usize = 0;
|
||||
while (i < self.count) : (i += 1) {
|
||||
const key = [_]u8{ self.prefix, @intCast(i) };
|
||||
batch.put(&key, "v") catch @panic("batch put failed");
|
||||
}
|
||||
self.guarded.apply(&batch) catch |err| {
|
||||
std.debug.panic("apply failed under contention: {s}", .{@errorName(err)});
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
var threads: [4]std.Thread = undefined;
|
||||
for (&threads, 0..) |*t, idx| {
|
||||
t.* = try std.Thread.spawn(.{}, ApplyWorker.run, .{ApplyWorker{
|
||||
.guarded = &guarded,
|
||||
.prefix = @intCast('a' + idx),
|
||||
.count = 50,
|
||||
}});
|
||||
}
|
||||
for (threads) |t| t.join();
|
||||
|
||||
inline for (.{ 'a', 'b', 'c', 'd' }) |prefix| {
|
||||
var iter = try guarded.db.scanPrefix(allocator, &[_]u8{prefix});
|
||||
defer iter.deinit();
|
||||
var n: usize = 0;
|
||||
while (try iter.next()) |_| n += 1;
|
||||
try std.testing.expectEqual(@as(usize, 50), n);
|
||||
}
|
||||
}
|
||||
|
||||
const TableTransferWorker = struct {
|
||||
guarded: *GuardedDb,
|
||||
rounds: usize,
|
||||
|
||||
/// Runs a full `Transaction` + `TableTxn` transfer lifecycle (begin,
|
||||
/// read both rows, stage both writes, commit) under a single
|
||||
/// `acquire()`/`release()` pair, closing the specific gap every prior
|
||||
/// session on this project flagged and deferred: "brak testu
|
||||
/// wspolbieznosci wielowatkowej (dwa watki komituja TableTxn na tym
|
||||
/// samym Db jednoczesnie)". `Db.get`/`Transaction.commit` touch shared
|
||||
/// mutable state (the memtable, the OCC read-set bookkeeping) with no
|
||||
/// locking of their own, so without the guard two real OS threads
|
||||
/// racing through `TableTxn.open`/`get`/`put`/`commit` at once would
|
||||
/// be exactly the kind of unguarded concurrent access every other
|
||||
/// module in this project takes its own mutex to avoid.
|
||||
fn run(self: TableTransferWorker) void {
|
||||
const allocator = std.heap.page_allocator;
|
||||
var i: usize = 0;
|
||||
while (i < self.rounds) : (i += 1) {
|
||||
const g = self.guarded.acquire();
|
||||
defer g.release();
|
||||
|
||||
var txn = Transaction.begin(g.db);
|
||||
defer txn.deinit();
|
||||
var t = TableTxn.open(&txn, allocator, "accounts") catch @panic("TableTxn.open failed");
|
||||
defer t.deinit();
|
||||
|
||||
var from = (t.get(.{ .int = 1 }) catch @panic("get 1 failed")).?;
|
||||
defer from.deinit();
|
||||
var to = (t.get(.{ .int = 2 }) catch @panic("get 2 failed")).?;
|
||||
defer to.deinit();
|
||||
if (from.get("balance").?.int < 1) continue;
|
||||
|
||||
t.put(.{ .int = 1 }, &.{.{ .name = "balance", .value = .{ .int = from.get("balance").?.int - 1 } }}) catch @panic("put 1 failed");
|
||||
t.put(.{ .int = 2 }, &.{.{ .name = "balance", .value = .{ .int = to.get("balance").?.int + 1 } }}) catch @panic("put 2 failed");
|
||||
|
||||
// The whole sequence runs under one lock, so no external writer
|
||||
// can ever appear mid-transfer here: commit() must never see a
|
||||
// real error.Conflict in this test, unlike the deliberately
|
||||
// racy scenario in txn_table.zig's own "all-or-nothing" test.
|
||||
txn.commit() catch |err| std.debug.panic("unexpected commit failure under serialized access: {s}", .{@errorName(err)});
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
test "TableTxn commits from many real OS threads never lose or duplicate a unit of balance" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
const startBalance: i64 = 2000;
|
||||
{
|
||||
var setup = try querymod.Table.open(&db, allocator, "accounts");
|
||||
defer setup.deinit();
|
||||
try setup.put(.{ .int = 1 }, &.{.{ .name = "balance", .value = .{ .int = startBalance } }});
|
||||
try setup.put(.{ .int = 2 }, &.{.{ .name = "balance", .value = .{ .int = 0 } }});
|
||||
}
|
||||
|
||||
var guarded = GuardedDb.init(&db);
|
||||
const threadCount = 4;
|
||||
const roundsPerThread = 200;
|
||||
var threads: [threadCount]std.Thread = undefined;
|
||||
for (&threads) |*t| {
|
||||
t.* = try std.Thread.spawn(.{}, TableTransferWorker.run, .{TableTransferWorker{
|
||||
.guarded = &guarded,
|
||||
.rounds = roundsPerThread,
|
||||
}});
|
||||
}
|
||||
for (threads) |t| t.join();
|
||||
|
||||
var verify = try querymod.Table.open(&db, allocator, "accounts");
|
||||
defer verify.deinit();
|
||||
var row1 = (try verify.get(.{ .int = 1 })).?;
|
||||
defer row1.deinit();
|
||||
var row2 = (try verify.get(.{ .int = 2 })).?;
|
||||
defer row2.deinit();
|
||||
const balance1 = row1.get("balance").?.int;
|
||||
const balance2 = row2.get("balance").?.int;
|
||||
|
||||
// Every transfer moves exactly 1 unit and every transfer's two writes
|
||||
// commit together or not at all, so the total must be exactly
|
||||
// preserved and account 2 must have gained exactly what account 1
|
||||
// lost - no interleaving of `get`+`put`+`commit` across threads can
|
||||
// produce any other pair of numbers.
|
||||
try std.testing.expectEqual(startBalance, balance1 + balance2);
|
||||
try std.testing.expectEqual(@as(i64, threadCount * roundsPerThread), balance2);
|
||||
try std.testing.expectEqual(startBalance - threadCount * roundsPerThread, balance1);
|
||||
}
|
||||
+269
-111
@@ -4,7 +4,10 @@ const memtable = @import("memtable.zig");
|
||||
const sstable = @import("sstable.zig");
|
||||
const cache = @import("cache.zig");
|
||||
const statsmod = @import("stats.zig");
|
||||
const Wal = @import("wal.zig").Wal;
|
||||
const heap = @import("heap.zig");
|
||||
const version = @import("version.zig");
|
||||
const walmod = @import("wal.zig");
|
||||
const Wal = walmod.Wal;
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
const Kind = encode.Kind;
|
||||
@@ -13,6 +16,8 @@ const Memtable = memtable.Memtable;
|
||||
const Reader = sstable.Reader;
|
||||
const Builder = sstable.Builder;
|
||||
const BlockCache = cache.BlockCache;
|
||||
const VersionSet = version.VersionSet;
|
||||
const VersionEdit = version.VersionEdit;
|
||||
|
||||
pub const WriteBatch = @import("batch.zig").WriteBatch;
|
||||
pub const Stats = statsmod.Stats;
|
||||
@@ -22,7 +27,6 @@ fn order(a: []const u8, b: []const u8) std.math.Order {
|
||||
}
|
||||
|
||||
const walName = "wal.log";
|
||||
const manifestName = "MANIFEST";
|
||||
|
||||
pub const Options = struct {
|
||||
flushThreshold: usize = 4 * 1024 * 1024,
|
||||
@@ -46,6 +50,8 @@ pub const Db = struct {
|
||||
wal: Wal,
|
||||
sstables: std.ArrayList(*Reader),
|
||||
blockCache: BlockCache,
|
||||
versionSet: VersionSet,
|
||||
versionState: std.AutoHashMap(u64, usize),
|
||||
nextFileId: u64,
|
||||
nextSeq: u64,
|
||||
flushThreshold: usize,
|
||||
@@ -91,10 +97,11 @@ pub const Db = struct {
|
||||
if (id > maxId) maxId = id;
|
||||
}
|
||||
|
||||
var levelMap = try loadManifest(allocator, dir);
|
||||
defer levelMap.deinit();
|
||||
var opened = try VersionSet.open(allocator, dir);
|
||||
errdefer opened.set.close();
|
||||
errdefer opened.state.deinit();
|
||||
for (sstables.items) |reader| {
|
||||
reader.level = levelMap.get(reader.fileId) orelse 0;
|
||||
reader.level = opened.state.get(reader.fileId) orelse 0;
|
||||
}
|
||||
|
||||
var table = try Memtable.init(allocator);
|
||||
@@ -104,6 +111,20 @@ pub const Db = struct {
|
||||
const walMaxSeq = try wal.replay(&table);
|
||||
if (walMaxSeq > maxSeq) maxSeq = walMaxSeq;
|
||||
|
||||
// The directory scan alone can regress: a leveled compaction that
|
||||
// fully garbage-collects a run of tombstones produces zero output
|
||||
// files (see runMerge), so no fresh fileId is minted and the
|
||||
// deleted inputs' seq/fileId watermark vanishes from disk. The
|
||||
// manifest's persisted counters (stamped on every flush/compact,
|
||||
// see VersionEdit.nextFileId/nextSeq) are the authoritative floor
|
||||
// for exactly that case; take whichever recovery source is higher
|
||||
// so this is purely additive safety, never a regression versus the
|
||||
// scan-only behaviour this project already relied on.
|
||||
const scanNextFileId = maxId + 1;
|
||||
const scanNextSeq = maxSeq + 1;
|
||||
const nextFileId = @max(scanNextFileId, opened.set.counters.nextFileId orelse 0);
|
||||
const nextSeq = @max(scanNextSeq, opened.set.counters.nextSeq orelse 0);
|
||||
|
||||
const dirPathCopy = try allocator.dupe(u8, dirPath);
|
||||
|
||||
return .{
|
||||
@@ -114,8 +135,10 @@ pub const Db = struct {
|
||||
.wal = wal,
|
||||
.sstables = sstables,
|
||||
.blockCache = BlockCache.init(allocator, options.blockCacheBytes),
|
||||
.nextFileId = maxId + 1,
|
||||
.nextSeq = maxSeq + 1,
|
||||
.versionSet = opened.set,
|
||||
.versionState = opened.state,
|
||||
.nextFileId = nextFileId,
|
||||
.nextSeq = nextSeq,
|
||||
.flushThreshold = options.flushThreshold,
|
||||
.l0Trigger = options.l0CompactionTrigger,
|
||||
.l1BaseBytes = options.l1BaseBytes,
|
||||
@@ -132,6 +155,8 @@ pub const Db = struct {
|
||||
}
|
||||
self.sstables.deinit();
|
||||
self.blockCache.deinit();
|
||||
self.versionSet.close();
|
||||
self.versionState.deinit();
|
||||
self.memtable.deinit();
|
||||
self.dir.close();
|
||||
self.allocator.free(self.dirPath);
|
||||
@@ -234,7 +259,14 @@ pub const Db = struct {
|
||||
self.nextFileId += 1;
|
||||
self.memtable.clear();
|
||||
try self.wal.reset();
|
||||
try self.writeManifest();
|
||||
_ = try walmod.bumpGeneration(self.dir);
|
||||
|
||||
var edit = VersionEdit.init(self.allocator);
|
||||
defer edit.deinit();
|
||||
try edit.addFile(fileId, 0);
|
||||
edit.setNextFileId(self.nextFileId);
|
||||
edit.setNextSeq(self.nextSeq);
|
||||
try self.versionSet.logAndApply(&edit, &self.versionState);
|
||||
std.posix.fsync(self.dir.fd) catch {};
|
||||
}
|
||||
|
||||
@@ -254,7 +286,9 @@ pub const Db = struct {
|
||||
{
|
||||
var builder = try Builder.init(self.allocator, self.dir, name, expected);
|
||||
defer builder.deinit();
|
||||
while (pickWinner(cursors)) |winner| {
|
||||
var merger = try SstableMerger.init(self.allocator, cursors);
|
||||
defer merger.deinit();
|
||||
while (try merger.next()) |winner| {
|
||||
if (winner.kind == .value) try builder.add(winner);
|
||||
}
|
||||
try builder.finish();
|
||||
@@ -265,7 +299,10 @@ pub const Db = struct {
|
||||
reader.* = try Reader.open(self.allocator, self.dir, name, fileId);
|
||||
errdefer reader.deinit();
|
||||
|
||||
var edit = VersionEdit.init(self.allocator);
|
||||
defer edit.deinit();
|
||||
for (self.sstables.items) |old| {
|
||||
try edit.deleteFile(old.fileId);
|
||||
var oldNameBuf: [64]u8 = undefined;
|
||||
const oldName = sstName(&oldNameBuf, old.fileId) catch unreachable;
|
||||
old.deinit();
|
||||
@@ -276,7 +313,10 @@ pub const Db = struct {
|
||||
reader.level = 0;
|
||||
try self.sstables.append(reader);
|
||||
self.nextFileId += 1;
|
||||
try self.writeManifest();
|
||||
try edit.addFile(fileId, 0);
|
||||
edit.setNextFileId(self.nextFileId);
|
||||
edit.setNextSeq(self.nextSeq);
|
||||
try self.versionSet.logAndApply(&edit, &self.versionState);
|
||||
std.posix.fsync(self.dir.fd) catch {};
|
||||
}
|
||||
|
||||
@@ -365,7 +405,9 @@ pub const Db = struct {
|
||||
var outputName: [64]u8 = undefined;
|
||||
errdefer if (builder) |*b| b.deinit();
|
||||
|
||||
while (pickWinner(cursors)) |winner| {
|
||||
var merger = try SstableMerger.init(self.allocator, cursors);
|
||||
defer merger.deinit();
|
||||
while (try merger.next()) |winner| {
|
||||
if (dropTombstones and winner.kind == .tombstone) continue;
|
||||
if (builder == null) {
|
||||
outputId = self.nextFileId;
|
||||
@@ -385,7 +427,10 @@ pub const Db = struct {
|
||||
try self.finishOutput(done, outputId, targetLevel, &outputs);
|
||||
}
|
||||
|
||||
var edit = VersionEdit.init(self.allocator);
|
||||
defer edit.deinit();
|
||||
for (inputs) |old| {
|
||||
try edit.deleteFile(old.fileId);
|
||||
self.removeReader(old);
|
||||
var oldNameBuf: [64]u8 = undefined;
|
||||
const oldName = sstName(&oldNameBuf, old.fileId) catch unreachable;
|
||||
@@ -393,10 +438,21 @@ pub const Db = struct {
|
||||
self.allocator.destroy(old);
|
||||
self.dir.deleteFile(oldName) catch {};
|
||||
}
|
||||
for (outputs.items) |produced| try self.sstables.append(produced);
|
||||
for (outputs.items) |produced| {
|
||||
try edit.addFile(produced.fileId, produced.level);
|
||||
try self.sstables.append(produced);
|
||||
}
|
||||
outputs.deinit();
|
||||
|
||||
try self.writeManifest();
|
||||
// Stamped even when this round produced zero output files (every
|
||||
// input was a dropped tombstone): that is exactly the case where no
|
||||
// fresh fileId gets minted and a directory scan on the next restart
|
||||
// could otherwise regress below the seq/fileId already consumed by
|
||||
// the now-deleted inputs (see the comment on VersionEdit.nextFileId
|
||||
// and Db.open's counter merge).
|
||||
edit.setNextFileId(self.nextFileId);
|
||||
edit.setNextSeq(self.nextSeq);
|
||||
try self.versionSet.logAndApply(&edit, &self.versionState);
|
||||
std.posix.fsync(self.dir.fd) catch {};
|
||||
}
|
||||
|
||||
@@ -465,32 +521,6 @@ pub const Db = struct {
|
||||
return false;
|
||||
}
|
||||
|
||||
fn writeManifest(self: *Db) !void {
|
||||
var buf = std.ArrayList(u8).init(self.allocator);
|
||||
defer buf.deinit();
|
||||
const bodyWriter = encode.ByteWriter.init(&buf);
|
||||
try bodyWriter.putU32(@intCast(self.sstables.items.len));
|
||||
for (self.sstables.items) |reader| {
|
||||
try bodyWriter.putU64(reader.fileId);
|
||||
try bodyWriter.putU32(@intCast(reader.level));
|
||||
}
|
||||
|
||||
var frame = std.ArrayList(u8).init(self.allocator);
|
||||
defer frame.deinit();
|
||||
const frameWriter = encode.ByteWriter.init(&frame);
|
||||
try frameWriter.putU32(encode.crc(buf.items));
|
||||
try frameWriter.slice(buf.items);
|
||||
|
||||
const tmpName = "MANIFEST.tmp";
|
||||
{
|
||||
const file = try self.dir.createFile(tmpName, .{ .truncate = true });
|
||||
defer file.close();
|
||||
try file.writeAll(frame.items);
|
||||
try file.sync();
|
||||
}
|
||||
try self.dir.rename(tmpName, manifestName);
|
||||
}
|
||||
|
||||
pub fn levelOf(self: *Db, fileId: u64) ?usize {
|
||||
for (self.sstables.items) |reader| {
|
||||
if (reader.fileId == fileId) return reader.level;
|
||||
@@ -548,6 +578,12 @@ pub const Db = struct {
|
||||
cursors[0] = .{ .mem = self.memtable.cursor(options.start) };
|
||||
for (self.sstables.items, 0..) |reader, i| cursors[i + 1] = .{ .sst = reader.cursor(options.start) };
|
||||
|
||||
var mergeHeap = AscKeyHeap.init(allocator, {});
|
||||
errdefer mergeHeap.deinit();
|
||||
for (cursors, 0..) |*c, i| {
|
||||
if (c.valid()) try mergeHeap.push(.{ .key = c.record().key, .idx = i });
|
||||
}
|
||||
|
||||
var startCopy: ?[]u8 = null;
|
||||
errdefer if (startCopy) |s| allocator.free(s);
|
||||
if (options.start) |bound| startCopy = try allocator.dupe(u8, bound);
|
||||
@@ -558,6 +594,7 @@ pub const Db = struct {
|
||||
.resultAllocator = allocator,
|
||||
.arena = std.heap.ArenaAllocator.init(allocator),
|
||||
.cursors = cursors,
|
||||
.heap = mergeHeap,
|
||||
.start = startCopy,
|
||||
.end = endCopy,
|
||||
.startInclusive = options.startInclusive,
|
||||
@@ -596,6 +633,12 @@ pub const Db = struct {
|
||||
built += 1;
|
||||
}
|
||||
|
||||
var mergeHeap = DescKeyHeap.init(allocator, {});
|
||||
errdefer mergeHeap.deinit();
|
||||
for (cursors, 0..) |*c, i| {
|
||||
if (c.valid()) try mergeHeap.push(.{ .key = c.record().key, .idx = i });
|
||||
}
|
||||
|
||||
var startCopy: ?[]u8 = null;
|
||||
if (options.start) |bound| startCopy = try allocator.dupe(u8, bound);
|
||||
|
||||
@@ -603,6 +646,7 @@ pub const Db = struct {
|
||||
.resultAllocator = allocator,
|
||||
.arena = std.heap.ArenaAllocator.init(allocator),
|
||||
.cursors = cursors,
|
||||
.heap = mergeHeap,
|
||||
.start = startCopy,
|
||||
.startInclusive = options.startInclusive,
|
||||
.maxSeq = maxSeq,
|
||||
@@ -678,30 +722,68 @@ const SourceCursor = union(enum) {
|
||||
}
|
||||
};
|
||||
|
||||
fn pickWinner(cursors: []sstable.Cursor) ?Record {
|
||||
var minKey: ?[]const u8 = null;
|
||||
for (cursors) |*c| {
|
||||
if (c.valid()) {
|
||||
const key = c.record().key;
|
||||
if (minKey == null or order(key, minKey.?) == .lt) minKey = key;
|
||||
/// One live position in a k-way merge: the key it currently points at
|
||||
/// (a snapshot valid for the underlying source's whole lifetime, since
|
||||
/// sstable/memtable cursors never overwrite already-decoded key bytes)
|
||||
/// plus the index of the source cursor it belongs to.
|
||||
const KeyedIndex = struct { key: []const u8, idx: usize };
|
||||
|
||||
fn keyIndexLess(_: void, a: KeyedIndex, b: KeyedIndex) bool {
|
||||
return order(a.key, b.key) == .lt;
|
||||
}
|
||||
fn keyIndexGreater(_: void, a: KeyedIndex, b: KeyedIndex) bool {
|
||||
return order(a.key, b.key) == .gt;
|
||||
}
|
||||
|
||||
/// Ascending k-way merge heap: root always holds the smallest live key.
|
||||
const AscKeyHeap = heap.MinHeap(KeyedIndex, void, keyIndexLess);
|
||||
/// Descending k-way merge heap: root always holds the largest live key.
|
||||
const DescKeyHeap = heap.MinHeap(KeyedIndex, void, keyIndexGreater);
|
||||
|
||||
/// Merges N sorted sstable cursors (compaction inputs) in ascending key
|
||||
/// order using a binary heap: O(log N) per advance instead of the O(N)
|
||||
/// full rescan a naive linear-scan merge needs for every emitted record.
|
||||
/// Ties (same key present in several inputs) resolve to the entry with
|
||||
/// the highest seq, matching last-writer-wins semantics.
|
||||
const SstableMerger = struct {
|
||||
cursors: []sstable.Cursor,
|
||||
heap: AscKeyHeap,
|
||||
|
||||
fn init(allocator: Allocator, cursors: []sstable.Cursor) !SstableMerger {
|
||||
var h = AscKeyHeap.init(allocator, {});
|
||||
errdefer h.deinit();
|
||||
for (cursors, 0..) |*c, i| {
|
||||
if (c.valid()) try h.push(.{ .key = c.record().key, .idx = i });
|
||||
}
|
||||
return .{ .cursors = cursors, .heap = h };
|
||||
}
|
||||
if (minKey == null) return null;
|
||||
var winner: ?Record = null;
|
||||
for (cursors) |*c| {
|
||||
if (c.valid() and order(c.record().key, minKey.?) == .eq) {
|
||||
|
||||
fn deinit(self: *SstableMerger) void {
|
||||
self.heap.deinit();
|
||||
}
|
||||
|
||||
fn next(self: *SstableMerger) !?Record {
|
||||
const top = self.heap.peek() orelse return null;
|
||||
const minKey = top.key;
|
||||
var winner: ?Record = null;
|
||||
while (self.heap.peek()) |entry| {
|
||||
if (order(entry.key, minKey) != .eq) break;
|
||||
const popped = self.heap.pop().?;
|
||||
var c = &self.cursors[popped.idx];
|
||||
const rec = c.record();
|
||||
if (winner == null or rec.seq > winner.?.seq) winner = rec;
|
||||
c.advance();
|
||||
if (c.valid()) try self.heap.push(.{ .key = c.record().key, .idx = popped.idx });
|
||||
}
|
||||
return winner.?;
|
||||
}
|
||||
return winner.?;
|
||||
}
|
||||
};
|
||||
|
||||
pub const RangeIterator = struct {
|
||||
resultAllocator: Allocator,
|
||||
arena: std.heap.ArenaAllocator,
|
||||
cursors: []SourceCursor,
|
||||
heap: AscKeyHeap,
|
||||
start: ?[]u8,
|
||||
end: ?[]u8,
|
||||
startInclusive: bool,
|
||||
@@ -710,26 +792,22 @@ pub const RangeIterator = struct {
|
||||
|
||||
pub fn next(self: *RangeIterator) !?Entry {
|
||||
while (true) {
|
||||
var minKey: ?[]const u8 = null;
|
||||
for (self.cursors) |*c| {
|
||||
if (c.valid()) {
|
||||
const key = c.record().key;
|
||||
if (minKey == null or order(key, minKey.?) == .lt) minKey = key;
|
||||
}
|
||||
}
|
||||
if (minKey == null) return null;
|
||||
const top = self.heap.peek() orelse return null;
|
||||
const minKey = top.key;
|
||||
if (self.end) |bound| {
|
||||
const past = if (self.endInclusive) order(minKey.?, bound) == .gt else order(minKey.?, bound) != .lt;
|
||||
const past = if (self.endInclusive) order(minKey, bound) == .gt else order(minKey, bound) != .lt;
|
||||
if (past) return null;
|
||||
}
|
||||
|
||||
var winner: ?Record = null;
|
||||
for (self.cursors) |*c| {
|
||||
if (c.valid() and order(c.record().key, minKey.?) == .eq) {
|
||||
const rec = c.record();
|
||||
if (rec.seq <= self.maxSeq and (winner == null or rec.seq > winner.?.seq)) winner = rec;
|
||||
c.advance();
|
||||
}
|
||||
while (self.heap.peek()) |entry| {
|
||||
if (order(entry.key, minKey) != .eq) break;
|
||||
const popped = self.heap.pop().?;
|
||||
var c = &self.cursors[popped.idx];
|
||||
const rec = c.record();
|
||||
if (rec.seq <= self.maxSeq and (winner == null or rec.seq > winner.?.seq)) winner = rec;
|
||||
c.advance();
|
||||
if (c.valid()) try self.heap.push(.{ .key = c.record().key, .idx = popped.idx });
|
||||
}
|
||||
if (winner == null) continue;
|
||||
if (!self.startInclusive) {
|
||||
@@ -749,6 +827,7 @@ pub const RangeIterator = struct {
|
||||
|
||||
pub fn deinit(self: *RangeIterator) void {
|
||||
self.arena.deinit();
|
||||
self.heap.deinit();
|
||||
if (self.start) |bound| self.resultAllocator.free(bound);
|
||||
if (self.end) |bound| self.resultAllocator.free(bound);
|
||||
self.resultAllocator.free(self.cursors);
|
||||
@@ -792,32 +871,29 @@ pub const ReverseRangeIterator = struct {
|
||||
resultAllocator: Allocator,
|
||||
arena: std.heap.ArenaAllocator,
|
||||
cursors: []ReverseSource,
|
||||
heap: DescKeyHeap,
|
||||
start: ?[]u8,
|
||||
startInclusive: bool,
|
||||
maxSeq: u64,
|
||||
|
||||
pub fn next(self: *ReverseRangeIterator) !?Entry {
|
||||
while (true) {
|
||||
var maxKey: ?[]const u8 = null;
|
||||
for (self.cursors) |*c| {
|
||||
if (c.valid()) {
|
||||
const key = c.record().key;
|
||||
if (maxKey == null or order(key, maxKey.?) == .gt) maxKey = key;
|
||||
}
|
||||
}
|
||||
if (maxKey == null) return null;
|
||||
const top = self.heap.peek() orelse return null;
|
||||
const maxKey = top.key;
|
||||
if (self.start) |bound| {
|
||||
const past = if (self.startInclusive) order(maxKey.?, bound) == .lt else order(maxKey.?, bound) != .gt;
|
||||
const past = if (self.startInclusive) order(maxKey, bound) == .lt else order(maxKey, bound) != .gt;
|
||||
if (past) return null;
|
||||
}
|
||||
|
||||
var winner: ?Record = null;
|
||||
for (self.cursors) |*c| {
|
||||
if (c.valid() and order(c.record().key, maxKey.?) == .eq) {
|
||||
const rec = c.record();
|
||||
if (rec.seq <= self.maxSeq and (winner == null or rec.seq > winner.?.seq)) winner = rec;
|
||||
c.advance();
|
||||
}
|
||||
while (self.heap.peek()) |entry| {
|
||||
if (order(entry.key, maxKey) != .eq) break;
|
||||
const popped = self.heap.pop().?;
|
||||
var c = &self.cursors[popped.idx];
|
||||
const rec = c.record();
|
||||
if (rec.seq <= self.maxSeq and (winner == null or rec.seq > winner.?.seq)) winner = rec;
|
||||
c.advance();
|
||||
if (c.valid()) try self.heap.push(.{ .key = c.record().key, .idx = popped.idx });
|
||||
}
|
||||
if (winner == null) continue;
|
||||
if (winner.?.kind == .tombstone) continue;
|
||||
@@ -832,6 +908,7 @@ pub const ReverseRangeIterator = struct {
|
||||
|
||||
pub fn deinit(self: *ReverseRangeIterator) void {
|
||||
self.arena.deinit();
|
||||
self.heap.deinit();
|
||||
for (self.cursors) |*c| c.deinit();
|
||||
if (self.start) |bound| self.resultAllocator.free(bound);
|
||||
self.resultAllocator.free(self.cursors);
|
||||
@@ -842,36 +919,6 @@ fn rangesOverlap(aMin: []const u8, aMax: []const u8, bMin: []const u8, bMax: []c
|
||||
return order(aMin, bMax) != .gt and order(bMin, aMax) != .gt;
|
||||
}
|
||||
|
||||
fn loadManifest(allocator: Allocator, dir: std.fs.Dir) !std.AutoHashMap(u64, usize) {
|
||||
var map = std.AutoHashMap(u64, usize).init(allocator);
|
||||
errdefer map.deinit();
|
||||
|
||||
const file = dir.openFile(manifestName, .{}) catch |err| switch (err) {
|
||||
error.FileNotFound => return map,
|
||||
else => return err,
|
||||
};
|
||||
defer file.close();
|
||||
const data = try file.readToEndAlloc(allocator, std.math.maxInt(usize));
|
||||
defer allocator.free(data);
|
||||
if (data.len < 8) return map;
|
||||
|
||||
const storedCrc = encode.readU32(data, 0);
|
||||
const body = data[4..];
|
||||
if (encode.crc(body) != storedCrc) return map;
|
||||
|
||||
const count = encode.readU32(body, 0);
|
||||
var pos: usize = 4;
|
||||
var i: u32 = 0;
|
||||
while (i < count) : (i += 1) {
|
||||
if (pos + 12 > body.len) break;
|
||||
const fileId = encode.readU64(body, pos);
|
||||
const level = encode.readU32(body, pos + 8);
|
||||
try map.put(fileId, level);
|
||||
pos += 12;
|
||||
}
|
||||
return map;
|
||||
}
|
||||
|
||||
fn sstName(buf: []u8, id: u64) ![]const u8 {
|
||||
return std.fmt.bufPrint(buf, "sst-{d:0>10}.cbt", .{id});
|
||||
}
|
||||
@@ -1234,6 +1281,67 @@ test "leveled compaction keeps levels disjoint and drops bottom tombstones" {
|
||||
try std.testing.expectEqual(@as(u64, 58), totalEntries);
|
||||
}
|
||||
|
||||
test "seq and fileId counters survive a compaction that garbage-collects every key" {
|
||||
// Reproduces a real recovery gap: when a bottom-level compaction merges
|
||||
// inputs that are entirely tombstones, runMerge produces zero output
|
||||
// files (see the `builder == null` guard around minting a fresh
|
||||
// fileId). If every file that ever carried a high seq/fileId gets
|
||||
// deleted in that round and nothing new gets flushed before a restart,
|
||||
// a directory-scan-only recovery would compute too low a floor and
|
||||
// risk reusing already-consumed seq numbers. The manifest's persisted
|
||||
// counters (VersionEdit.nextFileId/nextSeq, stamped unconditionally in
|
||||
// runMerge) must close that gap.
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
var beforeNextSeq: u64 = 0;
|
||||
var beforeNextFileId: u64 = 0;
|
||||
{
|
||||
var db = try openTemp(&tmp, allocator, .{
|
||||
.l0CompactionTrigger = 2,
|
||||
.l1BaseBytes = 4 * 1024 * 1024,
|
||||
});
|
||||
defer db.close();
|
||||
|
||||
var keyBuf: [16]u8 = undefined;
|
||||
var round: usize = 0;
|
||||
while (round < 2) : (round += 1) {
|
||||
const key = try std.fmt.bufPrint(&keyBuf, "gone:{d}", .{round});
|
||||
try db.put(key, "value");
|
||||
try db.delete(key);
|
||||
try db.flush();
|
||||
}
|
||||
try std.testing.expectEqual(@as(usize, 2), db.countLevel(0));
|
||||
|
||||
try db.maybeCompact();
|
||||
|
||||
// The whole point: every input was a dropped tombstone, so the
|
||||
// compaction round produced no output file at all, and no new
|
||||
// fileId was ever minted for it.
|
||||
try std.testing.expectEqual(@as(usize, 0), db.sstables.items.len);
|
||||
|
||||
beforeNextSeq = db.nextSeq;
|
||||
beforeNextFileId = db.nextFileId;
|
||||
try std.testing.expect(beforeNextSeq > 1);
|
||||
}
|
||||
|
||||
// On disk right now: no sstable files, an empty (reset) WAL. A
|
||||
// scan-only recovery would compute nextSeq=1/nextFileId=1 here.
|
||||
var reopened = try openTemp(&tmp, allocator, .{
|
||||
.l0CompactionTrigger = 2,
|
||||
.l1BaseBytes = 4 * 1024 * 1024,
|
||||
});
|
||||
defer reopened.close();
|
||||
try std.testing.expectEqual(beforeNextSeq, reopened.nextSeq);
|
||||
try std.testing.expectEqual(beforeNextFileId, reopened.nextFileId);
|
||||
|
||||
// Functional proof, not just counter equality: a fresh write after
|
||||
// reopen must not reuse a seq number already spent on the GC'd keys.
|
||||
try reopened.put("fresh", "value");
|
||||
try std.testing.expect(reopened.nextSeq > beforeNextSeq);
|
||||
}
|
||||
|
||||
fn lessByMinKey(_: void, a: *Reader, b: *Reader) bool {
|
||||
return order(a.minKey, b.minKey) == .lt;
|
||||
}
|
||||
@@ -1309,6 +1417,56 @@ test "prefix scan isolates a namespace" {
|
||||
try std.testing.expectEqual(expected.len, idx);
|
||||
}
|
||||
|
||||
test "scan merges a wide fan-in of uncompacted sstables with overlapping keys" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
var db = try openTemp(&tmp, allocator, .{ .flushThreshold = 1 });
|
||||
defer db.close();
|
||||
|
||||
const sstableCount = 14;
|
||||
var round: usize = 0;
|
||||
while (round < sstableCount) : (round += 1) {
|
||||
var keyBuf: [16]u8 = undefined;
|
||||
var valBuf: [24]u8 = undefined;
|
||||
var i: usize = 0;
|
||||
while (i < 5) : (i += 1) {
|
||||
const key = try std.fmt.bufPrint(&keyBuf, "shared:{d:0>2}", .{i});
|
||||
const value = try std.fmt.bufPrint(&valBuf, "round{d}", .{round});
|
||||
try db.put(key, value);
|
||||
}
|
||||
try db.flush();
|
||||
}
|
||||
try db.delete("shared:02");
|
||||
try db.flush();
|
||||
try std.testing.expect(db.sstables.items.len >= sstableCount + 1);
|
||||
|
||||
var forward = try db.scan(allocator, null, null);
|
||||
defer forward.deinit();
|
||||
var seen: usize = 0;
|
||||
while (try forward.next()) |entry| {
|
||||
if (std.mem.eql(u8, entry.key, "shared:02")) {
|
||||
try std.testing.expect(false);
|
||||
}
|
||||
try std.testing.expectEqualStrings("round13", entry.value);
|
||||
seen += 1;
|
||||
}
|
||||
try std.testing.expectEqual(@as(usize, 4), seen);
|
||||
|
||||
var reverse = try db.scanReverse(allocator, .{});
|
||||
defer reverse.deinit();
|
||||
var revSeen: usize = 0;
|
||||
var lastKey: ?[]const u8 = null;
|
||||
while (try reverse.next()) |entry| {
|
||||
if (lastKey) |prev| try std.testing.expect(order(entry.key, prev) == .lt);
|
||||
lastKey = entry.key;
|
||||
try std.testing.expectEqualStrings("round13", entry.value);
|
||||
revSeen += 1;
|
||||
}
|
||||
try std.testing.expectEqual(@as(usize, 4), revSeen);
|
||||
}
|
||||
|
||||
test "inclusive and exclusive bounds" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
|
||||
+596
@@ -0,0 +1,596 @@
|
||||
const std = @import("std");
|
||||
const dbmod = @import("db.zig");
|
||||
const rowmod = @import("row.zig");
|
||||
const querymod = @import("query.zig");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
const Db = dbmod.Db;
|
||||
const Value = rowmod.Value;
|
||||
const Field = rowmod.Field;
|
||||
const Table = querymod.Table;
|
||||
const pkField = querymod.pkField;
|
||||
|
||||
/// Table export/import to CSV and JSON, the two data-interchange formats
|
||||
/// every other tool speaks. Rows are schema-less (`row.zig`'s `Row` is an
|
||||
/// ordered field list, not a fixed struct), so both formats use a shared
|
||||
/// "column set" step: scan once, then take the union of every field name
|
||||
/// seen across every row (first-seen order, primary key first) as the
|
||||
/// header/key set - a row missing a given column just omits that cell/key
|
||||
/// rather than writing an empty or null placeholder, which is what makes
|
||||
/// import's "cell absent" case unambiguous from "cell present but empty
|
||||
/// text" (`s:` with a zero-length payload is a real, distinct encoding).
|
||||
///
|
||||
/// Every cell/value is written as a type-tagged string (`i:`/`f:`/`b:`/
|
||||
/// `s:`/`x:`/`n:`, mirroring `Value.parseCli`'s prefixes plus `x:` for raw
|
||||
/// bytes and `n:` for `.none`) rather than native CSV numbers or JSON
|
||||
/// numbers/booleans - this is what lets both formats share one
|
||||
/// encode/decode pair and round-trip every `ValueTag` exactly, including
|
||||
/// floats (encoded as the exact bit pattern in hex, not a decimal that
|
||||
/// could lose precision on the way back in).
|
||||
fn encodeCell(allocator: Allocator, value: Value) ![]u8 {
|
||||
return switch (value) {
|
||||
.none => try allocator.dupe(u8, "n:"),
|
||||
.boolean => |b| try allocator.dupe(u8, if (b) "b:true" else "b:false"),
|
||||
.int => |v| try std.fmt.allocPrint(allocator, "i:{d}", .{v}),
|
||||
.float => |v| try std.fmt.allocPrint(allocator, "f:{x}", .{@as(u64, @bitCast(v))}),
|
||||
.text => |s| try std.fmt.allocPrint(allocator, "s:{s}", .{s}),
|
||||
.bytes => |b| blk: {
|
||||
var out = std.ArrayList(u8).init(allocator);
|
||||
errdefer out.deinit();
|
||||
try out.appendSlice("x:");
|
||||
for (b) |byte| try out.writer().print("{x:0>2}", .{byte});
|
||||
break :blk out.toOwnedSlice();
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
fn decodeCell(allocator: Allocator, text: []const u8) !Value {
|
||||
if (text.len < 2 or text[1] != ':') return error.InvalidCell;
|
||||
const payload = text[2..];
|
||||
return switch (text[0]) {
|
||||
'n' => .{ .none = {} },
|
||||
'b' => if (std.mem.eql(u8, payload, "true"))
|
||||
.{ .boolean = true }
|
||||
else if (std.mem.eql(u8, payload, "false"))
|
||||
.{ .boolean = false }
|
||||
else
|
||||
error.InvalidCell,
|
||||
'i' => .{ .int = std.fmt.parseInt(i64, payload, 10) catch return error.InvalidCell },
|
||||
'f' => blk: {
|
||||
const bits = std.fmt.parseInt(u64, payload, 16) catch return error.InvalidCell;
|
||||
break :blk .{ .float = @as(f64, @bitCast(bits)) };
|
||||
},
|
||||
's' => .{ .text = try allocator.dupe(u8, payload) },
|
||||
'x' => blk: {
|
||||
if (payload.len % 2 != 0) return error.InvalidCell;
|
||||
const bytes = try allocator.alloc(u8, payload.len / 2);
|
||||
errdefer allocator.free(bytes);
|
||||
var i: usize = 0;
|
||||
while (i < bytes.len) : (i += 1) {
|
||||
bytes[i] = std.fmt.parseInt(u8, payload[i * 2 .. i * 2 + 2], 16) catch return error.InvalidCell;
|
||||
}
|
||||
break :blk .{ .bytes = bytes };
|
||||
},
|
||||
else => error.InvalidCell,
|
||||
};
|
||||
}
|
||||
|
||||
/// Union of every field name across `rows`, primary key first, everything
|
||||
/// else in first-seen order - the shared column list both CSV's header row
|
||||
/// and JSON's per-object keys iterate over.
|
||||
fn collectColumns(allocator: Allocator, rows: []const rowmod.Row) ![][]const u8 {
|
||||
var columns = std.ArrayList([]const u8).init(allocator);
|
||||
errdefer columns.deinit();
|
||||
try columns.append(pkField);
|
||||
for (rows) |r| {
|
||||
for (r.fields.items) |f| {
|
||||
if (std.mem.eql(u8, f.name, pkField)) continue;
|
||||
var seen = false;
|
||||
for (columns.items) |c| {
|
||||
if (std.mem.eql(u8, c, f.name)) {
|
||||
seen = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!seen) try columns.append(f.name);
|
||||
}
|
||||
}
|
||||
return columns.toOwnedSlice();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- CSV -----
|
||||
|
||||
fn writeCsvField(writer: anytype, field: []const u8) !void {
|
||||
const needsQuote = std.mem.indexOfAny(u8, field, ",\"\r\n") != null;
|
||||
if (!needsQuote) {
|
||||
try writer.writeAll(field);
|
||||
return;
|
||||
}
|
||||
try writer.writeByte('"');
|
||||
for (field) |c| {
|
||||
if (c == '"') try writer.writeByte('"');
|
||||
try writer.writeByte(c);
|
||||
}
|
||||
try writer.writeByte('"');
|
||||
}
|
||||
|
||||
pub fn exportCsvWriter(allocator: Allocator, table: *Table, writer: anytype) !usize {
|
||||
const rows = try table.scanAll();
|
||||
defer querymod.freeRows(allocator, rows);
|
||||
|
||||
const columns = try collectColumns(allocator, rows);
|
||||
defer allocator.free(columns);
|
||||
|
||||
for (columns, 0..) |c, i| {
|
||||
if (i > 0) try writer.writeByte(',');
|
||||
try writeCsvField(writer, c);
|
||||
}
|
||||
try writer.writeByte('\n');
|
||||
|
||||
for (rows) |r| {
|
||||
for (columns, 0..) |c, i| {
|
||||
if (i > 0) try writer.writeByte(',');
|
||||
if (r.get(c)) |v| {
|
||||
const cell = try encodeCell(allocator, v);
|
||||
defer allocator.free(cell);
|
||||
try writeCsvField(writer, cell);
|
||||
}
|
||||
}
|
||||
try writer.writeByte('\n');
|
||||
}
|
||||
return rows.len;
|
||||
}
|
||||
|
||||
pub fn exportCsv(allocator: Allocator, table: *Table, path: []const u8) !usize {
|
||||
const file = try std.fs.cwd().createFile(path, .{ .truncate = true });
|
||||
defer file.close();
|
||||
var buffered = std.io.bufferedWriter(file.writer());
|
||||
const n = try exportCsvWriter(allocator, table, buffered.writer());
|
||||
try buffered.flush();
|
||||
return n;
|
||||
}
|
||||
|
||||
/// A minimal RFC4180-shaped reader: comma-separated fields, `"..."`
|
||||
/// quoting with `""` as an escaped quote, records terminated by `\n` or
|
||||
/// `\r\n`. Fields are freshly allocated (never slices into `data`) so a
|
||||
/// quoted field's unescaped content can differ in length from its source
|
||||
/// bytes.
|
||||
const CsvReader = struct {
|
||||
data: []const u8,
|
||||
pos: usize = 0,
|
||||
|
||||
fn atEnd(self: *const CsvReader) bool {
|
||||
return self.pos >= self.data.len;
|
||||
}
|
||||
|
||||
fn readField(self: *CsvReader, allocator: Allocator) ![]u8 {
|
||||
var out = std.ArrayList(u8).init(allocator);
|
||||
errdefer out.deinit();
|
||||
if (!self.atEnd() and self.data[self.pos] == '"') {
|
||||
self.pos += 1;
|
||||
while (true) {
|
||||
if (self.atEnd()) return error.MalformedCsv;
|
||||
const c = self.data[self.pos];
|
||||
if (c == '"') {
|
||||
self.pos += 1;
|
||||
if (!self.atEnd() and self.data[self.pos] == '"') {
|
||||
try out.append('"');
|
||||
self.pos += 1;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
try out.append(c);
|
||||
self.pos += 1;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
while (!self.atEnd()) {
|
||||
const c = self.data[self.pos];
|
||||
if (c == ',' or c == '\n' or c == '\r') break;
|
||||
try out.append(c);
|
||||
self.pos += 1;
|
||||
}
|
||||
}
|
||||
return out.toOwnedSlice();
|
||||
}
|
||||
|
||||
/// Returns the next record (list of owned field slices), or null at
|
||||
/// end of input. Caller owns every field and the returned slice.
|
||||
fn nextRecord(self: *CsvReader, allocator: Allocator) !?[][]u8 {
|
||||
if (self.atEnd()) return null;
|
||||
var fields = std.ArrayList([]u8).init(allocator);
|
||||
errdefer {
|
||||
for (fields.items) |f| allocator.free(f);
|
||||
fields.deinit();
|
||||
}
|
||||
|
||||
while (true) {
|
||||
try fields.append(try self.readField(allocator));
|
||||
if (self.atEnd()) break;
|
||||
const c = self.data[self.pos];
|
||||
if (c == ',') {
|
||||
self.pos += 1;
|
||||
continue;
|
||||
} else if (c == '\n') {
|
||||
self.pos += 1;
|
||||
break;
|
||||
} else if (c == '\r') {
|
||||
self.pos += 1;
|
||||
if (!self.atEnd() and self.data[self.pos] == '\n') self.pos += 1;
|
||||
break;
|
||||
}
|
||||
return error.MalformedCsv;
|
||||
}
|
||||
return try fields.toOwnedSlice();
|
||||
}
|
||||
};
|
||||
|
||||
pub fn importCsvContent(allocator: Allocator, table: *Table, content: []const u8) !usize {
|
||||
var csv = CsvReader{ .data = content };
|
||||
const header = try csv.nextRecord(allocator) orelse return 0;
|
||||
defer {
|
||||
for (header) |h| allocator.free(h);
|
||||
allocator.free(header);
|
||||
}
|
||||
|
||||
var pkIndex: ?usize = null;
|
||||
for (header, 0..) |h, i| {
|
||||
if (std.mem.eql(u8, h, pkField)) {
|
||||
pkIndex = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
const pkIdx = pkIndex orelse return error.MissingPrimaryKey;
|
||||
|
||||
var count: usize = 0;
|
||||
while (try csv.nextRecord(allocator)) |record| {
|
||||
defer {
|
||||
for (record) |f| allocator.free(f);
|
||||
allocator.free(record);
|
||||
}
|
||||
if (record.len != header.len) return error.MalformedCsv;
|
||||
|
||||
var pk: ?Value = null;
|
||||
defer if (pk) |p| p.free(allocator);
|
||||
var fields = std.ArrayList(Field).init(allocator);
|
||||
defer {
|
||||
for (fields.items) |f| f.value.free(allocator);
|
||||
fields.deinit();
|
||||
}
|
||||
|
||||
for (record, 0..) |cell, i| {
|
||||
if (cell.len == 0) continue;
|
||||
const value = try decodeCell(allocator, cell);
|
||||
if (i == pkIdx) {
|
||||
pk = value;
|
||||
} else {
|
||||
try fields.append(.{ .name = header[i], .value = value });
|
||||
}
|
||||
}
|
||||
const pkValue = pk orelse return error.MissingPrimaryKey;
|
||||
try table.put(pkValue, fields.items);
|
||||
count += 1;
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
pub fn importCsv(allocator: Allocator, table: *Table, path: []const u8) !usize {
|
||||
const content = try std.fs.cwd().readFileAlloc(allocator, path, 1 << 30);
|
||||
defer allocator.free(content);
|
||||
return importCsvContent(allocator, table, content);
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------- JSON -----
|
||||
|
||||
pub fn exportJsonWriter(allocator: Allocator, table: *Table, writer: anytype) !usize {
|
||||
const rows = try table.scanAll();
|
||||
defer querymod.freeRows(allocator, rows);
|
||||
|
||||
const columns = try collectColumns(allocator, rows);
|
||||
defer allocator.free(columns);
|
||||
|
||||
var arena = std.heap.ArenaAllocator.init(allocator);
|
||||
defer arena.deinit();
|
||||
const arenaAllocator = arena.allocator();
|
||||
|
||||
var array = std.json.Array.init(arenaAllocator);
|
||||
for (rows) |r| {
|
||||
var object = std.json.ObjectMap.init(arenaAllocator);
|
||||
for (columns) |c| {
|
||||
if (r.get(c)) |v| {
|
||||
const cell = try encodeCell(arenaAllocator, v);
|
||||
try object.put(c, .{ .string = cell });
|
||||
}
|
||||
}
|
||||
try array.append(.{ .object = object });
|
||||
}
|
||||
|
||||
try std.json.stringify(std.json.Value{ .array = array }, .{ .whitespace = .indent_2 }, writer);
|
||||
try writer.writeByte('\n');
|
||||
return rows.len;
|
||||
}
|
||||
|
||||
pub fn exportJson(allocator: Allocator, table: *Table, path: []const u8) !usize {
|
||||
const file = try std.fs.cwd().createFile(path, .{ .truncate = true });
|
||||
defer file.close();
|
||||
var buffered = std.io.bufferedWriter(file.writer());
|
||||
const n = try exportJsonWriter(allocator, table, buffered.writer());
|
||||
try buffered.flush();
|
||||
return n;
|
||||
}
|
||||
|
||||
pub fn importJsonContent(allocator: Allocator, table: *Table, content: []const u8) !usize {
|
||||
var parsed = std.json.parseFromSlice(std.json.Value, allocator, content, .{}) catch return error.MalformedJson;
|
||||
defer parsed.deinit();
|
||||
|
||||
const array = switch (parsed.value) {
|
||||
.array => |a| a,
|
||||
else => return error.MalformedJson,
|
||||
};
|
||||
|
||||
var count: usize = 0;
|
||||
for (array.items) |entry| {
|
||||
const object = switch (entry) {
|
||||
.object => |o| o,
|
||||
else => return error.MalformedJson,
|
||||
};
|
||||
|
||||
var pk: ?Value = null;
|
||||
defer if (pk) |p| p.free(allocator);
|
||||
var fields = std.ArrayList(Field).init(allocator);
|
||||
defer {
|
||||
for (fields.items) |f| f.value.free(allocator);
|
||||
fields.deinit();
|
||||
}
|
||||
|
||||
for (object.keys(), object.values()) |name, jv| {
|
||||
const cell = switch (jv) {
|
||||
.string => |s| s,
|
||||
else => return error.MalformedJson,
|
||||
};
|
||||
const value = try decodeCell(allocator, cell);
|
||||
if (std.mem.eql(u8, name, pkField)) {
|
||||
pk = value;
|
||||
} else {
|
||||
try fields.append(.{ .name = name, .value = value });
|
||||
}
|
||||
}
|
||||
const pkValue = pk orelse return error.MissingPrimaryKey;
|
||||
try table.put(pkValue, fields.items);
|
||||
count += 1;
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
pub fn importJson(allocator: Allocator, table: *Table, path: []const u8) !usize {
|
||||
const content = try std.fs.cwd().readFileAlloc(allocator, path, 1 << 30);
|
||||
defer allocator.free(content);
|
||||
return importJsonContent(allocator, table, content);
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------- tests ----
|
||||
|
||||
const testing = std.testing;
|
||||
|
||||
test "encodeCell decodeCell round trip every value tag" {
|
||||
const allocator = testing.allocator;
|
||||
const values = [_]Value{
|
||||
.{ .none = {} },
|
||||
.{ .boolean = true },
|
||||
.{ .boolean = false },
|
||||
.{ .int = -42 },
|
||||
.{ .int = 0 },
|
||||
.{ .float = 3.5 },
|
||||
.{ .float = -0.001 },
|
||||
.{ .text = "hello, \"world\"\n" },
|
||||
.{ .text = "" },
|
||||
.{ .bytes = &.{ 0, 1, 2, 255 } },
|
||||
};
|
||||
for (values) |v| {
|
||||
const cell = try encodeCell(allocator, v);
|
||||
defer allocator.free(cell);
|
||||
const decoded = try decodeCell(allocator, cell);
|
||||
defer decoded.free(allocator);
|
||||
try testing.expectEqual(@as(rowmod.ValueTag, v), @as(rowmod.ValueTag, decoded));
|
||||
try testing.expect(Value.eql(v, decoded));
|
||||
}
|
||||
}
|
||||
|
||||
test "decodeCell rejects malformed cells" {
|
||||
const allocator = testing.allocator;
|
||||
try testing.expectError(error.InvalidCell, decodeCell(allocator, ""));
|
||||
try testing.expectError(error.InvalidCell, decodeCell(allocator, "x"));
|
||||
try testing.expectError(error.InvalidCell, decodeCell(allocator, "q:1"));
|
||||
try testing.expectError(error.InvalidCell, decodeCell(allocator, "i:notanumber"));
|
||||
try testing.expectError(error.InvalidCell, decodeCell(allocator, "b:maybe"));
|
||||
try testing.expectError(error.InvalidCell, decodeCell(allocator, "x:abc"));
|
||||
try testing.expectError(error.InvalidCell, decodeCell(allocator, "x:zz"));
|
||||
}
|
||||
|
||||
fn openTemp(tmp: *std.testing.TmpDir, allocator: Allocator) !Db {
|
||||
const dbPath = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(dbPath);
|
||||
return Db.open(allocator, dbPath, .{});
|
||||
}
|
||||
|
||||
fn seedTable(table: *Table) !void {
|
||||
try table.put(.{ .int = 1 }, &.{
|
||||
.{ .name = "city", .value = .{ .text = "warsaw" } },
|
||||
.{ .name = "age", .value = .{ .int = 30 } },
|
||||
.{ .name = "active", .value = .{ .boolean = true } },
|
||||
});
|
||||
try table.put(.{ .int = 2 }, &.{
|
||||
.{ .name = "city", .value = .{ .text = "krakow, poland \"old town\"\nsecond line" } },
|
||||
.{ .name = "score", .value = .{ .float = -12.5 } },
|
||||
});
|
||||
try table.put(.{ .int = 3 }, &.{});
|
||||
}
|
||||
|
||||
fn expectRoundTripped(table: *Table) !void {
|
||||
var r1 = (try table.get(.{ .int = 1 })).?;
|
||||
defer r1.deinit();
|
||||
try testing.expectEqualStrings("warsaw", r1.get("city").?.text);
|
||||
try testing.expectEqual(@as(i64, 30), r1.get("age").?.int);
|
||||
try testing.expectEqual(true, r1.get("active").?.boolean);
|
||||
|
||||
var r2 = (try table.get(.{ .int = 2 })).?;
|
||||
defer r2.deinit();
|
||||
try testing.expectEqualStrings("krakow, poland \"old town\"\nsecond line", r2.get("city").?.text);
|
||||
try testing.expectEqual(@as(f64, -12.5), r2.get("score").?.float);
|
||||
try testing.expect(r2.get("age") == null);
|
||||
|
||||
// Row 3 was put with zero extra fields, so its only field is the
|
||||
// primary key itself (`put` always stores that under `pkField`).
|
||||
var r3 = (try table.get(.{ .int = 3 })).?;
|
||||
defer r3.deinit();
|
||||
try testing.expectEqual(@as(usize, 1), r3.count());
|
||||
try testing.expectEqual(@as(i64, 3), r3.get(pkField).?.int);
|
||||
}
|
||||
|
||||
test "csv export then import round trips every row and field exactly" {
|
||||
const allocator = testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
var table = try Table.open(&db, allocator, "people");
|
||||
defer table.deinit();
|
||||
try seedTable(&table);
|
||||
|
||||
var buf = std.ArrayList(u8).init(allocator);
|
||||
defer buf.deinit();
|
||||
const exported = try exportCsvWriter(allocator, &table, buf.writer());
|
||||
try testing.expectEqual(@as(usize, 3), exported);
|
||||
|
||||
var dstTmp = std.testing.tmpDir(.{});
|
||||
defer dstTmp.cleanup();
|
||||
var dstDb = try openTemp(&dstTmp, allocator);
|
||||
defer dstDb.close();
|
||||
var dstTable = try Table.open(&dstDb, allocator, "people");
|
||||
defer dstTable.deinit();
|
||||
|
||||
const imported = try importCsvContent(allocator, &dstTable, buf.items);
|
||||
try testing.expectEqual(@as(usize, 3), imported);
|
||||
|
||||
try expectRoundTripped(&dstTable);
|
||||
}
|
||||
|
||||
test "json export then import round trips every row and field exactly" {
|
||||
const allocator = testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
var table = try Table.open(&db, allocator, "people");
|
||||
defer table.deinit();
|
||||
try seedTable(&table);
|
||||
|
||||
var buf = std.ArrayList(u8).init(allocator);
|
||||
defer buf.deinit();
|
||||
const exported = try exportJsonWriter(allocator, &table, buf.writer());
|
||||
try testing.expectEqual(@as(usize, 3), exported);
|
||||
|
||||
var dstTmp = std.testing.tmpDir(.{});
|
||||
defer dstTmp.cleanup();
|
||||
var dstDb = try openTemp(&dstTmp, allocator);
|
||||
defer dstDb.close();
|
||||
var dstTable = try Table.open(&dstDb, allocator, "people");
|
||||
defer dstTable.deinit();
|
||||
|
||||
const imported = try importJsonContent(allocator, &dstTable, buf.items);
|
||||
try testing.expectEqual(@as(usize, 3), imported);
|
||||
|
||||
try expectRoundTripped(&dstTable);
|
||||
}
|
||||
|
||||
test "csv export of an empty table is just a header row and imports back to zero rows" {
|
||||
const allocator = testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
var table = try Table.open(&db, allocator, "empty");
|
||||
defer table.deinit();
|
||||
|
||||
var buf = std.ArrayList(u8).init(allocator);
|
||||
defer buf.deinit();
|
||||
const exported = try exportCsvWriter(allocator, &table, buf.writer());
|
||||
try testing.expectEqual(@as(usize, 0), exported);
|
||||
try testing.expectEqualStrings("__pk\n", buf.items);
|
||||
|
||||
const imported = try importCsvContent(allocator, &table, buf.items);
|
||||
try testing.expectEqual(@as(usize, 0), imported);
|
||||
}
|
||||
|
||||
test "csv import rejects a header without the primary key column" {
|
||||
const allocator = testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
var table = try Table.open(&db, allocator, "people");
|
||||
defer table.deinit();
|
||||
|
||||
try testing.expectError(error.MissingPrimaryKey, importCsvContent(allocator, &table, "city\ns:warsaw\n"));
|
||||
}
|
||||
|
||||
test "csv import rejects a record whose field count does not match the header" {
|
||||
const allocator = testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
var table = try Table.open(&db, allocator, "people");
|
||||
defer table.deinit();
|
||||
|
||||
try testing.expectError(error.MalformedCsv, importCsvContent(allocator, &table, "__pk,city\ni:1,s:warsaw,s:extra\n"));
|
||||
}
|
||||
|
||||
test "json import rejects a non-array root and a non-string cell value" {
|
||||
const allocator = testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
var table = try Table.open(&db, allocator, "people");
|
||||
defer table.deinit();
|
||||
|
||||
try testing.expectError(error.MalformedJson, importJsonContent(allocator, &table, "{}"));
|
||||
try testing.expectError(error.MalformedJson, importJsonContent(allocator, &table, "[{\"__pk\": 1}]"));
|
||||
}
|
||||
|
||||
test "export csv and json to real files then import from real files round trips" {
|
||||
const allocator = testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
var table = try Table.open(&db, allocator, "people");
|
||||
defer table.deinit();
|
||||
try seedTable(&table);
|
||||
|
||||
const csvPath = "test-dump-people.csv";
|
||||
const jsonPath = "test-dump-people.json";
|
||||
defer std.fs.cwd().deleteFile(csvPath) catch {};
|
||||
defer std.fs.cwd().deleteFile(jsonPath) catch {};
|
||||
|
||||
try testing.expectEqual(@as(usize, 3), try exportCsv(allocator, &table, csvPath));
|
||||
try testing.expectEqual(@as(usize, 3), try exportJson(allocator, &table, jsonPath));
|
||||
|
||||
var dstTmp = std.testing.tmpDir(.{});
|
||||
defer dstTmp.cleanup();
|
||||
var dstDb = try openTemp(&dstTmp, allocator);
|
||||
defer dstDb.close();
|
||||
var csvTable = try Table.open(&dstDb, allocator, "from_csv");
|
||||
defer csvTable.deinit();
|
||||
var jsonTable = try Table.open(&dstDb, allocator, "from_json");
|
||||
defer jsonTable.deinit();
|
||||
|
||||
try testing.expectEqual(@as(usize, 3), try importCsv(allocator, &csvTable, csvPath));
|
||||
try testing.expectEqual(@as(usize, 3), try importJson(allocator, &jsonTable, jsonPath));
|
||||
try expectRoundTripped(&csvTable);
|
||||
try expectRoundTripped(&jsonTable);
|
||||
}
|
||||
+567
@@ -0,0 +1,567 @@
|
||||
const std = @import("std");
|
||||
const rowmod = @import("row.zig");
|
||||
const idx = @import("index.zig");
|
||||
const querymod = @import("query.zig");
|
||||
const dbmod = @import("db.zig");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
const Predicate = querymod.Predicate;
|
||||
const Condition = querymod.Condition;
|
||||
|
||||
/// Errors a caller can act on when a `table select` expression is rejected.
|
||||
/// Parsing never panics on malformed input - every failure mode a user can
|
||||
/// trigger from the CLI maps to one of these.
|
||||
pub const ParseError = error{
|
||||
UnexpectedToken,
|
||||
UnexpectedEnd,
|
||||
InvalidIdentifier,
|
||||
UnknownOperator,
|
||||
InvalidValue,
|
||||
} || Allocator.Error;
|
||||
|
||||
/// Owns every allocation a parsed Predicate tree depends on (field-name and
|
||||
/// text-value copies, and_/or_ slices, not_ boxes) in one arena, so freeing
|
||||
/// is a single `deinit()` regardless of how deep the tree is.
|
||||
pub const Parsed = struct {
|
||||
arena: std.heap.ArenaAllocator,
|
||||
predicate: Predicate,
|
||||
|
||||
pub fn deinit(self: *Parsed) void {
|
||||
self.arena.deinit();
|
||||
}
|
||||
};
|
||||
|
||||
const TokenKind = enum { lparen, rparen, word, eof };
|
||||
|
||||
const Token = struct {
|
||||
kind: TokenKind,
|
||||
text: []const u8,
|
||||
};
|
||||
|
||||
const Lexer = struct {
|
||||
src: []const u8,
|
||||
pos: usize = 0,
|
||||
|
||||
fn next(self: *Lexer) Token {
|
||||
while (self.pos < self.src.len and std.ascii.isWhitespace(self.src[self.pos])) : (self.pos += 1) {}
|
||||
if (self.pos >= self.src.len) return .{ .kind = .eof, .text = "" };
|
||||
|
||||
const c = self.src[self.pos];
|
||||
if (c == '(') {
|
||||
self.pos += 1;
|
||||
return .{ .kind = .lparen, .text = "(" };
|
||||
}
|
||||
if (c == ')') {
|
||||
self.pos += 1;
|
||||
return .{ .kind = .rparen, .text = ")" };
|
||||
}
|
||||
|
||||
const start = self.pos;
|
||||
while (self.pos < self.src.len) {
|
||||
const ch = self.src[self.pos];
|
||||
if (std.ascii.isWhitespace(ch) or ch == '(' or ch == ')') break;
|
||||
self.pos += 1;
|
||||
}
|
||||
return .{ .kind = .word, .text = self.src[start..self.pos] };
|
||||
}
|
||||
};
|
||||
|
||||
fn isKeyword(tok: Token, kw: []const u8) bool {
|
||||
return tok.kind == .word and std.mem.eql(u8, tok.text, kw);
|
||||
}
|
||||
|
||||
/// Recursive-descent parser for the `table select` predicate language:
|
||||
///
|
||||
/// expr := and ("or" and)*
|
||||
/// and := unary ("and" unary)*
|
||||
/// unary := "not" unary | primary
|
||||
/// primary := "(" expr ")" | field op value | field ("is_null" | "is_not_null")
|
||||
/// op := eq | ne | lt | lte | gt | gte | like
|
||||
///
|
||||
/// Precedence (loosest to tightest) is or < and < not, matching every other
|
||||
/// C-family boolean expression language, so `a eq 1 or b eq 2 and not c eq 3`
|
||||
/// parses as `a eq 1 or (b eq 2 and (not (c eq 3)))`.
|
||||
const Parser = struct {
|
||||
lexer: Lexer,
|
||||
allocator: Allocator,
|
||||
cur: Token,
|
||||
|
||||
fn advance(self: *Parser) void {
|
||||
self.cur = self.lexer.next();
|
||||
}
|
||||
|
||||
fn parseOr(self: *Parser) ParseError!Predicate {
|
||||
var list = std.ArrayList(Predicate).init(self.allocator);
|
||||
try list.append(try self.parseAnd());
|
||||
while (isKeyword(self.cur, "or")) {
|
||||
self.advance();
|
||||
try list.append(try self.parseAnd());
|
||||
}
|
||||
if (list.items.len == 1) return list.items[0];
|
||||
return .{ .or_ = try list.toOwnedSlice() };
|
||||
}
|
||||
|
||||
fn parseAnd(self: *Parser) ParseError!Predicate {
|
||||
var list = std.ArrayList(Predicate).init(self.allocator);
|
||||
try list.append(try self.parseUnary());
|
||||
while (isKeyword(self.cur, "and")) {
|
||||
self.advance();
|
||||
try list.append(try self.parseUnary());
|
||||
}
|
||||
if (list.items.len == 1) return list.items[0];
|
||||
return .{ .and_ = try list.toOwnedSlice() };
|
||||
}
|
||||
|
||||
fn parseUnary(self: *Parser) ParseError!Predicate {
|
||||
if (isKeyword(self.cur, "not")) {
|
||||
self.advance();
|
||||
const inner = try self.parseUnary();
|
||||
const boxed = try self.allocator.create(Predicate);
|
||||
boxed.* = inner;
|
||||
return .{ .not_ = boxed };
|
||||
}
|
||||
return self.parsePrimary();
|
||||
}
|
||||
|
||||
fn parsePrimary(self: *Parser) ParseError!Predicate {
|
||||
if (self.cur.kind == .lparen) {
|
||||
self.advance();
|
||||
const inner = try self.parseOr();
|
||||
if (self.cur.kind != .rparen) return error.UnexpectedToken;
|
||||
self.advance();
|
||||
return inner;
|
||||
}
|
||||
return self.parseCondition();
|
||||
}
|
||||
|
||||
fn parseCondition(self: *Parser) ParseError!Predicate {
|
||||
if (self.cur.kind != .word) return if (self.cur.kind == .eof) error.UnexpectedEnd else error.UnexpectedToken;
|
||||
const field = self.cur.text;
|
||||
if (!idx.validIdentifier(field)) return error.InvalidIdentifier;
|
||||
const fieldCopy = try self.allocator.dupe(u8, field);
|
||||
self.advance();
|
||||
|
||||
if (self.cur.kind != .word) return error.UnexpectedEnd;
|
||||
const op = self.cur.text;
|
||||
self.advance();
|
||||
|
||||
if (std.mem.eql(u8, op, "in")) return self.parseIn(fieldCopy);
|
||||
if (std.mem.eql(u8, op, "between")) return self.parseBetween(fieldCopy);
|
||||
if (std.mem.eql(u8, op, "is_null")) return .{ .is_null = fieldCopy };
|
||||
if (std.mem.eql(u8, op, "is_not_null")) return .{ .is_not_null = fieldCopy };
|
||||
|
||||
if (self.cur.kind != .word) return error.UnexpectedEnd;
|
||||
const rawValue = rowmod.Value.parseCli(self.cur.text) catch return error.InvalidValue;
|
||||
const value = try rawValue.dupe(self.allocator);
|
||||
self.advance();
|
||||
|
||||
const cond = Condition{ .field = fieldCopy, .value = value };
|
||||
if (std.mem.eql(u8, op, "eq")) return .{ .eq = cond };
|
||||
if (std.mem.eql(u8, op, "ne")) return .{ .ne = cond };
|
||||
if (std.mem.eql(u8, op, "lt")) return .{ .lt = cond };
|
||||
if (std.mem.eql(u8, op, "lte")) return .{ .lte = cond };
|
||||
if (std.mem.eql(u8, op, "gt")) return .{ .gt = cond };
|
||||
if (std.mem.eql(u8, op, "gte")) return .{ .gte = cond };
|
||||
if (std.mem.eql(u8, op, "like")) {
|
||||
if (value != .text) return error.InvalidValue;
|
||||
return .{ .like = cond };
|
||||
}
|
||||
return error.UnknownOperator;
|
||||
}
|
||||
|
||||
/// `field in (v1 v2 ...)`: desugars to an `or_` of per-value `eq`
|
||||
/// leaves against the same field, sharing the one field-name copy the
|
||||
/// caller already made (Condition never owns `field`, so aliasing it
|
||||
/// across every leaf is safe). A single value collapses to a bare `eq`
|
||||
/// leaf, the same collapse `parseAnd`/`parseOr` already do for a
|
||||
/// one-element and/or chain - lets a one-value `in` use an index like
|
||||
/// `eq` does, since `findIndexableLeaf`/`findCompositeMatch` never look
|
||||
/// inside `or_`. An empty list, `field in ()`, is not a syntax error:
|
||||
/// it desugars to an `or_` over zero leaves, which `Predicate.evaluate`
|
||||
/// already folds to `false` for an empty `preds` slice - matching SQL's
|
||||
/// "IN () is always false" without a special case here.
|
||||
fn parseIn(self: *Parser, field: []const u8) ParseError!Predicate {
|
||||
if (self.cur.kind != .lparen) return if (self.cur.kind == .eof) error.UnexpectedEnd else error.UnexpectedToken;
|
||||
self.advance();
|
||||
|
||||
var list = std.ArrayList(Predicate).init(self.allocator);
|
||||
while (self.cur.kind != .rparen) {
|
||||
if (self.cur.kind != .word) return if (self.cur.kind == .eof) error.UnexpectedEnd else error.UnexpectedToken;
|
||||
const rawValue = rowmod.Value.parseCli(self.cur.text) catch return error.InvalidValue;
|
||||
const value = try rawValue.dupe(self.allocator);
|
||||
self.advance();
|
||||
try list.append(.{ .eq = .{ .field = field, .value = value } });
|
||||
}
|
||||
self.advance();
|
||||
|
||||
if (list.items.len == 1) return list.items[0];
|
||||
return .{ .or_ = try list.toOwnedSlice() };
|
||||
}
|
||||
|
||||
/// `field between lo and hi`: desugars to `field gte lo and field lte
|
||||
/// hi` (inclusive both ends, matching SQL BETWEEN). The `and` here is
|
||||
/// consumed as part of this leaf's own grammar, not the outer
|
||||
/// `parseAnd` binary connective, so it never nests a further `and_`
|
||||
/// level - by the time control returns to the caller, `self.cur` is
|
||||
/// whatever token follows `hi` (rparen/eof/"and"/"or"), decided purely
|
||||
/// by having advanced past every token this leaf owns.
|
||||
fn parseBetween(self: *Parser, field: []const u8) ParseError!Predicate {
|
||||
if (self.cur.kind != .word) return error.UnexpectedEnd;
|
||||
const rawLow = rowmod.Value.parseCli(self.cur.text) catch return error.InvalidValue;
|
||||
const low = try rawLow.dupe(self.allocator);
|
||||
self.advance();
|
||||
|
||||
if (!isKeyword(self.cur, "and")) return if (self.cur.kind == .eof) error.UnexpectedEnd else error.UnexpectedToken;
|
||||
self.advance();
|
||||
|
||||
if (self.cur.kind != .word) return error.UnexpectedEnd;
|
||||
const rawHigh = rowmod.Value.parseCli(self.cur.text) catch return error.InvalidValue;
|
||||
const high = try rawHigh.dupe(self.allocator);
|
||||
self.advance();
|
||||
|
||||
const bounds = try self.allocator.alloc(Predicate, 2);
|
||||
bounds[0] = .{ .gte = .{ .field = field, .value = low } };
|
||||
bounds[1] = .{ .lte = .{ .field = field, .value = high } };
|
||||
return .{ .and_ = bounds };
|
||||
}
|
||||
};
|
||||
|
||||
/// Parses a full `table select` predicate expression, e.g.
|
||||
/// `city eq s:warsaw and (age gte i:30 or not age lt i:10)`. The returned
|
||||
/// Parsed owns every slice/pointer the Predicate tree references - `text`
|
||||
/// itself does not need to outlive the call.
|
||||
pub fn parse(backing: Allocator, text: []const u8) ParseError!Parsed {
|
||||
var arena = std.heap.ArenaAllocator.init(backing);
|
||||
errdefer arena.deinit();
|
||||
const allocator = arena.allocator();
|
||||
|
||||
var p = Parser{ .lexer = .{ .src = text }, .allocator = allocator, .cur = undefined };
|
||||
p.cur = p.lexer.next();
|
||||
|
||||
const pred = try p.parseOr();
|
||||
if (p.cur.kind != .eof) return error.UnexpectedToken;
|
||||
|
||||
return .{ .arena = arena, .predicate = pred };
|
||||
}
|
||||
|
||||
fn evalStr(allocator: Allocator, text: []const u8, row: *const rowmod.Row) !bool {
|
||||
var parsed = try parse(allocator, text);
|
||||
defer parsed.deinit();
|
||||
return parsed.predicate.evaluate(row);
|
||||
}
|
||||
|
||||
test "parses a single condition" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = rowmod.Row.init(allocator);
|
||||
defer row.deinit();
|
||||
try row.set("city", .{ .text = "warsaw" });
|
||||
|
||||
try std.testing.expect(try evalStr(allocator, "city eq s:warsaw", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "city eq s:krakow", &row));
|
||||
}
|
||||
|
||||
test "and binds tighter than or" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = rowmod.Row.init(allocator);
|
||||
defer row.deinit();
|
||||
try row.set("a", .{ .int = 1 });
|
||||
try row.set("b", .{ .int = 2 });
|
||||
try row.set("c", .{ .int = 99 });
|
||||
|
||||
// a eq 1 or b eq 2 and c eq 3 == (a eq 1) or ((b eq 2) and (c eq 3))
|
||||
// b eq 2 is true but c eq 3 is false, so the and-branch is false; the
|
||||
// overall result must still be true because of the leading `a eq 1`.
|
||||
try std.testing.expect(try evalStr(allocator, "a eq i:1 or b eq i:2 and c eq i:3", &row));
|
||||
|
||||
var row2 = rowmod.Row.init(allocator);
|
||||
defer row2.deinit();
|
||||
try row2.set("a", .{ .int = 0 });
|
||||
try row2.set("b", .{ .int = 2 });
|
||||
try row2.set("c", .{ .int = 99 });
|
||||
try std.testing.expect(!try evalStr(allocator, "a eq i:1 or b eq i:2 and c eq i:3", &row2));
|
||||
}
|
||||
|
||||
test "parentheses override default precedence" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = rowmod.Row.init(allocator);
|
||||
defer row.deinit();
|
||||
try row.set("a", .{ .int = 0 });
|
||||
try row.set("b", .{ .int = 2 });
|
||||
try row.set("c", .{ .int = 5 });
|
||||
|
||||
// (a eq 1 or b eq 2) and c eq 5 - forces the or to evaluate before and.
|
||||
try std.testing.expect(try evalStr(allocator, "(a eq i:1 or b eq i:2) and c eq i:5", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "(a eq i:1 or b eq i:2) and c eq i:6", &row));
|
||||
}
|
||||
|
||||
test "not negates the following term" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = rowmod.Row.init(allocator);
|
||||
defer row.deinit();
|
||||
try row.set("age", .{ .int = 20 });
|
||||
|
||||
try std.testing.expect(try evalStr(allocator, "not age eq i:10", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "not age eq i:20", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "not (age eq i:10 or age eq i:20)", &row));
|
||||
}
|
||||
|
||||
test "nested parentheses and multiple not" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = rowmod.Row.init(allocator);
|
||||
defer row.deinit();
|
||||
try row.set("x", .{ .int = 5 });
|
||||
|
||||
try std.testing.expect(try evalStr(allocator, "not not x eq i:5", &row));
|
||||
try std.testing.expect(try evalStr(allocator, "((x eq i:5))", &row));
|
||||
}
|
||||
|
||||
test "range and inequality operators parse and evaluate" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = rowmod.Row.init(allocator);
|
||||
defer row.deinit();
|
||||
try row.set("score", .{ .int = 42 });
|
||||
|
||||
try std.testing.expect(try evalStr(allocator, "score gt i:10", &row));
|
||||
try std.testing.expect(try evalStr(allocator, "score gte i:42", &row));
|
||||
try std.testing.expect(try evalStr(allocator, "score lt i:100", &row));
|
||||
try std.testing.expect(try evalStr(allocator, "score lte i:42", &row));
|
||||
try std.testing.expect(try evalStr(allocator, "score ne i:0", &row));
|
||||
}
|
||||
|
||||
test "like matches SQL-style wildcards and rejects non-text patterns" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = rowmod.Row.init(allocator);
|
||||
defer row.deinit();
|
||||
try row.set("name", .{ .text = "Ada Lovelace" });
|
||||
try row.set("age", .{ .int = 36 });
|
||||
|
||||
try std.testing.expect(try evalStr(allocator, "name like s:Ada%", &row));
|
||||
try std.testing.expect(try evalStr(allocator, "name like s:%Lovelace", &row));
|
||||
try std.testing.expect(try evalStr(allocator, "name like s:%Love%", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "name like s:Boris%", &row));
|
||||
// like against a non-text row field never matches, it does not error.
|
||||
try std.testing.expect(!try evalStr(allocator, "age like s:3%", &row));
|
||||
// like combines with and/or/not/parens exactly like every other operator.
|
||||
try std.testing.expect(try evalStr(allocator, "name like s:Ada% and age eq i:36", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "not name like s:Ada%", &row));
|
||||
|
||||
try std.testing.expectError(error.InvalidValue, parse(allocator, "name like i:5"));
|
||||
}
|
||||
|
||||
test "is_null matches a row missing the field, is_not_null matches a row that has it" {
|
||||
const allocator = std.testing.allocator;
|
||||
var withCity = rowmod.Row.init(allocator);
|
||||
defer withCity.deinit();
|
||||
try withCity.set("city", .{ .text = "warsaw" });
|
||||
|
||||
var withoutCity = rowmod.Row.init(allocator);
|
||||
defer withoutCity.deinit();
|
||||
try withoutCity.set("age", .{ .int = 30 });
|
||||
|
||||
try std.testing.expect(!try evalStr(allocator, "city is_null", &withCity));
|
||||
try std.testing.expect(try evalStr(allocator, "city is_not_null", &withCity));
|
||||
try std.testing.expect(try evalStr(allocator, "city is_null", &withoutCity));
|
||||
try std.testing.expect(!try evalStr(allocator, "city is_not_null", &withoutCity));
|
||||
}
|
||||
|
||||
test "is_null treats an explicit Value.none the same as a missing field" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = rowmod.Row.init(allocator);
|
||||
defer row.deinit();
|
||||
try row.set("nickname", .{ .none = {} });
|
||||
|
||||
try std.testing.expect(try evalStr(allocator, "nickname is_null", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "nickname is_not_null", &row));
|
||||
}
|
||||
|
||||
test "is_null and is_not_null take no trailing value and combine with and/or/not" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = rowmod.Row.init(allocator);
|
||||
defer row.deinit();
|
||||
try row.set("age", .{ .int = 30 });
|
||||
|
||||
try std.testing.expect(try evalStr(allocator, "city is_null and age eq i:30", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "city is_not_null or age eq i:99", &row));
|
||||
try std.testing.expect(try evalStr(allocator, "not city is_not_null", &row));
|
||||
|
||||
var parsed = try parse(allocator, "city is_null");
|
||||
defer parsed.deinit();
|
||||
try std.testing.expect(parsed.predicate == .is_null);
|
||||
try std.testing.expectEqualStrings("city", parsed.predicate.is_null);
|
||||
}
|
||||
|
||||
test "field name copy and value copy outlive the source text" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = rowmod.Row.init(allocator);
|
||||
defer row.deinit();
|
||||
try row.set("name", .{ .text = "ada" });
|
||||
|
||||
var buf: [64]u8 = undefined;
|
||||
const src = try std.fmt.bufPrint(&buf, "name eq s:ada", .{});
|
||||
var parsed = try parse(allocator, src);
|
||||
// Overwrite the source buffer entirely: if the parser had kept slices
|
||||
// into `src` instead of arena-owned copies, this would corrupt the
|
||||
// predicate's field name/value in place.
|
||||
@memset(&buf, 'X');
|
||||
try std.testing.expect(parsed.predicate.evaluate(&row));
|
||||
parsed.deinit();
|
||||
}
|
||||
|
||||
test "rejects malformed expressions" {
|
||||
const allocator = std.testing.allocator;
|
||||
try std.testing.expectError(error.UnexpectedEnd, parse(allocator, "age eq"));
|
||||
try std.testing.expectError(error.UnexpectedEnd, parse(allocator, "age"));
|
||||
try std.testing.expectError(error.UnknownOperator, parse(allocator, "age foo i:1"));
|
||||
try std.testing.expectError(error.InvalidValue, parse(allocator, "age eq i:notanumber"));
|
||||
try std.testing.expectError(error.UnexpectedToken, parse(allocator, "(age eq i:1"));
|
||||
try std.testing.expectError(error.UnexpectedToken, parse(allocator, "age eq i:1)"));
|
||||
try std.testing.expectError(error.UnexpectedEnd, parse(allocator, ""));
|
||||
try std.testing.expectError(error.UnexpectedToken, parse(allocator, "age eq i:1 age eq i:2"));
|
||||
}
|
||||
|
||||
test "in matches any value in the list" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = rowmod.Row.init(allocator);
|
||||
defer row.deinit();
|
||||
try row.set("city", .{ .text = "krakow" });
|
||||
|
||||
try std.testing.expect(try evalStr(allocator, "city in (s:warsaw s:krakow s:gdansk)", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "city in (s:warsaw s:gdansk)", &row));
|
||||
// Empty list is always false, mirroring SQL's `IN ()`.
|
||||
try std.testing.expect(!try evalStr(allocator, "city in ()", &row));
|
||||
// Combines with and/or/not/parens like every other operator.
|
||||
try std.testing.expect(try evalStr(allocator, "city in (s:krakow) and not city eq s:warsaw", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "not city in (s:warsaw s:krakow)", &row));
|
||||
}
|
||||
|
||||
test "in with a single value collapses to a bare eq leaf" {
|
||||
const allocator = std.testing.allocator;
|
||||
var parsed = try parse(allocator, "city in (s:warsaw)");
|
||||
defer parsed.deinit();
|
||||
try std.testing.expect(parsed.predicate == .eq);
|
||||
try std.testing.expectEqualStrings("city", parsed.predicate.eq.field);
|
||||
}
|
||||
|
||||
test "between is inclusive on both ends" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = rowmod.Row.init(allocator);
|
||||
defer row.deinit();
|
||||
try row.set("age", .{ .int = 30 });
|
||||
|
||||
try std.testing.expect(try evalStr(allocator, "age between i:30 and i:40", &row));
|
||||
try std.testing.expect(try evalStr(allocator, "age between i:20 and i:30", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "age between i:31 and i:40", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "age between i:10 and i:29", &row));
|
||||
// Combines with and/or/not/parens, and the leaf's own "and" does not
|
||||
// leak into the outer and/or grammar as a second top-level conjunct.
|
||||
try std.testing.expect(try evalStr(allocator, "age between i:25 and i:35 and age ne i:0", &row));
|
||||
try std.testing.expect(!try evalStr(allocator, "not age between i:25 and i:35", &row));
|
||||
}
|
||||
|
||||
test "between desugars to a two-leaf and_ of gte/lte" {
|
||||
const allocator = std.testing.allocator;
|
||||
var parsed = try parse(allocator, "age between i:1 and i:2");
|
||||
defer parsed.deinit();
|
||||
try std.testing.expect(parsed.predicate == .and_);
|
||||
try std.testing.expectEqual(@as(usize, 2), parsed.predicate.and_.len);
|
||||
try std.testing.expect(parsed.predicate.and_[0] == .gte);
|
||||
try std.testing.expect(parsed.predicate.and_[1] == .lte);
|
||||
}
|
||||
|
||||
test "in and between reject malformed input" {
|
||||
const allocator = std.testing.allocator;
|
||||
try std.testing.expectError(error.UnexpectedToken, parse(allocator, "city in s:warsaw"));
|
||||
try std.testing.expectError(error.UnexpectedEnd, parse(allocator, "city in (s:warsaw"));
|
||||
try std.testing.expectError(error.InvalidValue, parse(allocator, "age in (i:notanumber)"));
|
||||
try std.testing.expectError(error.UnexpectedEnd, parse(allocator, "age between i:1"));
|
||||
try std.testing.expectError(error.UnexpectedToken, parse(allocator, "age between i:1 or i:2"));
|
||||
try std.testing.expectError(error.InvalidValue, parse(allocator, "age between i:1 and i:notanumber"));
|
||||
}
|
||||
|
||||
fn openTemp(tmp: *std.testing.TmpDir, allocator: Allocator) !dbmod.Db {
|
||||
const dbPath = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(dbPath);
|
||||
return dbmod.Db.open(allocator, dbPath, .{});
|
||||
}
|
||||
|
||||
test "a single-value in reaches Table.select through an index like eq does" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
var table = try querymod.Table.open(&db, allocator, "cities");
|
||||
defer table.deinit();
|
||||
try table.createIndex("city");
|
||||
try table.put(.{ .int = 1 }, &.{.{ .name = "city", .value = .{ .text = "oslo" } }});
|
||||
try table.put(.{ .int = 2 }, &.{.{ .name = "city", .value = .{ .text = "bergen" } }});
|
||||
|
||||
var parsed = try parse(allocator, "city in (s:oslo)");
|
||||
defer parsed.deinit();
|
||||
|
||||
var plan = try table.explain(allocator, parsed.predicate);
|
||||
defer plan.deinit(allocator);
|
||||
try std.testing.expectEqual(querymod.PlanKind.index_eq, plan.kind);
|
||||
try std.testing.expectEqualStrings("city", plan.indexName.?);
|
||||
|
||||
const results = try table.select(parsed.predicate);
|
||||
defer querymod.freeRows(allocator, results);
|
||||
try std.testing.expectEqual(@as(usize, 1), results.len);
|
||||
try std.testing.expectEqual(@as(i64, 1), results[0].get(querymod.pkField).?.int);
|
||||
}
|
||||
|
||||
test "a multi-value in falls back to a full scan and still matches every listed value" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
var table = try querymod.Table.open(&db, allocator, "cities");
|
||||
defer table.deinit();
|
||||
try table.createIndex("city");
|
||||
try table.put(.{ .int = 1 }, &.{.{ .name = "city", .value = .{ .text = "oslo" } }});
|
||||
try table.put(.{ .int = 2 }, &.{.{ .name = "city", .value = .{ .text = "bergen" } }});
|
||||
try table.put(.{ .int = 3 }, &.{.{ .name = "city", .value = .{ .text = "trondheim" } }});
|
||||
|
||||
var parsed = try parse(allocator, "city in (s:oslo s:trondheim)");
|
||||
defer parsed.deinit();
|
||||
|
||||
var plan = try table.explain(allocator, parsed.predicate);
|
||||
defer plan.deinit(allocator);
|
||||
try std.testing.expectEqual(querymod.PlanKind.full_scan, plan.kind);
|
||||
|
||||
const results = try table.select(parsed.predicate);
|
||||
defer querymod.freeRows(allocator, results);
|
||||
try std.testing.expectEqual(@as(usize, 2), results.len);
|
||||
}
|
||||
|
||||
test "between reaches Table.select through an index range like gte/lte does" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
var table = try querymod.Table.open(&db, allocator, "people");
|
||||
defer table.deinit();
|
||||
try table.createIndex("age");
|
||||
var i: i64 = 0;
|
||||
while (i < 10) : (i += 1) {
|
||||
try table.put(.{ .int = i }, &.{.{ .name = "age", .value = .{ .int = i * 10 } }});
|
||||
}
|
||||
|
||||
var parsed = try parse(allocator, "age between i:20 and i:50");
|
||||
defer parsed.deinit();
|
||||
|
||||
var plan = try table.explain(allocator, parsed.predicate);
|
||||
defer plan.deinit(allocator);
|
||||
try std.testing.expectEqual(querymod.PlanKind.index_range, plan.kind);
|
||||
try std.testing.expectEqualStrings("age", plan.indexName.?);
|
||||
|
||||
const results = try table.select(parsed.predicate);
|
||||
defer querymod.freeRows(allocator, results);
|
||||
// ages 20, 30, 40, 50 - inclusive both ends.
|
||||
try std.testing.expectEqual(@as(usize, 4), results.len);
|
||||
}
|
||||
+265
@@ -0,0 +1,265 @@
|
||||
const std = @import("std");
|
||||
const Allocator = std.mem.Allocator;
|
||||
|
||||
/// Binary heap ordered by `lessThan(context, a, b)`: parent never compares
|
||||
/// greater than either child under that relation. Pass a comparator that
|
||||
/// returns `a > b` to get a max-heap out of the same implementation.
|
||||
pub fn MinHeap(comptime T: type, comptime Context: type, comptime lessThan: fn (context: Context, a: T, b: T) bool) type {
|
||||
return struct {
|
||||
items: std.ArrayList(T),
|
||||
context: Context,
|
||||
|
||||
const Self = @This();
|
||||
|
||||
pub fn init(allocator: Allocator, context: Context) Self {
|
||||
return .{ .items = std.ArrayList(T).init(allocator), .context = context };
|
||||
}
|
||||
|
||||
pub fn deinit(self: *Self) void {
|
||||
self.items.deinit();
|
||||
}
|
||||
|
||||
pub fn len(self: *const Self) usize {
|
||||
return self.items.items.len;
|
||||
}
|
||||
|
||||
pub fn peek(self: *const Self) ?T {
|
||||
if (self.items.items.len == 0) return null;
|
||||
return self.items.items[0];
|
||||
}
|
||||
|
||||
pub fn push(self: *Self, item: T) !void {
|
||||
try self.items.append(item);
|
||||
self.siftUp(self.items.items.len - 1);
|
||||
}
|
||||
|
||||
pub fn pop(self: *Self) ?T {
|
||||
if (self.items.items.len == 0) return null;
|
||||
const top = self.items.items[0];
|
||||
const last = self.items.pop();
|
||||
if (self.items.items.len > 0) {
|
||||
self.items.items[0] = last;
|
||||
self.siftDown(0);
|
||||
}
|
||||
return top;
|
||||
}
|
||||
|
||||
/// Swap the root for `item` and restore the heap property in one
|
||||
/// siftDown, instead of a pop() (siftDown) followed by a push()
|
||||
/// (siftUp). Caller must ensure the heap is non-empty.
|
||||
pub fn replaceTop(self: *Self, item: T) T {
|
||||
const old = self.items.items[0];
|
||||
self.items.items[0] = item;
|
||||
self.siftDown(0);
|
||||
return old;
|
||||
}
|
||||
|
||||
fn siftUp(self: *Self, startIdx: usize) void {
|
||||
var idx = startIdx;
|
||||
while (idx > 0) {
|
||||
const parent = (idx - 1) / 2;
|
||||
if (!lessThan(self.context, self.items.items[idx], self.items.items[parent])) break;
|
||||
std.mem.swap(T, &self.items.items[idx], &self.items.items[parent]);
|
||||
idx = parent;
|
||||
}
|
||||
}
|
||||
|
||||
fn siftDown(self: *Self, startIdx: usize) void {
|
||||
var idx = startIdx;
|
||||
const n = self.items.items.len;
|
||||
while (true) {
|
||||
const left = idx * 2 + 1;
|
||||
const right = idx * 2 + 2;
|
||||
var smallest = idx;
|
||||
if (left < n and lessThan(self.context, self.items.items[left], self.items.items[smallest])) smallest = left;
|
||||
if (right < n and lessThan(self.context, self.items.items[right], self.items.items[smallest])) smallest = right;
|
||||
if (smallest == idx) break;
|
||||
std.mem.swap(T, &self.items.items[idx], &self.items.items[smallest]);
|
||||
idx = smallest;
|
||||
}
|
||||
}
|
||||
|
||||
/// Debug helper: verifies the heap invariant holds at every node.
|
||||
pub fn isValid(self: *const Self) bool {
|
||||
const n = self.items.items.len;
|
||||
var idx: usize = 0;
|
||||
while (idx < n) : (idx += 1) {
|
||||
const left = idx * 2 + 1;
|
||||
const right = idx * 2 + 2;
|
||||
if (left < n and lessThan(self.context, self.items.items[left], self.items.items[idx])) return false;
|
||||
if (right < n and lessThan(self.context, self.items.items[right], self.items.items[idx])) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
fn u32Less(_: void, a: u32, b: u32) bool {
|
||||
return a < b;
|
||||
}
|
||||
const U32Heap = MinHeap(u32, void, u32Less);
|
||||
|
||||
test "heap: empty peek/pop return null" {
|
||||
var heap = U32Heap.init(std.testing.allocator, {});
|
||||
defer heap.deinit();
|
||||
try std.testing.expectEqual(@as(?u32, null), heap.peek());
|
||||
try std.testing.expectEqual(@as(?u32, null), heap.pop());
|
||||
try std.testing.expectEqual(@as(usize, 0), heap.len());
|
||||
}
|
||||
|
||||
test "heap: push/pop yields ascending order" {
|
||||
var heap = U32Heap.init(std.testing.allocator, {});
|
||||
defer heap.deinit();
|
||||
const values = [_]u32{ 5, 1, 9, 3, 3, 7, 0, 42, 2 };
|
||||
for (values) |v| try heap.push(v);
|
||||
try std.testing.expectEqual(values.len, heap.len());
|
||||
|
||||
var sorted = values;
|
||||
std.mem.sort(u32, &sorted, {}, std.sort.asc(u32));
|
||||
|
||||
for (sorted) |expected| {
|
||||
try std.testing.expectEqual(@as(?u32, expected), heap.peek());
|
||||
try std.testing.expectEqual(@as(?u32, expected), heap.pop());
|
||||
}
|
||||
try std.testing.expectEqual(@as(usize, 0), heap.len());
|
||||
}
|
||||
|
||||
test "heap: peek does not remove" {
|
||||
var heap = U32Heap.init(std.testing.allocator, {});
|
||||
defer heap.deinit();
|
||||
try heap.push(10);
|
||||
try heap.push(4);
|
||||
try std.testing.expectEqual(@as(?u32, 4), heap.peek());
|
||||
try std.testing.expectEqual(@as(?u32, 4), heap.peek());
|
||||
try std.testing.expectEqual(@as(usize, 2), heap.len());
|
||||
}
|
||||
|
||||
test "heap: replaceTop matches pop-then-push" {
|
||||
var a = U32Heap.init(std.testing.allocator, {});
|
||||
defer a.deinit();
|
||||
var b = U32Heap.init(std.testing.allocator, {});
|
||||
defer b.deinit();
|
||||
for ([_]u32{ 8, 2, 6, 1, 9 }) |v| {
|
||||
try a.push(v);
|
||||
try b.push(v);
|
||||
}
|
||||
|
||||
const old = a.replaceTop(5);
|
||||
try std.testing.expectEqual(@as(u32, 1), old);
|
||||
_ = b.pop();
|
||||
try b.push(5);
|
||||
|
||||
while (a.len() > 0) {
|
||||
try std.testing.expectEqual(a.pop().?, b.pop().?);
|
||||
}
|
||||
}
|
||||
|
||||
test "heap: random stress matches sort reference and stays valid" {
|
||||
var prng = std.Random.DefaultPrng.init(0xC0BA17);
|
||||
const random = prng.random();
|
||||
var heap = U32Heap.init(std.testing.allocator, {});
|
||||
defer heap.deinit();
|
||||
|
||||
var reference = std.ArrayList(u32).init(std.testing.allocator);
|
||||
defer reference.deinit();
|
||||
|
||||
var i: usize = 0;
|
||||
while (i < 500) : (i += 1) {
|
||||
const v = random.intRangeAtMost(u32, 0, 1000);
|
||||
try heap.push(v);
|
||||
try reference.append(v);
|
||||
try std.testing.expect(heap.isValid());
|
||||
}
|
||||
|
||||
std.mem.sort(u32, reference.items, {}, std.sort.asc(u32));
|
||||
for (reference.items) |expected| {
|
||||
try std.testing.expectEqual(@as(?u32, expected), heap.pop());
|
||||
try std.testing.expect(heap.isValid());
|
||||
}
|
||||
try std.testing.expectEqual(@as(usize, 0), heap.len());
|
||||
}
|
||||
|
||||
test "heap: interleaved push and pop stays valid" {
|
||||
var prng = std.Random.DefaultPrng.init(0xFEED);
|
||||
const random = prng.random();
|
||||
var heap = U32Heap.init(std.testing.allocator, {});
|
||||
defer heap.deinit();
|
||||
|
||||
var count: usize = 0;
|
||||
var i: usize = 0;
|
||||
while (i < 2000) : (i += 1) {
|
||||
if (count == 0 or random.boolean()) {
|
||||
try heap.push(random.intRangeAtMost(u32, 0, 100));
|
||||
count += 1;
|
||||
} else {
|
||||
const popped = heap.pop();
|
||||
try std.testing.expect(popped != null);
|
||||
count -= 1;
|
||||
}
|
||||
try std.testing.expect(heap.isValid());
|
||||
try std.testing.expectEqual(count, heap.len());
|
||||
}
|
||||
}
|
||||
|
||||
fn descU32(_: void, a: u32, b: u32) bool {
|
||||
return a > b;
|
||||
}
|
||||
const MaxU32Heap = MinHeap(u32, void, descU32);
|
||||
|
||||
test "heap: inverted comparator behaves as a max-heap" {
|
||||
var heap = MaxU32Heap.init(std.testing.allocator, {});
|
||||
defer heap.deinit();
|
||||
for ([_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 }) |v| try heap.push(v);
|
||||
|
||||
var sorted = [_]u32{ 3, 1, 4, 1, 5, 9, 2, 6 };
|
||||
std.mem.sort(u32, &sorted, {}, std.sort.desc(u32));
|
||||
for (sorted) |expected| try std.testing.expectEqual(@as(?u32, expected), heap.pop());
|
||||
}
|
||||
|
||||
const Weighted = struct { key: []const u8, seq: u32 };
|
||||
|
||||
fn keyLess(_: void, a: Weighted, b: Weighted) bool {
|
||||
return std.mem.order(u8, a.key, b.key) == .lt;
|
||||
}
|
||||
const KeyHeap = MinHeap(Weighted, void, keyLess);
|
||||
|
||||
test "heap: k-way merge of pre-sorted lists matches full sort" {
|
||||
const allocator = std.testing.allocator;
|
||||
const lists = [_][]const []const u8{
|
||||
&[_][]const u8{ "apple", "cherry", "fig", "kiwi" },
|
||||
&[_][]const u8{ "banana", "date", "eggplant" },
|
||||
&[_][]const u8{"lemon"},
|
||||
&[_][]const u8{},
|
||||
};
|
||||
|
||||
const Cursor = struct { list: []const []const u8, idx: usize };
|
||||
var cursors: [lists.len]Cursor = undefined;
|
||||
for (lists, 0..) |list, i| cursors[i] = .{ .list = list, .idx = 0 };
|
||||
|
||||
var heap = KeyHeap.init(allocator, {});
|
||||
defer heap.deinit();
|
||||
for (cursors, 0..) |c, i| {
|
||||
if (c.idx < c.list.len) try heap.push(.{ .key = c.list[c.idx], .seq = @intCast(i) });
|
||||
}
|
||||
|
||||
var merged = std.ArrayList([]const u8).init(allocator);
|
||||
defer merged.deinit();
|
||||
while (heap.pop()) |top| {
|
||||
try merged.append(top.key);
|
||||
var c = &cursors[top.seq];
|
||||
c.idx += 1;
|
||||
if (c.idx < c.list.len) try heap.push(.{ .key = c.list[c.idx], .seq = top.seq });
|
||||
}
|
||||
|
||||
var expected = std.ArrayList([]const u8).init(allocator);
|
||||
defer expected.deinit();
|
||||
for (lists) |list| try expected.appendSlice(list);
|
||||
std.mem.sort([]const u8, expected.items, {}, struct {
|
||||
fn lt(_: void, a: []const u8, b: []const u8) bool {
|
||||
return std.mem.order(u8, a, b) == .lt;
|
||||
}
|
||||
}.lt);
|
||||
|
||||
try std.testing.expectEqual(expected.items.len, merged.items.len);
|
||||
for (expected.items, merged.items) |e, m| try std.testing.expectEqualStrings(e, m);
|
||||
}
|
||||
+1937
File diff suppressed because it is too large.
Load diff
+1158
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,96 @@
|
||||
const std = @import("std");
|
||||
const auth = @import("auth.zig");
|
||||
const index = @import("index.zig");
|
||||
const collections = @import("collections.zig");
|
||||
const collttl = @import("collttl.zig");
|
||||
const ttl = @import("ttl.zig");
|
||||
|
||||
/// Single source of truth for every top-level byte tag a module uses to
|
||||
/// namespace its own keys in the flat `Db` keyspace (auth.zig/index.zig/
|
||||
/// collections.zig/collttl.zig/ttl.zig each own one plain byte as the
|
||||
/// first byte of every key they write, and rely on every OTHER module
|
||||
/// picking a different one - the entire scheme is "gentleman's agreement
|
||||
/// by comment" with nothing that actually enforces it).
|
||||
///
|
||||
/// Session #29 hit the failure mode this closes: collttl.zig originally
|
||||
/// reused ttl.zig's tag 'T' (only byte-disjoint via an extra marker byte,
|
||||
/// not prefix-disjoint), so ttl.zig's `sweepExpired` prefix-scan over "T"
|
||||
/// silently matched collttl.zig's markers too, corrupting sweep counts and
|
||||
/// leaving collections un-swept. It was caught by a flaky-looking test
|
||||
/// failure, not by anything structural - nothing stopped a *third* module
|
||||
/// from making the same mistake tomorrow.
|
||||
///
|
||||
/// Each entry below points at the module's own `pub const ...Tag`
|
||||
/// definition instead of retyping the byte, so this list is derived FROM
|
||||
/// the real definitions and cannot silently drift out of sync with them.
|
||||
/// `assertAllDistinct` runs at comptime (see the `comptime` block below,
|
||||
/// pulled in transitively by every build via root.zig) - a future module
|
||||
/// that picks an already-taken byte fails the BUILD, not a test someone
|
||||
/// has to remember to write.
|
||||
pub const Entry = struct { owner: []const u8, tag: u8 };
|
||||
|
||||
pub const registry = [_]Entry{
|
||||
.{ .owner = "auth", .tag = auth.nsTag },
|
||||
.{ .owner = "index.row", .tag = index.rowTag },
|
||||
.{ .owner = "index.index", .tag = index.indexTag },
|
||||
.{ .owner = "index.compositeIndex", .tag = index.compositeIndexTag },
|
||||
.{ .owner = "index.meta", .tag = index.metaTag },
|
||||
.{ .owner = "collections", .tag = collections.nsTag },
|
||||
.{ .owner = "collttl", .tag = collttl.nsTag },
|
||||
.{ .owner = "ttl", .tag = ttl.nsTag },
|
||||
};
|
||||
|
||||
fn assertAllDistinct() void {
|
||||
@setEvalBranchQuota(10_000);
|
||||
for (registry, 0..) |a, i| {
|
||||
for (registry[i + 1 ..]) |b| {
|
||||
if (a.tag == b.tag) {
|
||||
@compileError("keyspace tag collision: '" ++ &[_]u8{a.tag} ++ "' claimed by both '" ++ a.owner ++ "' and '" ++ b.owner ++ "' - pick a different top-level byte for one of them");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
comptime {
|
||||
assertAllDistinct();
|
||||
}
|
||||
|
||||
/// `true` iff no other registered owner already claims `tag` - the check a
|
||||
/// brand new module should run against its OWN candidate byte before
|
||||
/// wiring it into `registry` above, so the collision surfaces while
|
||||
/// picking the byte instead of after writing a whole module around it.
|
||||
pub fn isFree(tag: u8) bool {
|
||||
for (registry) |e| {
|
||||
if (e.tag == tag) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
test "registry has no duplicate top-level tags (runtime mirror of the comptime check)" {
|
||||
for (registry, 0..) |a, i| {
|
||||
for (registry[i + 1 ..]) |b| {
|
||||
try std.testing.expect(a.tag != b.tag);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
test "registry pins the exact set of tags currently in use" {
|
||||
var seen = std.AutoHashMap(u8, void).init(std.testing.allocator);
|
||||
defer seen.deinit();
|
||||
for (registry) |e| try seen.put(e.tag, {});
|
||||
// A change here (tag added, removed, or reassigned) must be a
|
||||
// deliberate edit to this file's `registry` literal, not something
|
||||
// that happened to fall out of touching an unrelated module.
|
||||
try std.testing.expectEqual(@as(usize, 8), seen.count());
|
||||
inline for (.{ 'A', 'R', 'I', 'J', 'M', 'X', 'Y', 'T' }) |expected| {
|
||||
try std.testing.expect(seen.contains(expected));
|
||||
}
|
||||
}
|
||||
|
||||
test "isFree reports the registered tags as taken and an unused byte as free" {
|
||||
try std.testing.expect(!isFree('A'));
|
||||
try std.testing.expect(!isFree('T'));
|
||||
try std.testing.expect(!isFree('Y'));
|
||||
try std.testing.expect(isFree('Z'));
|
||||
try std.testing.expect(isFree(0));
|
||||
}
|
||||
+3545
-5
File diff suppressed because it is too large.
Load diff
+600
@@ -0,0 +1,600 @@
|
||||
const std = @import("std");
|
||||
const stats_mod = @import("stats.zig");
|
||||
|
||||
// Metrics primitives (counter/gauge/histogram), two independent rate-limiter
|
||||
// algorithms (token bucket, sliding window counter), and a text-exposition
|
||||
// registry. Fully additive module: does not touch Db/txn/storage internals,
|
||||
// only depends on stats.zig's public Stats type for the export helper.
|
||||
|
||||
pub const MetricCounter = struct {
|
||||
value: std.atomic.Value(u64) = std.atomic.Value(u64).init(0),
|
||||
|
||||
pub fn incr(self: *MetricCounter) void {
|
||||
self.incrBy(1);
|
||||
}
|
||||
|
||||
pub fn incrBy(self: *MetricCounter, delta: u64) void {
|
||||
_ = self.value.fetchAdd(delta, .monotonic);
|
||||
}
|
||||
|
||||
pub fn get(self: *const MetricCounter) u64 {
|
||||
return self.value.load(.monotonic);
|
||||
}
|
||||
|
||||
pub fn reset(self: *MetricCounter) void {
|
||||
self.value.store(0, .monotonic);
|
||||
}
|
||||
};
|
||||
|
||||
pub const Gauge = struct {
|
||||
bits: std.atomic.Value(u64) = std.atomic.Value(u64).init(0),
|
||||
|
||||
pub fn set(self: *Gauge, v: f64) void {
|
||||
self.bits.store(@bitCast(v), .monotonic);
|
||||
}
|
||||
|
||||
pub fn get(self: *const Gauge) f64 {
|
||||
const raw = self.bits.load(.monotonic);
|
||||
return @bitCast(raw);
|
||||
}
|
||||
|
||||
pub fn add(self: *Gauge, delta: f64) void {
|
||||
while (true) {
|
||||
const cur = self.bits.load(.monotonic);
|
||||
const curF: f64 = @bitCast(cur);
|
||||
const newBits: u64 = @bitCast(curF + delta);
|
||||
if (self.bits.cmpxchgWeak(cur, newBits, .monotonic, .monotonic) == null) return;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
pub const HistogramError = error{EmptyBounds};
|
||||
|
||||
/// Fixed-bucket histogram with Prometheus-style cumulative-bucket quantile
|
||||
/// estimation. `bounds` must be strictly increasing and is borrowed (caller
|
||||
/// keeps it alive for the histogram's lifetime); values greater than the
|
||||
/// last bound land in an implicit +Inf bucket.
|
||||
pub const Histogram = struct {
|
||||
bounds: []const f64,
|
||||
counts: []u64,
|
||||
sum: f64 = 0,
|
||||
total: u64 = 0,
|
||||
mutex: std.Thread.Mutex = .{},
|
||||
allocator: std.mem.Allocator,
|
||||
|
||||
pub fn init(allocator: std.mem.Allocator, bounds: []const f64) !Histogram {
|
||||
if (bounds.len == 0) return HistogramError.EmptyBounds;
|
||||
var i: usize = 1;
|
||||
while (i < bounds.len) : (i += 1) {
|
||||
std.debug.assert(bounds[i] > bounds[i - 1]);
|
||||
}
|
||||
const counts = try allocator.alloc(u64, bounds.len + 1);
|
||||
@memset(counts, 0);
|
||||
return .{ .bounds = bounds, .counts = counts, .allocator = allocator };
|
||||
}
|
||||
|
||||
pub fn deinit(self: *Histogram) void {
|
||||
self.allocator.free(self.counts);
|
||||
self.* = undefined;
|
||||
}
|
||||
|
||||
pub fn observe(self: *Histogram, value: f64) void {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
var idx: usize = self.bounds.len;
|
||||
for (self.bounds, 0..) |bound, i| {
|
||||
if (value <= bound) {
|
||||
idx = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
self.counts[idx] += 1;
|
||||
self.sum += value;
|
||||
self.total += 1;
|
||||
}
|
||||
|
||||
pub fn count(self: *const Histogram) u64 {
|
||||
return self.total;
|
||||
}
|
||||
|
||||
pub fn mean(self: *const Histogram) f64 {
|
||||
if (self.total == 0) return 0;
|
||||
return self.sum / @as(f64, @floatFromInt(self.total));
|
||||
}
|
||||
|
||||
/// Linear-interpolation quantile estimate over cumulative bucket counts,
|
||||
/// same algorithm Prometheus uses for `histogram_quantile`. q in [0, 1].
|
||||
pub fn quantile(self: *Histogram, q: f64) f64 {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
if (self.total == 0 or self.bounds.len == 0) return 0;
|
||||
const clamped = std.math.clamp(q, 0, 1);
|
||||
const target = clamped * @as(f64, @floatFromInt(self.total));
|
||||
|
||||
var cumulative: u64 = 0;
|
||||
var i: usize = 0;
|
||||
while (i < self.counts.len) : (i += 1) {
|
||||
const prevCumulative = cumulative;
|
||||
cumulative += self.counts[i];
|
||||
if (@as(f64, @floatFromInt(cumulative)) >= target) {
|
||||
if (i == 0) return self.bounds[0];
|
||||
if (i == self.bounds.len) return self.bounds[self.bounds.len - 1];
|
||||
const lower = self.bounds[i - 1];
|
||||
const upper = self.bounds[i];
|
||||
const countInBucket = cumulative - prevCumulative;
|
||||
if (countInBucket == 0) return upper;
|
||||
const fraction = (target - @as(f64, @floatFromInt(prevCumulative))) / @as(f64, @floatFromInt(countInBucket));
|
||||
return lower + fraction * (upper - lower);
|
||||
}
|
||||
}
|
||||
return self.bounds[self.bounds.len - 1];
|
||||
}
|
||||
|
||||
pub fn bucketCount(self: *Histogram, idx: usize) u64 {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
return self.counts[idx];
|
||||
}
|
||||
};
|
||||
|
||||
pub fn linearBuckets(allocator: std.mem.Allocator, start: f64, width: f64, n: usize) ![]f64 {
|
||||
std.debug.assert(n > 0);
|
||||
std.debug.assert(width > 0);
|
||||
const out = try allocator.alloc(f64, n);
|
||||
var i: usize = 0;
|
||||
while (i < n) : (i += 1) out[i] = start + width * @as(f64, @floatFromInt(i));
|
||||
return out;
|
||||
}
|
||||
|
||||
pub fn exponentialBuckets(allocator: std.mem.Allocator, start: f64, factor: f64, n: usize) ![]f64 {
|
||||
std.debug.assert(n > 0);
|
||||
std.debug.assert(start > 0);
|
||||
std.debug.assert(factor > 1);
|
||||
const out = try allocator.alloc(f64, n);
|
||||
var cur = start;
|
||||
var i: usize = 0;
|
||||
while (i < n) : (i += 1) {
|
||||
out[i] = cur;
|
||||
cur *= factor;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/// Classic token-bucket rate limiter. Time is passed in explicitly as
|
||||
/// nanoseconds so the algorithm is deterministic and testable; `tryAcquireNow`
|
||||
/// wraps it with the wall clock for real use.
|
||||
pub const TokenBucket = struct {
|
||||
capacity: f64,
|
||||
refillPerSecond: f64,
|
||||
tokens: f64,
|
||||
lastRefillNanos: u64,
|
||||
mutex: std.Thread.Mutex = .{},
|
||||
|
||||
pub fn init(capacity: f64, refillPerSecond: f64, nowNanos: u64) TokenBucket {
|
||||
std.debug.assert(capacity > 0);
|
||||
std.debug.assert(refillPerSecond > 0);
|
||||
return .{
|
||||
.capacity = capacity,
|
||||
.refillPerSecond = refillPerSecond,
|
||||
.tokens = capacity,
|
||||
.lastRefillNanos = nowNanos,
|
||||
};
|
||||
}
|
||||
|
||||
fn refillLocked(self: *TokenBucket, nowNanos: u64) void {
|
||||
if (nowNanos <= self.lastRefillNanos) return;
|
||||
const elapsedNanos = nowNanos - self.lastRefillNanos;
|
||||
const elapsedSeconds = @as(f64, @floatFromInt(elapsedNanos)) / std.time.ns_per_s;
|
||||
self.tokens = @min(self.capacity, self.tokens + elapsedSeconds * self.refillPerSecond);
|
||||
self.lastRefillNanos = nowNanos;
|
||||
}
|
||||
|
||||
/// Attempts to consume `n` tokens at time `nowNanos`. All-or-nothing:
|
||||
/// never partially consumes on failure.
|
||||
pub fn tryAcquireAt(self: *TokenBucket, n: f64, nowNanos: u64) bool {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
self.refillLocked(nowNanos);
|
||||
if (self.tokens >= n) {
|
||||
self.tokens -= n;
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
pub fn tryAcquireNow(self: *TokenBucket, n: f64) bool {
|
||||
return self.tryAcquireAt(n, @intCast(std.time.nanoTimestamp()));
|
||||
}
|
||||
|
||||
pub fn availableAt(self: *TokenBucket, nowNanos: u64) f64 {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
self.refillLocked(nowNanos);
|
||||
return self.tokens;
|
||||
}
|
||||
};
|
||||
|
||||
/// Sliding-window-counter rate limiter (the "Cloudflare" approximation):
|
||||
/// tracks a count for the current fixed window and the previous one, and
|
||||
/// weights the previous window's count by how much of it still overlaps the
|
||||
/// trailing `windowNanos` interval. O(1) memory unlike a sliding log, more
|
||||
/// accurate at window boundaries than a plain fixed window.
|
||||
pub const SlidingWindowLimiter = struct {
|
||||
limit: u64,
|
||||
windowNanos: u64,
|
||||
currentWindowStart: u64,
|
||||
currentCount: u64 = 0,
|
||||
previousCount: u64 = 0,
|
||||
mutex: std.Thread.Mutex = .{},
|
||||
|
||||
pub fn init(limit: u64, windowNanos: u64, nowNanos: u64) SlidingWindowLimiter {
|
||||
std.debug.assert(limit > 0);
|
||||
std.debug.assert(windowNanos > 0);
|
||||
return .{
|
||||
.limit = limit,
|
||||
.windowNanos = windowNanos,
|
||||
.currentWindowStart = alignedWindowStart(nowNanos, windowNanos),
|
||||
};
|
||||
}
|
||||
|
||||
fn alignedWindowStart(nowNanos: u64, windowNanos: u64) u64 {
|
||||
return (nowNanos / windowNanos) * windowNanos;
|
||||
}
|
||||
|
||||
fn rotateLocked(self: *SlidingWindowLimiter, nowNanos: u64) void {
|
||||
const windowStart = alignedWindowStart(nowNanos, self.windowNanos);
|
||||
if (windowStart == self.currentWindowStart) return;
|
||||
const windowsElapsed = (windowStart - self.currentWindowStart) / self.windowNanos;
|
||||
if (windowsElapsed == 1) {
|
||||
self.previousCount = self.currentCount;
|
||||
} else {
|
||||
self.previousCount = 0;
|
||||
}
|
||||
self.currentCount = 0;
|
||||
self.currentWindowStart = windowStart;
|
||||
}
|
||||
|
||||
pub fn allowAt(self: *SlidingWindowLimiter, nowNanos: u64) bool {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
self.rotateLocked(nowNanos);
|
||||
const elapsedInWindow = nowNanos - self.currentWindowStart;
|
||||
const overlap = 1.0 - @as(f64, @floatFromInt(elapsedInWindow)) / @as(f64, @floatFromInt(self.windowNanos));
|
||||
const clampedOverlap = std.math.clamp(overlap, 0, 1);
|
||||
const weighted = @as(f64, @floatFromInt(self.previousCount)) * clampedOverlap + @as(f64, @floatFromInt(self.currentCount));
|
||||
if (weighted >= @as(f64, @floatFromInt(self.limit))) return false;
|
||||
self.currentCount += 1;
|
||||
return true;
|
||||
}
|
||||
|
||||
pub fn allowNow(self: *SlidingWindowLimiter) bool {
|
||||
return self.allowAt(@intCast(std.time.nanoTimestamp()));
|
||||
}
|
||||
};
|
||||
|
||||
const MetricKind = enum { counter, gauge, histogram };
|
||||
|
||||
const MetricRef = union(MetricKind) {
|
||||
counter: *const MetricCounter,
|
||||
gauge: *const Gauge,
|
||||
histogram: *Histogram,
|
||||
};
|
||||
|
||||
/// Named-metric registry with a Prometheus-style text exporter. Stores only
|
||||
/// borrowed pointers (owned name copies) - callers keep the metrics alive.
|
||||
pub const Registry = struct {
|
||||
allocator: std.mem.Allocator,
|
||||
entries: std.StringHashMapUnmanaged(MetricRef) = .{},
|
||||
|
||||
pub fn init(allocator: std.mem.Allocator) Registry {
|
||||
return .{ .allocator = allocator };
|
||||
}
|
||||
|
||||
pub fn deinit(self: *Registry) void {
|
||||
var it = self.entries.keyIterator();
|
||||
while (it.next()) |key| self.allocator.free(key.*);
|
||||
self.entries.deinit(self.allocator);
|
||||
self.* = undefined;
|
||||
}
|
||||
|
||||
fn put(self: *Registry, name: []const u8, ref: MetricRef) !void {
|
||||
if (self.entries.getPtr(name)) |existing| {
|
||||
existing.* = ref;
|
||||
return;
|
||||
}
|
||||
const owned = try self.allocator.dupe(u8, name);
|
||||
errdefer self.allocator.free(owned);
|
||||
try self.entries.put(self.allocator, owned, ref);
|
||||
}
|
||||
|
||||
pub fn registerCounter(self: *Registry, name: []const u8, c: *const MetricCounter) !void {
|
||||
try self.put(name, .{ .counter = c });
|
||||
}
|
||||
|
||||
pub fn registerGauge(self: *Registry, name: []const u8, g: *const Gauge) !void {
|
||||
try self.put(name, .{ .gauge = g });
|
||||
}
|
||||
|
||||
pub fn registerHistogram(self: *Registry, name: []const u8, h: *Histogram) !void {
|
||||
try self.put(name, .{ .histogram = h });
|
||||
}
|
||||
|
||||
pub fn count(self: *const Registry) usize {
|
||||
return self.entries.count();
|
||||
}
|
||||
|
||||
/// Writes every registered metric in Prometheus exposition format.
|
||||
/// Output is sorted by name so it is deterministic (and diffable).
|
||||
pub fn exportText(self: *Registry, allocator: std.mem.Allocator, writer: anytype) !void {
|
||||
var names = try allocator.alloc([]const u8, self.entries.count());
|
||||
defer allocator.free(names);
|
||||
var it = self.entries.keyIterator();
|
||||
var i: usize = 0;
|
||||
while (it.next()) |key| : (i += 1) names[i] = key.*;
|
||||
std.mem.sort([]const u8, names, {}, lessThanStr);
|
||||
|
||||
for (names) |name| {
|
||||
const ref = self.entries.get(name).?;
|
||||
switch (ref) {
|
||||
.counter => |c| try writer.print("{s} {d}\n", .{ name, c.get() }),
|
||||
.gauge => |g| try writer.print("{s} {d}\n", .{ name, g.get() }),
|
||||
.histogram => |h| {
|
||||
var cumulative: u64 = 0;
|
||||
for (h.bounds, 0..) |bound, idx| {
|
||||
cumulative += h.counts[idx];
|
||||
try writer.print("{s}_bucket{{le=\"{d}\"}} {d}\n", .{ name, bound, cumulative });
|
||||
}
|
||||
cumulative += h.counts[h.bounds.len];
|
||||
try writer.print("{s}_bucket{{le=\"+Inf\"}} {d}\n", .{ name, cumulative });
|
||||
try writer.print("{s}_sum {d}\n", .{ name, h.sum });
|
||||
try writer.print("{s}_count {d}\n", .{ name, h.total });
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
fn lessThanStr(_: void, a: []const u8, b: []const u8) bool {
|
||||
return std.mem.order(u8, a, b) == .lt;
|
||||
}
|
||||
|
||||
/// Renders a Db Stats snapshot in the same Prometheus text format used by
|
||||
/// Registry.exportText, so `stats` output and ad-hoc registry metrics can be
|
||||
/// scraped through one shared shape.
|
||||
pub fn formatDbStats(writer: anytype, prefix: []const u8, s: stats_mod.Stats) !void {
|
||||
try writer.print("{s}_memtable_entries {d}\n", .{ prefix, s.memtableEntries });
|
||||
try writer.print("{s}_memtable_bytes {d}\n", .{ prefix, s.memtableBytes });
|
||||
try writer.print("{s}_sstable_count {d}\n", .{ prefix, s.sstableCount });
|
||||
try writer.print("{s}_live_entries {d}\n", .{ prefix, s.liveEntries });
|
||||
try writer.print("{s}_physical_bytes {d}\n", .{ prefix, s.physicalBytes });
|
||||
try writer.print("{s}_cache_hits {d}\n", .{ prefix, s.cacheHits });
|
||||
try writer.print("{s}_cache_misses {d}\n", .{ prefix, s.cacheMisses });
|
||||
try writer.print("{s}_cache_hit_rate {d:.6}\n", .{ prefix, s.cacheHitRate() });
|
||||
try writer.print("{s}_next_seq {d}\n", .{ prefix, s.nextSeq });
|
||||
const deepest = s.deepestLevel();
|
||||
var i: usize = 0;
|
||||
while (i <= deepest) : (i += 1) {
|
||||
const level = s.levels[i];
|
||||
try writer.print("{s}_level_tables{{level=\"{d}\"}} {d}\n", .{ prefix, i, level.tables });
|
||||
try writer.print("{s}_level_entries{{level=\"{d}\"}} {d}\n", .{ prefix, i, level.entries });
|
||||
}
|
||||
}
|
||||
|
||||
test "counter incr and incrBy accumulate" {
|
||||
var c = MetricCounter{};
|
||||
c.incr();
|
||||
c.incrBy(41);
|
||||
try std.testing.expectEqual(@as(u64, 42), c.get());
|
||||
c.reset();
|
||||
try std.testing.expectEqual(@as(u64, 0), c.get());
|
||||
}
|
||||
|
||||
test "gauge set/add round trip through bit reinterpretation" {
|
||||
var g = Gauge{};
|
||||
g.set(3.5);
|
||||
try std.testing.expectApproxEqAbs(@as(f64, 3.5), g.get(), 1e-12);
|
||||
g.add(-1.25);
|
||||
try std.testing.expectApproxEqAbs(@as(f64, 2.25), g.get(), 1e-12);
|
||||
}
|
||||
|
||||
test "histogram observe buckets and mean" {
|
||||
const bounds = [_]f64{ 10, 50, 100 };
|
||||
var h = try Histogram.init(std.testing.allocator, &bounds);
|
||||
defer h.deinit();
|
||||
|
||||
h.observe(5);
|
||||
h.observe(20);
|
||||
h.observe(20);
|
||||
h.observe(75);
|
||||
h.observe(1000);
|
||||
|
||||
try std.testing.expectEqual(@as(u64, 1), h.bucketCount(0));
|
||||
try std.testing.expectEqual(@as(u64, 2), h.bucketCount(1));
|
||||
try std.testing.expectEqual(@as(u64, 1), h.bucketCount(2));
|
||||
try std.testing.expectEqual(@as(u64, 1), h.bucketCount(3));
|
||||
try std.testing.expectEqual(@as(u64, 5), h.count());
|
||||
try std.testing.expectApproxEqAbs(@as(f64, (5 + 20 + 20 + 75 + 1000) / 5.0), h.mean(), 1e-9);
|
||||
}
|
||||
|
||||
test "histogram quantile interpolates within a bucket" {
|
||||
const bounds = [_]f64{ 100, 200 };
|
||||
var h = try Histogram.init(std.testing.allocator, &bounds);
|
||||
defer h.deinit();
|
||||
|
||||
// 4 values land in the (0,100] bucket, evenly spread by construction;
|
||||
// linear interpolation only sees prev/cur cumulative counts, so the
|
||||
// median of an all-in-one-bucket histogram is the bucket's own bound.
|
||||
var i: usize = 0;
|
||||
while (i < 4) : (i += 1) h.observe(50);
|
||||
try std.testing.expectApproxEqAbs(@as(f64, 100), h.quantile(0.5), 1e-9);
|
||||
|
||||
// Now spread evenly across both finite buckets: p50 should land exactly
|
||||
// on the boundary between them.
|
||||
h.observe(150);
|
||||
h.observe(150);
|
||||
h.observe(150);
|
||||
h.observe(150);
|
||||
// counts: bucket0=4 (<=100), bucket1=4 (<=200) => cumulative [4,8], total 8
|
||||
const q50 = h.quantile(0.5);
|
||||
try std.testing.expectApproxEqAbs(@as(f64, 100), q50, 1e-9);
|
||||
}
|
||||
|
||||
test "histogram quantile falls back to last bound for the +Inf bucket" {
|
||||
const bounds = [_]f64{ 10, 20 };
|
||||
var h = try Histogram.init(std.testing.allocator, &bounds);
|
||||
defer h.deinit();
|
||||
h.observe(5);
|
||||
h.observe(9999);
|
||||
try std.testing.expectApproxEqAbs(@as(f64, 20), h.quantile(1.0), 1e-9);
|
||||
}
|
||||
|
||||
test "linearBuckets and exponentialBuckets produce expected boundaries" {
|
||||
const lin = try linearBuckets(std.testing.allocator, 0, 10, 5);
|
||||
defer std.testing.allocator.free(lin);
|
||||
try std.testing.expectEqualSlices(f64, &[_]f64{ 0, 10, 20, 30, 40 }, lin);
|
||||
|
||||
const exp = try exponentialBuckets(std.testing.allocator, 1, 2, 5);
|
||||
defer std.testing.allocator.free(exp);
|
||||
try std.testing.expectEqualSlices(f64, &[_]f64{ 1, 2, 4, 8, 16 }, exp);
|
||||
}
|
||||
|
||||
test "token bucket refills over time and denies when empty" {
|
||||
var b = TokenBucket.init(10, 5, 0); // 10 capacity, 5 tokens/sec
|
||||
try std.testing.expect(b.tryAcquireAt(10, 0));
|
||||
try std.testing.expect(!b.tryAcquireAt(1, 0));
|
||||
|
||||
// half a second later: +2.5 tokens
|
||||
try std.testing.expect(!b.tryAcquireAt(3, std.time.ns_per_s / 2));
|
||||
try std.testing.expectApproxEqAbs(@as(f64, 2.5), b.availableAt(std.time.ns_per_s / 2), 1e-9);
|
||||
|
||||
// two full seconds later than that: refill caps at capacity
|
||||
const later = std.time.ns_per_s / 2 + 2 * std.time.ns_per_s;
|
||||
try std.testing.expect(b.tryAcquireAt(10, later));
|
||||
try std.testing.expect(!b.tryAcquireAt(1, later));
|
||||
}
|
||||
|
||||
test "token bucket never partially consumes on a failed acquire" {
|
||||
var b = TokenBucket.init(5, 1, 0);
|
||||
try std.testing.expect(!b.tryAcquireAt(6, 0));
|
||||
try std.testing.expectApproxEqAbs(@as(f64, 5), b.availableAt(0), 1e-9);
|
||||
}
|
||||
|
||||
test "sliding window limiter smooths across a window boundary" {
|
||||
const windowNs: u64 = std.time.ns_per_s;
|
||||
var lim = SlidingWindowLimiter.init(10, windowNs, 0);
|
||||
|
||||
var i: usize = 0;
|
||||
while (i < 10) : (i += 1) try std.testing.expect(lim.allowAt(0));
|
||||
try std.testing.expect(!lim.allowAt(0)); // exhausted current window
|
||||
|
||||
// Halfway into the next window: previous window's 10 requests are
|
||||
// weighted at 50% overlap (5), so only 5 more should be allowed before
|
||||
// the weighted total reaches the limit of 10.
|
||||
const halfwayNextWindow = windowNs + windowNs / 2;
|
||||
var allowed: usize = 0;
|
||||
i = 0;
|
||||
while (i < 20) : (i += 1) {
|
||||
if (lim.allowAt(halfwayNextWindow)) allowed += 1;
|
||||
}
|
||||
try std.testing.expectEqual(@as(usize, 5), allowed);
|
||||
}
|
||||
|
||||
test "sliding window limiter resets fully after a gap of more than one window" {
|
||||
const windowNs: u64 = std.time.ns_per_s;
|
||||
var lim = SlidingWindowLimiter.init(3, windowNs, 0);
|
||||
try std.testing.expect(lim.allowAt(0));
|
||||
try std.testing.expect(lim.allowAt(0));
|
||||
try std.testing.expect(lim.allowAt(0));
|
||||
try std.testing.expect(!lim.allowAt(0));
|
||||
|
||||
// Jump 10 windows ahead: previous window count must not leak in.
|
||||
const farFuture = 10 * windowNs;
|
||||
try std.testing.expect(lim.allowAt(farFuture));
|
||||
try std.testing.expect(lim.allowAt(farFuture));
|
||||
try std.testing.expect(lim.allowAt(farFuture));
|
||||
try std.testing.expect(!lim.allowAt(farFuture));
|
||||
}
|
||||
|
||||
test "registry exports counters gauges and histograms as sorted prometheus text" {
|
||||
var c = MetricCounter{};
|
||||
c.incrBy(7);
|
||||
var g = Gauge{};
|
||||
g.set(2.5);
|
||||
const bounds = [_]f64{ 5, 10 };
|
||||
var h = try Histogram.init(std.testing.allocator, &bounds);
|
||||
defer h.deinit();
|
||||
h.observe(3);
|
||||
h.observe(8);
|
||||
|
||||
var reg = Registry.init(std.testing.allocator);
|
||||
defer reg.deinit();
|
||||
try reg.registerCounter("requests_total", &c);
|
||||
try reg.registerGauge("temperature", &g);
|
||||
try reg.registerHistogram("latency", &h);
|
||||
try std.testing.expectEqual(@as(usize, 3), reg.count());
|
||||
|
||||
var buf = std.ArrayList(u8).init(std.testing.allocator);
|
||||
defer buf.deinit();
|
||||
try reg.exportText(std.testing.allocator, buf.writer());
|
||||
|
||||
const text = buf.items;
|
||||
// sorted alphabetically: latency < requests_total < temperature
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "latency_bucket{le=\"5\"} 1\n") != null);
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "latency_bucket{le=\"10\"} 2\n") != null);
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "latency_bucket{le=\"+Inf\"} 2\n") != null);
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "latency_sum 11") != null);
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "latency_count 2\n") != null);
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "requests_total 7\n") != null);
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "temperature 2.5\n") != null);
|
||||
|
||||
const latencyPos = std.mem.indexOf(u8, text, "latency_bucket").?;
|
||||
const requestsPos = std.mem.indexOf(u8, text, "requests_total").?;
|
||||
const temperaturePos = std.mem.indexOf(u8, text, "temperature").?;
|
||||
try std.testing.expect(latencyPos < requestsPos);
|
||||
try std.testing.expect(requestsPos < temperaturePos);
|
||||
}
|
||||
|
||||
test "registry re-registering the same name overwrites rather than duplicating" {
|
||||
var c1 = MetricCounter{};
|
||||
c1.incrBy(1);
|
||||
var c2 = MetricCounter{};
|
||||
c2.incrBy(99);
|
||||
|
||||
var reg = Registry.init(std.testing.allocator);
|
||||
defer reg.deinit();
|
||||
try reg.registerCounter("x", &c1);
|
||||
try reg.registerCounter("x", &c2);
|
||||
try std.testing.expectEqual(@as(usize, 1), reg.count());
|
||||
|
||||
var buf = std.ArrayList(u8).init(std.testing.allocator);
|
||||
defer buf.deinit();
|
||||
try reg.exportText(std.testing.allocator, buf.writer());
|
||||
try std.testing.expectEqualStrings("x 99\n", buf.items);
|
||||
}
|
||||
|
||||
test "formatDbStats renders levels up to the deepest populated one" {
|
||||
var s = stats_mod.Stats{};
|
||||
s.memtableEntries = 3;
|
||||
s.memtableBytes = 128;
|
||||
s.sstableCount = 2;
|
||||
s.cacheHits = 9;
|
||||
s.cacheMisses = 1;
|
||||
s.levels[0].tables = 1;
|
||||
s.levels[0].entries = 10;
|
||||
s.levels[2].tables = 4;
|
||||
s.levels[2].entries = 40;
|
||||
|
||||
var buf = std.ArrayList(u8).init(std.testing.allocator);
|
||||
defer buf.deinit();
|
||||
try formatDbStats(buf.writer(), "cobalt", s);
|
||||
|
||||
const text = buf.items;
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "cobalt_memtable_entries 3\n") != null);
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "cobalt_cache_hit_rate 0.900000\n") != null);
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "cobalt_level_tables{level=\"0\"} 1\n") != null);
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "cobalt_level_tables{level=\"2\"} 4\n") != null);
|
||||
// level 1 is empty but must still appear since deepestLevel is 2
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "cobalt_level_tables{level=\"1\"} 0\n") != null);
|
||||
// level 3+ must not appear at all
|
||||
try std.testing.expect(std.mem.indexOf(u8, text, "level=\"3\"") == null);
|
||||
}
|
||||
+3998
File diff suppressed because it is too large.
Load diff
+848
@@ -0,0 +1,848 @@
|
||||
const std = @import("std");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
|
||||
/// Default bound on how many undelivered messages a single subscriber's
|
||||
/// mailbox holds before the oldest is dropped to make room for the newest -
|
||||
/// see `Subscriber.push`'s doc comment for why drop-oldest (never blocking
|
||||
/// the publisher) is the chosen backpressure policy.
|
||||
pub const defaultCapacity: usize = 256;
|
||||
|
||||
/// One delivered message: which channel it arrived on (a subscriber can
|
||||
/// listen to several, so the tag travels with the payload) and the payload
|
||||
/// bytes themselves, both independently owned copies - `publish` never
|
||||
/// hands the same buffer to two subscribers.
|
||||
pub const Message = struct {
|
||||
channel: []u8,
|
||||
payload: []u8,
|
||||
|
||||
pub fn deinit(self: Message, allocator: Allocator) void {
|
||||
allocator.free(self.channel);
|
||||
allocator.free(self.payload);
|
||||
}
|
||||
};
|
||||
|
||||
/// One subscriber's private mailbox: a bounded FIFO of `Message`s plus the
|
||||
/// synchronization a consumer thread needs to block until either a new
|
||||
/// message arrives or the subscription is closed. A `Subscriber` is always
|
||||
/// created by `Hub.subscribe` and must be retired through
|
||||
/// `Hub.closeSubscriber` then `Hub.destroySubscriber`, in that order (see
|
||||
/// their doc comments for why the two steps are separate).
|
||||
pub const Subscriber = struct {
|
||||
allocator: Allocator,
|
||||
mutex: std.Thread.Mutex = .{},
|
||||
cond: std.Thread.Condition = .{},
|
||||
queue: std.ArrayList(Message),
|
||||
capacity: usize,
|
||||
dropped: usize = 0,
|
||||
closed: bool = false,
|
||||
/// Assigned by `Hub.subscribe` (monotonically increasing per `Hub`
|
||||
/// instance, never reused within its lifetime) so a control connection
|
||||
/// that never saw this `Subscriber`'s pointer - a separate `unsubscribe`
|
||||
/// request, possibly on a different TCP connection entirely - can still
|
||||
/// name it. Only meaningful together with the `Hub` that assigned it;
|
||||
/// not persisted, not unique across process restarts.
|
||||
id: u64 = 0,
|
||||
/// Channel names this subscriber is registered under in whichever `Hub`
|
||||
/// created it - copied here (not just left as entries in
|
||||
/// `Hub.channels`) so `Hub.closeSubscriber`/`Hub.unsubscribe` can find
|
||||
/// every channel to unregister from without the caller having to
|
||||
/// remember the original subscribe list.
|
||||
channels: [][]u8,
|
||||
|
||||
fn init(allocator: Allocator, capacity: usize, channels: [][]u8) Subscriber {
|
||||
return .{
|
||||
.allocator = allocator,
|
||||
.queue = std.ArrayList(Message).init(allocator),
|
||||
.capacity = capacity,
|
||||
.channels = channels,
|
||||
};
|
||||
}
|
||||
|
||||
/// Delivers one message, taking ownership of it. Never blocks: a
|
||||
/// `capacity` of zero drops every message immediately (a legal, if
|
||||
/// degenerate, "nobody's listening yet" mailbox); otherwise, once the
|
||||
/// mailbox is at `capacity`, the oldest queued message is dropped (and
|
||||
/// freed) to make room, `dropped` is incremented, and the new message
|
||||
/// is appended - so `waitMessage` always eventually surfaces the most
|
||||
/// RECENT messages a slow subscriber missed, never the oldest stale
|
||||
/// ones. A message pushed after `close()` is silently freed and does
|
||||
/// not count as delivered or dropped - a closed subscriber is never
|
||||
/// coming back to read it, and `dropped` exists to describe backpressure
|
||||
/// on a live subscriber, not the ordinary end of one.
|
||||
fn push(self: *Subscriber, msg: Message) void {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
if (self.closed) {
|
||||
msg.deinit(self.allocator);
|
||||
return;
|
||||
}
|
||||
if (self.capacity == 0) {
|
||||
msg.deinit(self.allocator);
|
||||
self.dropped += 1;
|
||||
return;
|
||||
}
|
||||
if (self.queue.items.len >= self.capacity) {
|
||||
const old = self.queue.orderedRemove(0);
|
||||
old.deinit(self.allocator);
|
||||
self.dropped += 1;
|
||||
}
|
||||
self.queue.append(msg) catch {
|
||||
msg.deinit(self.allocator);
|
||||
self.dropped += 1;
|
||||
return;
|
||||
};
|
||||
self.cond.signal();
|
||||
}
|
||||
|
||||
/// Blocks until a message is available or the subscription is closed
|
||||
/// with an empty mailbox, whichever happens first. A close with
|
||||
/// messages still queued is drained first - `waitMessage` keeps
|
||||
/// returning them in order before finally returning `null`, so a
|
||||
/// consumer never loses a message that was successfully enqueued
|
||||
/// before `close()` ran. Ownership of the returned `Message` passes to
|
||||
/// the caller, who must `deinit` it.
|
||||
pub fn waitMessage(self: *Subscriber) ?Message {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
while (self.queue.items.len == 0 and !self.closed) {
|
||||
self.cond.wait(&self.mutex);
|
||||
}
|
||||
if (self.queue.items.len == 0) return null;
|
||||
return self.queue.orderedRemove(0);
|
||||
}
|
||||
|
||||
/// Idempotent: marks the mailbox closed and wakes every thread blocked
|
||||
/// in `waitMessage`. Safe to call from any thread, any number of times
|
||||
/// concurrently - a watcher thread noticing a client disconnect and a
|
||||
/// publisher noticing a dead connection can both call it without
|
||||
/// coordinating between themselves.
|
||||
pub fn close(self: *Subscriber) void {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
if (self.closed) return;
|
||||
self.closed = true;
|
||||
self.cond.broadcast();
|
||||
}
|
||||
|
||||
pub fn isClosed(self: *Subscriber) bool {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
return self.closed;
|
||||
}
|
||||
|
||||
pub fn droppedCount(self: *Subscriber) usize {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
return self.dropped;
|
||||
}
|
||||
|
||||
/// Frees the mailbox's own queued messages and the subscriber's
|
||||
/// channel-name copies, then the caller frees the `Subscriber` struct
|
||||
/// itself via `Hub.destroySubscriber`. Only the thread that owns this
|
||||
/// subscriber past its last `waitMessage`/close call may reach this -
|
||||
/// see `Hub.destroySubscriber`'s doc comment for the ownership handoff
|
||||
/// it depends on.
|
||||
fn deinitContents(self: *Subscriber) void {
|
||||
for (self.queue.items) |m| m.deinit(self.allocator);
|
||||
self.queue.deinit();
|
||||
for (self.channels) |c| self.allocator.free(c);
|
||||
self.allocator.free(self.channels);
|
||||
}
|
||||
};
|
||||
|
||||
/// The channel registry every `Subscriber` is filed under. Completely
|
||||
/// independent of `Db` - publish/subscribe never touches storage, so a
|
||||
/// `Hub`'s own mutex is never held anywhere near `Db`'s (see `server.zig`'s
|
||||
/// `dispatch`, which handles the `subscribe`/`publish` opcodes before it
|
||||
/// ever reaches for `Server.mutex`, so a slow or permanently-open
|
||||
/// subscriber connection can never stall a read or write against the
|
||||
/// store).
|
||||
pub const Hub = struct {
|
||||
allocator: Allocator,
|
||||
mutex: std.Thread.Mutex = .{},
|
||||
channels: std.StringHashMap(std.ArrayList(*Subscriber)),
|
||||
/// Every currently-live subscriber, keyed by the `id` `subscribe`
|
||||
/// assigned it - lets `unsubscribe` (see its doc comment) look one up by
|
||||
/// the bare number a client saw in its subscribe ack, without needing
|
||||
/// the original `*Subscriber` pointer, which never crosses a second TCP
|
||||
/// connection. Entries are removed exactly when `closeSubscriber` runs
|
||||
/// (or `unsubscribe` empties a subscriber's channel set and closes it
|
||||
/// itself), so a stale or already-closed id reliably misses here rather
|
||||
/// than risking a dangling pointer.
|
||||
byId: std.AutoHashMap(u64, *Subscriber),
|
||||
nextId: u64 = 0,
|
||||
|
||||
pub fn init(allocator: Allocator) Hub {
|
||||
return .{
|
||||
.allocator = allocator,
|
||||
.channels = std.StringHashMap(std.ArrayList(*Subscriber)).init(allocator),
|
||||
.byId = std.AutoHashMap(u64, *Subscriber).init(allocator),
|
||||
};
|
||||
}
|
||||
|
||||
/// Frees every remaining channel's subscriber list and key copy, and the
|
||||
/// id registry. Does NOT touch any `Subscriber` still registered -
|
||||
/// callers must `closeSubscriber` (which also unregisters, so in an
|
||||
/// orderly shutdown this map is already empty by the time `deinit` runs)
|
||||
/// and `destroySubscriber` every subscriber they created first, the same
|
||||
/// "caller drains, then frees the container" ordering `Db`/`Server`
|
||||
/// already use elsewhere in this codebase.
|
||||
pub fn deinit(self: *Hub) void {
|
||||
var it = self.channels.iterator();
|
||||
while (it.next()) |entry| {
|
||||
entry.value_ptr.deinit();
|
||||
self.allocator.free(entry.key_ptr.*);
|
||||
}
|
||||
self.channels.deinit();
|
||||
self.byId.deinit();
|
||||
}
|
||||
|
||||
/// Creates a new subscriber registered under every name in
|
||||
/// `channelNames` (duplicate names collapse to one registration, so a
|
||||
/// message published on a channel a caller listed twice is still only
|
||||
/// ever delivered once). `capacity` bounds its mailbox (see
|
||||
/// `Subscriber.push`). Returns a pointer the caller must eventually
|
||||
/// pass to `closeSubscriber` then `destroySubscriber`, in that order.
|
||||
pub fn subscribe(self: *Hub, channelNames: []const []const u8, capacity: usize) !*Subscriber {
|
||||
var unique = std.ArrayList([]const u8).init(self.allocator);
|
||||
defer unique.deinit();
|
||||
for (channelNames) |name| {
|
||||
var seen = false;
|
||||
for (unique.items) |u| {
|
||||
if (std.mem.eql(u8, u, name)) {
|
||||
seen = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!seen) try unique.append(name);
|
||||
}
|
||||
|
||||
const owned = try self.allocator.alloc([]u8, unique.items.len);
|
||||
var ownedCount: usize = 0;
|
||||
errdefer {
|
||||
for (owned[0..ownedCount]) |c| self.allocator.free(c);
|
||||
self.allocator.free(owned);
|
||||
}
|
||||
for (unique.items, 0..) |name, i| {
|
||||
owned[i] = try self.allocator.dupe(u8, name);
|
||||
ownedCount += 1;
|
||||
}
|
||||
|
||||
const sub = try self.allocator.create(Subscriber);
|
||||
errdefer self.allocator.destroy(sub);
|
||||
sub.* = Subscriber.init(self.allocator, capacity, owned);
|
||||
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
|
||||
const id = self.nextId;
|
||||
self.nextId += 1;
|
||||
sub.id = id;
|
||||
try self.byId.put(id, sub);
|
||||
errdefer _ = self.byId.remove(id);
|
||||
|
||||
var registered: usize = 0;
|
||||
errdefer {
|
||||
for (owned[0..registered]) |name| {
|
||||
if (self.channels.getPtr(name)) |list| {
|
||||
for (list.items, 0..) |s, idx| {
|
||||
if (s == sub) {
|
||||
_ = list.swapRemove(idx);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
for (owned) |name| {
|
||||
const gop = try self.channels.getOrPut(name);
|
||||
if (!gop.found_existing) {
|
||||
gop.key_ptr.* = self.allocator.dupe(u8, name) catch |err| {
|
||||
// The map entry was already inserted with `name`'s
|
||||
// pointer as a placeholder key by `getOrPut` itself;
|
||||
// remove it before propagating so `deinit`/`subscribe`
|
||||
// never iterate over a key pointing into memory this
|
||||
// function is about to free via `owned`'s errdefer.
|
||||
_ = self.channels.remove(name);
|
||||
return err;
|
||||
};
|
||||
gop.value_ptr.* = std.ArrayList(*Subscriber).init(self.allocator);
|
||||
}
|
||||
try gop.value_ptr.append(sub);
|
||||
registered += 1;
|
||||
}
|
||||
|
||||
return sub;
|
||||
}
|
||||
|
||||
/// Stops delivering new messages to `sub` and unregisters it from every
|
||||
/// channel it was subscribed to, wakes any thread blocked in its
|
||||
/// `waitMessage`, and drops a channel's map entry entirely once its
|
||||
/// subscriber list becomes empty (an unsubscribed-from channel is a
|
||||
/// completely normal, permanent steady state, not a reason to keep an
|
||||
/// empty list around forever). Safe to call more than once (from a
|
||||
/// watcher thread and a failed publish both noticing the same dead
|
||||
/// connection, say) - both `Subscriber.close` and the unregister loop
|
||||
/// below are no-ops the second time. Does NOT free `sub` itself - see
|
||||
/// `destroySubscriber`.
|
||||
pub fn closeSubscriber(self: *Hub, sub: *Subscriber) void {
|
||||
sub.close();
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
_ = self.byId.remove(sub.id);
|
||||
for (sub.channels) |name| {
|
||||
if (self.channels.getPtr(name)) |list| {
|
||||
for (list.items, 0..) |s, idx| {
|
||||
if (s == sub) {
|
||||
_ = list.swapRemove(idx);
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (list.items.len == 0) {
|
||||
if (self.channels.fetchRemove(name)) |kv| {
|
||||
kv.value.deinit();
|
||||
self.allocator.free(kv.key);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Removes every name in `channelNames` from `id`'s subscription set
|
||||
/// without disturbing any channel it stays registered on, so a client
|
||||
/// that subscribed to several channels can drop a subset from a SEPARATE
|
||||
/// connection - the control channel `handleSubscribe`'s doc comment
|
||||
/// says a live subscribe connection can never multiplex further
|
||||
/// requests onto itself. `channelNames` naming a channel `id` was never
|
||||
/// on is harmless and simply does not count toward the returned total
|
||||
/// (matching `publish`-to-an-empty-channel's "no-op, not an error"
|
||||
/// posture). Returns `error.UnknownSubscriber` if `id` names no
|
||||
/// currently-live subscriber - already closed (a normal disconnect
|
||||
/// race, not a caller bug) or never issued (a caller typo) look
|
||||
/// identical from here, both correctly rejected as "there is nothing
|
||||
/// left to unsubscribe."
|
||||
///
|
||||
/// If this empties `id`'s entire channel set, the subscriber is closed
|
||||
/// outright (`sub.close()`, unregistered from `byId`) rather than left
|
||||
/// alive with nothing left to ever receive - the same terminal state a
|
||||
/// disconnect reaches through `closeSubscriber`, just reached from the
|
||||
/// opposite direction. `handleSubscribe`'s push loop notices this
|
||||
/// exactly like any other close: `sub.waitMessage()` returns `null`, the
|
||||
/// loop ends, and the connection's own `closeSubscriber`/
|
||||
/// `destroySubscriber` pair runs as usual - by then `sub.channels` is
|
||||
/// already empty, so that second `closeSubscriber` call's unregister
|
||||
/// loop simply does nothing (idempotent, like every other double-close
|
||||
/// path in this file).
|
||||
pub fn unsubscribe(self: *Hub, id: u64, channelNames: []const []const u8) !usize {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
const sub = self.byId.get(id) orelse return error.UnknownSubscriber;
|
||||
|
||||
var kept = try std.ArrayList([]u8).initCapacity(self.allocator, sub.channels.len);
|
||||
errdefer kept.deinit();
|
||||
var removed: usize = 0;
|
||||
for (sub.channels) |name| {
|
||||
var matches = false;
|
||||
for (channelNames) |target| {
|
||||
if (std.mem.eql(u8, name, target)) {
|
||||
matches = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!matches) {
|
||||
kept.appendAssumeCapacity(name);
|
||||
continue;
|
||||
}
|
||||
removed += 1;
|
||||
if (self.channels.getPtr(name)) |list| {
|
||||
for (list.items, 0..) |s, idx| {
|
||||
if (s == sub) {
|
||||
_ = list.swapRemove(idx);
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (list.items.len == 0) {
|
||||
if (self.channels.fetchRemove(name)) |kv| {
|
||||
kv.value.deinit();
|
||||
self.allocator.free(kv.key);
|
||||
}
|
||||
}
|
||||
}
|
||||
self.allocator.free(name);
|
||||
}
|
||||
self.allocator.free(sub.channels);
|
||||
sub.channels = try kept.toOwnedSlice();
|
||||
|
||||
if (sub.channels.len == 0) {
|
||||
_ = self.byId.remove(id);
|
||||
sub.close();
|
||||
}
|
||||
return removed;
|
||||
}
|
||||
|
||||
/// Frees `sub`'s own memory (mailbox contents, channel-name copies, the
|
||||
/// struct itself). Must only be called after `closeSubscriber(sub)` has
|
||||
/// already run and every thread that might still call `sub.waitMessage`/
|
||||
/// `sub.push` (directly or through `publish`) has stopped touching it -
|
||||
/// `closeSubscriber` unregisters `sub` from `Hub.channels` precisely so
|
||||
/// `publish` can never hand this function a dangling pointer to race
|
||||
/// against.
|
||||
pub fn destroySubscriber(self: *Hub, sub: *Subscriber) void {
|
||||
sub.deinitContents();
|
||||
self.allocator.destroy(sub);
|
||||
}
|
||||
|
||||
/// Number of subscribers currently registered under `channel` (0 if it
|
||||
/// has none, including a channel name nobody has ever used). Mainly for
|
||||
/// tests and metrics - `publish` recomputes its own snapshot rather than
|
||||
/// calling this, to avoid taking `self.mutex` twice per publish.
|
||||
pub fn channelSubscriberCount(self: *Hub, channel: []const u8) usize {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
if (self.channels.get(channel)) |list| return list.items.len;
|
||||
return 0;
|
||||
}
|
||||
|
||||
/// Delivers an independent copy of `payload` to every subscriber
|
||||
/// currently registered under `channel`, returning how many subscribers
|
||||
/// it was handed to (0 if the channel currently has none - not an
|
||||
/// error, matching real pub/sub semantics where publishing to an empty
|
||||
/// channel is a routine no-op, not a failure). Never blocks on a slow
|
||||
/// subscriber - see `Subscriber.push`. `allocator` is used only for the
|
||||
/// short-lived snapshot of subscriber pointers, not the per-subscriber
|
||||
/// message copies, which are always allocated from `self.allocator`
|
||||
/// (each subscriber frees its own mailbox contents independently, so
|
||||
/// they must not depend on a caller-supplied allocator outliving this
|
||||
/// call).
|
||||
pub fn publish(self: *Hub, allocator: Allocator, channel: []const u8, payload: []const u8) !usize {
|
||||
self.mutex.lock();
|
||||
errdefer self.mutex.unlock();
|
||||
const subs = if (self.channels.get(channel)) |list| list.items else &[_]*Subscriber{};
|
||||
// Snapshot the subscriber list under the lock, but do the actual
|
||||
// per-subscriber copy+push OUTSIDE it: `Subscriber.push` takes its
|
||||
// own independent mutex, never `self.mutex`, so holding this one
|
||||
// across every subscriber's push would serialize every publish
|
||||
// behind every subscriber's own lock contention for no reason.
|
||||
const snapshot = try allocator.alloc(*Subscriber, subs.len);
|
||||
defer allocator.free(snapshot);
|
||||
@memcpy(snapshot, subs);
|
||||
self.mutex.unlock();
|
||||
|
||||
for (snapshot) |sub| {
|
||||
const channelCopy = self.allocator.dupe(u8, channel) catch continue;
|
||||
const payloadCopy = self.allocator.dupe(u8, payload) catch {
|
||||
self.allocator.free(channelCopy);
|
||||
continue;
|
||||
};
|
||||
sub.push(.{ .channel = channelCopy, .payload = payloadCopy });
|
||||
}
|
||||
return snapshot.len;
|
||||
}
|
||||
};
|
||||
|
||||
fn destroy(hub: *Hub, sub: *Subscriber) void {
|
||||
hub.closeSubscriber(sub);
|
||||
hub.destroySubscriber(sub);
|
||||
}
|
||||
|
||||
test "publish delivers a copy to a single subscriber on the channel it named" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{"weather"}, defaultCapacity);
|
||||
defer destroy(&hub, sub);
|
||||
|
||||
const delivered = try hub.publish(allocator, "weather", "sunny");
|
||||
try std.testing.expectEqual(@as(usize, 1), delivered);
|
||||
|
||||
const msg = sub.waitMessage().?;
|
||||
defer msg.deinit(allocator);
|
||||
try std.testing.expectEqualStrings("weather", msg.channel);
|
||||
try std.testing.expectEqualStrings("sunny", msg.payload);
|
||||
}
|
||||
|
||||
test "publish to a channel with no subscribers returns 0 without crashing" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const delivered = try hub.publish(allocator, "nobody-home", "hello");
|
||||
try std.testing.expectEqual(@as(usize, 0), delivered);
|
||||
}
|
||||
|
||||
test "every subscriber on a channel gets its own independent copy" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const a = try hub.subscribe(&.{"orders"}, defaultCapacity);
|
||||
defer destroy(&hub, a);
|
||||
const b = try hub.subscribe(&.{"orders"}, defaultCapacity);
|
||||
defer destroy(&hub, b);
|
||||
|
||||
const delivered = try hub.publish(allocator, "orders", "order-42");
|
||||
try std.testing.expectEqual(@as(usize, 2), delivered);
|
||||
|
||||
const ma = a.waitMessage().?;
|
||||
defer ma.deinit(allocator);
|
||||
const mb = b.waitMessage().?;
|
||||
defer mb.deinit(allocator);
|
||||
try std.testing.expectEqualStrings("order-42", ma.payload);
|
||||
try std.testing.expectEqualStrings("order-42", mb.payload);
|
||||
// Independent allocations, not aliased - freeing one must not affect
|
||||
// the other's still-valid pointer.
|
||||
try std.testing.expect(ma.payload.ptr != mb.payload.ptr);
|
||||
}
|
||||
|
||||
test "a subscriber on several channels gets messages tagged with the right channel" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{ "a", "b" }, defaultCapacity);
|
||||
defer destroy(&hub, sub);
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 1), try hub.publish(allocator, "b", "from-b"));
|
||||
try std.testing.expectEqual(@as(usize, 1), try hub.publish(allocator, "a", "from-a"));
|
||||
|
||||
const first = sub.waitMessage().?;
|
||||
defer first.deinit(allocator);
|
||||
try std.testing.expectEqualStrings("b", first.channel);
|
||||
try std.testing.expectEqualStrings("from-b", first.payload);
|
||||
|
||||
const second = sub.waitMessage().?;
|
||||
defer second.deinit(allocator);
|
||||
try std.testing.expectEqualStrings("a", second.channel);
|
||||
try std.testing.expectEqualStrings("from-a", second.payload);
|
||||
}
|
||||
|
||||
test "subscribing to the same channel name twice still delivers only once" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{ "dup", "dup" }, defaultCapacity);
|
||||
defer destroy(&hub, sub);
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 1), try hub.publish(allocator, "dup", "x"));
|
||||
const msg = sub.waitMessage().?;
|
||||
defer msg.deinit(allocator);
|
||||
// Only one message was ever queued - a second `waitMessage` in a
|
||||
// non-blocking check would hang, so instead assert there is exactly
|
||||
// one by closing right after and expecting a drained-then-null result.
|
||||
sub.close();
|
||||
try std.testing.expect(sub.waitMessage() == null);
|
||||
}
|
||||
|
||||
test "closing a subscriber unregisters it - a later publish no longer reaches it" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{"ch"}, defaultCapacity);
|
||||
hub.closeSubscriber(sub);
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 0), try hub.publish(allocator, "ch", "too-late"));
|
||||
try std.testing.expectEqual(@as(usize, 0), hub.channelSubscriberCount("ch"));
|
||||
|
||||
hub.destroySubscriber(sub);
|
||||
}
|
||||
|
||||
test "waitMessage drains already-queued messages before returning null on close" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{"ch"}, defaultCapacity);
|
||||
defer hub.destroySubscriber(sub);
|
||||
|
||||
_ = try hub.publish(allocator, "ch", "one");
|
||||
_ = try hub.publish(allocator, "ch", "two");
|
||||
hub.closeSubscriber(sub);
|
||||
|
||||
const m1 = sub.waitMessage().?;
|
||||
defer m1.deinit(allocator);
|
||||
try std.testing.expectEqualStrings("one", m1.payload);
|
||||
const m2 = sub.waitMessage().?;
|
||||
defer m2.deinit(allocator);
|
||||
try std.testing.expectEqualStrings("two", m2.payload);
|
||||
try std.testing.expect(sub.waitMessage() == null);
|
||||
}
|
||||
|
||||
test "a full mailbox drops the oldest message, not the newest" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{"ch"}, 2);
|
||||
defer destroy(&hub, sub);
|
||||
|
||||
_ = try hub.publish(allocator, "ch", "1");
|
||||
_ = try hub.publish(allocator, "ch", "2");
|
||||
_ = try hub.publish(allocator, "ch", "3");
|
||||
try std.testing.expectEqual(@as(usize, 1), sub.droppedCount());
|
||||
|
||||
const m1 = sub.waitMessage().?;
|
||||
defer m1.deinit(allocator);
|
||||
try std.testing.expectEqualStrings("2", m1.payload);
|
||||
const m2 = sub.waitMessage().?;
|
||||
defer m2.deinit(allocator);
|
||||
try std.testing.expectEqualStrings("3", m2.payload);
|
||||
}
|
||||
|
||||
test "capacity zero drops every message without ever panicking" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{"ch"}, 0);
|
||||
defer destroy(&hub, sub);
|
||||
|
||||
_ = try hub.publish(allocator, "ch", "gone");
|
||||
_ = try hub.publish(allocator, "ch", "also-gone");
|
||||
try std.testing.expectEqual(@as(usize, 2), sub.droppedCount());
|
||||
sub.close();
|
||||
try std.testing.expect(sub.waitMessage() == null);
|
||||
}
|
||||
|
||||
test "closing a subscriber wakes a thread already blocked in waitMessage" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{"ch"}, defaultCapacity);
|
||||
defer hub.destroySubscriber(sub);
|
||||
|
||||
const Waiter = struct {
|
||||
fn run(s: *Subscriber, result: *?Message) void {
|
||||
result.* = s.waitMessage();
|
||||
}
|
||||
};
|
||||
var result: ?Message = undefined;
|
||||
result = Message{ .channel = &.{}, .payload = &.{} };
|
||||
const t = try std.Thread.spawn(.{}, Waiter.run, .{ sub, &result });
|
||||
hub.closeSubscriber(sub);
|
||||
t.join();
|
||||
try std.testing.expect(result == null);
|
||||
}
|
||||
|
||||
test "publishing after closeSubscriber but before destroySubscriber never touches freed memory" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{"ch"}, defaultCapacity);
|
||||
hub.closeSubscriber(sub);
|
||||
// The channel entry is gone entirely now (empty list cleanup), so this
|
||||
// publish must find zero subscribers rather than dereferencing `sub`.
|
||||
try std.testing.expectEqual(@as(usize, 0), try hub.publish(allocator, "ch", "x"));
|
||||
hub.destroySubscriber(sub);
|
||||
}
|
||||
|
||||
test "concurrent publishers hammering one subscriber deliver every message intact" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const publisherCount = 4;
|
||||
const perPublisher = 200;
|
||||
const sub = try hub.subscribe(&.{"ch"}, publisherCount * perPublisher);
|
||||
defer destroy(&hub, sub);
|
||||
|
||||
const Publisher = struct {
|
||||
fn run(h: *Hub, alloc: Allocator, tag: usize) void {
|
||||
var buf: [32]u8 = undefined;
|
||||
var i: usize = 0;
|
||||
while (i < perPublisher) : (i += 1) {
|
||||
const text = std.fmt.bufPrint(&buf, "{d}:{d}", .{ tag, i }) catch unreachable;
|
||||
_ = h.publish(alloc, "ch", text) catch unreachable;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
var threads: [publisherCount]std.Thread = undefined;
|
||||
for (0..publisherCount) |i| {
|
||||
threads[i] = try std.Thread.spawn(.{}, Publisher.run, .{ &hub, allocator, i });
|
||||
}
|
||||
for (threads) |t| t.join();
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 0), sub.droppedCount());
|
||||
|
||||
var seenPerTag = [_]usize{0} ** publisherCount;
|
||||
var total: usize = 0;
|
||||
while (sub.waitMessage()) |msg| {
|
||||
defer msg.deinit(allocator);
|
||||
total += 1;
|
||||
var it = std.mem.splitScalar(u8, msg.payload, ':');
|
||||
const tag = std.fmt.parseInt(usize, it.next().?, 10) catch unreachable;
|
||||
seenPerTag[tag] += 1;
|
||||
if (total == publisherCount * perPublisher) break;
|
||||
}
|
||||
try std.testing.expectEqual(@as(usize, publisherCount * perPublisher), total);
|
||||
for (seenPerTag) |count| try std.testing.expectEqual(@as(usize, perPublisher), count);
|
||||
}
|
||||
|
||||
test "subscribe leaves nothing registered when it fails partway through" {
|
||||
var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 });
|
||||
var hub = Hub.init(failing.allocator());
|
||||
defer hub.deinit();
|
||||
|
||||
try std.testing.expectError(error.OutOfMemory, hub.subscribe(&.{ "a", "b", "c", "d" }, defaultCapacity));
|
||||
// Nothing should be left registered anywhere after the unwind.
|
||||
try std.testing.expectEqual(@as(usize, 0), hub.channelSubscriberCount("a"));
|
||||
try std.testing.expectEqual(@as(usize, 0), hub.channelSubscriberCount("b"));
|
||||
}
|
||||
|
||||
test "push under memory pressure drops the message and counts it instead of leaking" {
|
||||
var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 });
|
||||
const allocator = failing.allocator();
|
||||
const channels = try std.testing.allocator.alloc([]u8, 0);
|
||||
var sub = Subscriber.init(allocator, defaultCapacity, channels);
|
||||
defer sub.deinitContents();
|
||||
|
||||
const channel = try std.testing.allocator.dupe(u8, "ch");
|
||||
const payload = try std.testing.allocator.dupe(u8, "x");
|
||||
sub.push(.{ .channel = channel, .payload = payload });
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 1), sub.droppedCount());
|
||||
try std.testing.expectEqual(@as(usize, 0), sub.queue.items.len);
|
||||
}
|
||||
|
||||
test "publish tolerates an out-of-memory snapshot allocation without leaking" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{"ch"}, defaultCapacity);
|
||||
defer destroy(&hub, sub);
|
||||
|
||||
var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 });
|
||||
try std.testing.expectError(error.OutOfMemory, hub.publish(failing.allocator(), "ch", "x"));
|
||||
}
|
||||
|
||||
test "subscribe assigns each subscriber a distinct, increasing id" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const a = try hub.subscribe(&.{"ch"}, defaultCapacity);
|
||||
defer destroy(&hub, a);
|
||||
const b = try hub.subscribe(&.{"ch"}, defaultCapacity);
|
||||
defer destroy(&hub, b);
|
||||
|
||||
try std.testing.expect(a.id != b.id);
|
||||
try std.testing.expect(b.id > a.id);
|
||||
}
|
||||
|
||||
test "unsubscribe from a subset of channels leaves the rest delivering" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{ "weather", "orders", "alerts" }, defaultCapacity);
|
||||
defer destroy(&hub, sub);
|
||||
|
||||
const removed = try hub.unsubscribe(sub.id, &.{"orders"});
|
||||
try std.testing.expectEqual(@as(usize, 1), removed);
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 0), try hub.publish(allocator, "orders", "too-late"));
|
||||
try std.testing.expectEqual(@as(usize, 1), try hub.publish(allocator, "weather", "sunny"));
|
||||
try std.testing.expectEqual(@as(usize, 1), try hub.publish(allocator, "alerts", "fire"));
|
||||
|
||||
const m1 = sub.waitMessage().?;
|
||||
defer m1.deinit(allocator);
|
||||
try std.testing.expectEqualStrings("sunny", m1.payload);
|
||||
const m2 = sub.waitMessage().?;
|
||||
defer m2.deinit(allocator);
|
||||
try std.testing.expectEqualStrings("fire", m2.payload);
|
||||
|
||||
// The channel-name registry itself must not have been disturbed, so a
|
||||
// brand new subscriber can still register under the vacated channel.
|
||||
try std.testing.expectEqual(@as(usize, 0), hub.channelSubscriberCount("orders"));
|
||||
try std.testing.expectEqual(@as(usize, 1), hub.channelSubscriberCount("weather"));
|
||||
}
|
||||
|
||||
test "unsubscribe naming a channel the subscriber was never on counts zero, no error" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{"weather"}, defaultCapacity);
|
||||
defer destroy(&hub, sub);
|
||||
|
||||
const removed = try hub.unsubscribe(sub.id, &.{"nonexistent"});
|
||||
try std.testing.expectEqual(@as(usize, 0), removed);
|
||||
try std.testing.expectEqual(@as(usize, 1), try hub.publish(allocator, "weather", "still-here"));
|
||||
const msg = sub.waitMessage().?;
|
||||
defer msg.deinit(allocator);
|
||||
}
|
||||
|
||||
test "unsubscribe from every remaining channel closes the subscriber outright" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{ "a", "b" }, defaultCapacity);
|
||||
defer hub.destroySubscriber(sub);
|
||||
|
||||
const removed = try hub.unsubscribe(sub.id, &.{ "a", "b" });
|
||||
try std.testing.expectEqual(@as(usize, 2), removed);
|
||||
|
||||
try std.testing.expect(sub.isClosed());
|
||||
try std.testing.expect(sub.waitMessage() == null);
|
||||
try std.testing.expectEqual(@as(usize, 0), hub.channelSubscriberCount("a"));
|
||||
try std.testing.expectEqual(@as(usize, 0), hub.channelSubscriberCount("b"));
|
||||
|
||||
// The id is gone too - a second unsubscribe by the same id must not
|
||||
// find a dangling subscriber.
|
||||
try std.testing.expectError(error.UnknownSubscriber, hub.unsubscribe(sub.id, &.{"a"}));
|
||||
|
||||
// Mirrors the normal disconnect path: `closeSubscriber` still runs
|
||||
// afterward (e.g. from `handleSubscribe`'s push loop noticing the
|
||||
// close) and must be a harmless no-op over an already-empty channel set.
|
||||
hub.closeSubscriber(sub);
|
||||
}
|
||||
|
||||
test "unsubscribe wakes a thread already blocked in waitMessage when it empties the set" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{"ch"}, defaultCapacity);
|
||||
defer hub.destroySubscriber(sub);
|
||||
|
||||
const Waiter = struct {
|
||||
fn run(s: *Subscriber, result: *?Message) void {
|
||||
result.* = s.waitMessage();
|
||||
}
|
||||
};
|
||||
var result: ?Message = undefined;
|
||||
result = Message{ .channel = &.{}, .payload = &.{} };
|
||||
const t = try std.Thread.spawn(.{}, Waiter.run, .{ sub, &result });
|
||||
_ = try hub.unsubscribe(sub.id, &.{"ch"});
|
||||
t.join();
|
||||
try std.testing.expect(result == null);
|
||||
}
|
||||
|
||||
test "unsubscribe rejects an id nobody ever issued" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
try std.testing.expectError(error.UnknownSubscriber, hub.unsubscribe(999, &.{"ch"}));
|
||||
}
|
||||
|
||||
test "closeSubscriber removes the id, so a later unsubscribe by that id is rejected" {
|
||||
const allocator = std.testing.allocator;
|
||||
var hub = Hub.init(allocator);
|
||||
defer hub.deinit();
|
||||
|
||||
const sub = try hub.subscribe(&.{"ch"}, defaultCapacity);
|
||||
const id = sub.id;
|
||||
hub.closeSubscriber(sub);
|
||||
defer hub.destroySubscriber(sub);
|
||||
|
||||
try std.testing.expectError(error.UnknownSubscriber, hub.unsubscribe(id, &.{"ch"}));
|
||||
}
|
||||
+3868
File diff suppressed because it is too large.
Load diff
+2938
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,808 @@
|
||||
const std = @import("std");
|
||||
const encode = @import("encode.zig");
|
||||
const dbmod = @import("db.zig");
|
||||
const walmod = @import("wal.zig");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
const Kind = encode.Kind;
|
||||
const Db = dbmod.Db;
|
||||
const WriteBatch = dbmod.WriteBatch;
|
||||
|
||||
// Physical WAL-tailing replication: reads the primary's on-disk wal.log file
|
||||
// directly (same binary frame format as wal.zig: [crc32 u32][seq u64][kind u8]
|
||||
// [keyLen u32][valLen u32][key][value], or a batch frame marked by kind byte
|
||||
// 0xB0 holding a count then per-op [kind][seq][keyLen][valLen][key][value]),
|
||||
// independently of wal.zig's private replay-into-memtable path, and applies
|
||||
// newly appended, crc-verified, complete frames onto a separate replica Db.
|
||||
//
|
||||
// This mirrors real physical log shipping (e.g. Postgres WAL streaming): it
|
||||
// only sees what is still in the WAL file. Db.flush() truncates the primary's
|
||||
// wal.log to empty once its memtable is durable in an SSTable (see db.zig
|
||||
// flush()), so a tailer that has fallen behind a flush can no longer recover
|
||||
// the truncated records from the WAL alone - poll() detects that condition
|
||||
// (current file size smaller than the tailer's already-consumed offset) and
|
||||
// fails loudly with error.WalTruncatedAheadOfTailer instead of silently
|
||||
// resuming from zero and dropping data. The caller is expected to recover via
|
||||
// fullResync(), which re-bases the replica from a live Snapshot (paying the
|
||||
// same "logical snapshot, not a physical replica of on-disk levels/history"
|
||||
// cost already documented for backup.zig) and restarts tailing from the
|
||||
// (now known-empty) WAL, exactly mirroring a real base-backup-plus-streaming
|
||||
// recovery: bulk-copy current state, then resume incremental log shipping
|
||||
// from the offset the bulk copy was taken at.
|
||||
|
||||
const headerLen = 8 + 1 + 4 + 4;
|
||||
const batchMarker: u8 = 0xB0;
|
||||
const batchHeader = 8 + 1 + 4;
|
||||
const opHeader = 1 + 8 + 4 + 4;
|
||||
const walName = "wal.log";
|
||||
const maxTailRead: usize = 256 * 1024 * 1024;
|
||||
|
||||
pub const ReplicateError = error{ WalTruncatedAheadOfTailer, WalGenerationChanged };
|
||||
|
||||
/// Reads the WAL generation stamp (see wal.zig's currentGeneration) for a
|
||||
/// primary directory given only its path, for callers (WalTailer.poll,
|
||||
/// Replicator construction from a persisted offset) that do not hold a live
|
||||
/// *Db handle to the primary and so cannot reuse its already-open Dir.
|
||||
fn readGenerationByPath(primaryDirPath: []const u8) !u64 {
|
||||
var dir = try std.fs.cwd().openDir(primaryDirPath, .{});
|
||||
defer dir.close();
|
||||
return walmod.currentGeneration(dir);
|
||||
}
|
||||
|
||||
pub const OwnedRecord = struct {
|
||||
kind: Kind,
|
||||
seq: u64,
|
||||
key: []u8,
|
||||
value: []u8,
|
||||
|
||||
pub fn deinit(self: OwnedRecord, allocator: Allocator) void {
|
||||
allocator.free(self.key);
|
||||
allocator.free(self.value);
|
||||
}
|
||||
};
|
||||
|
||||
fn freeAll(allocator: Allocator, records: []OwnedRecord) void {
|
||||
for (records) |record| record.deinit(allocator);
|
||||
allocator.free(records);
|
||||
}
|
||||
|
||||
/// Tails a single WAL file from a byte offset, decoding whole, crc-verified
|
||||
/// frames appended since the last poll. Never advances past an incomplete or
|
||||
/// corrupt tail frame (mirrors wal.zig's own torn-write tolerance), so a
|
||||
/// concurrent writer mid-append is simply not observed until its frame is
|
||||
/// fully flushed and synced.
|
||||
pub const WalTailer = struct {
|
||||
allocator: Allocator,
|
||||
file: std.fs.File,
|
||||
offset: u64,
|
||||
primaryDir: []u8,
|
||||
generation: u64,
|
||||
|
||||
pub fn open(allocator: Allocator, primaryDirPath: []const u8, startOffset: u64, startGeneration: u64) !WalTailer {
|
||||
const path = try std.fs.path.join(allocator, &.{ primaryDirPath, walName });
|
||||
defer allocator.free(path);
|
||||
const file = try std.fs.cwd().openFile(path, .{});
|
||||
errdefer file.close();
|
||||
const dirCopy = try allocator.dupe(u8, primaryDirPath);
|
||||
return .{ .allocator = allocator, .file = file, .offset = startOffset, .primaryDir = dirCopy, .generation = startGeneration };
|
||||
}
|
||||
|
||||
pub fn deinit(self: *WalTailer) void {
|
||||
self.file.close();
|
||||
self.allocator.free(self.primaryDir);
|
||||
}
|
||||
|
||||
/// Returns newly observed records in append order. Caller owns the
|
||||
/// returned slice and each record's key/value (free via freeAll or by
|
||||
/// calling deinit() on each record then freeing the slice). Checks the
|
||||
/// primary's WAL generation stamp before anything else: a changed
|
||||
/// generation means the primary truncated (flushed) since the last poll,
|
||||
/// which the caller must recover from via fullResync() regardless of how
|
||||
/// the raw byte offsets happen to compare (see wal.zig's bumpGeneration
|
||||
/// doc comment for why the byte-size comparison alone is not enough).
|
||||
pub fn poll(self: *WalTailer) ![]OwnedRecord {
|
||||
const nowGeneration = try readGenerationByPath(self.primaryDir);
|
||||
if (nowGeneration != self.generation) return ReplicateError.WalGenerationChanged;
|
||||
|
||||
const endPos = try self.file.getEndPos();
|
||||
if (endPos < self.offset) return ReplicateError.WalTruncatedAheadOfTailer;
|
||||
if (endPos == self.offset) return &.{};
|
||||
|
||||
try self.file.seekTo(self.offset);
|
||||
const data = try self.file.readToEndAlloc(self.allocator, maxTailRead);
|
||||
defer self.allocator.free(data);
|
||||
|
||||
var out = std.ArrayList(OwnedRecord).init(self.allocator);
|
||||
errdefer {
|
||||
for (out.items) |record| record.deinit(self.allocator);
|
||||
out.deinit();
|
||||
}
|
||||
|
||||
var pos: usize = 0;
|
||||
while (pos + 4 + batchHeader <= data.len) {
|
||||
const storedCrc = encode.readU32(data, pos);
|
||||
const bodyStart = pos + 4;
|
||||
const bodyLen = frameBodyLen(data, bodyStart) orelse break;
|
||||
if (bodyStart + bodyLen > data.len) break;
|
||||
const body = data[bodyStart .. bodyStart + bodyLen];
|
||||
if (encode.crc(body) != storedCrc) break;
|
||||
|
||||
try decodeFrame(self.allocator, body, &out);
|
||||
|
||||
pos = bodyStart + bodyLen;
|
||||
}
|
||||
|
||||
self.offset += pos;
|
||||
return try out.toOwnedSlice();
|
||||
}
|
||||
};
|
||||
|
||||
fn frameBodyLen(data: []const u8, bodyStart: usize) ?usize {
|
||||
const marker = data[bodyStart + 8];
|
||||
if (marker == batchMarker) {
|
||||
if (bodyStart + batchHeader > data.len) return null;
|
||||
const opCount = encode.readU32(data, bodyStart + 9);
|
||||
var p = bodyStart + batchHeader;
|
||||
var n: u32 = 0;
|
||||
while (n < opCount) : (n += 1) {
|
||||
if (p + opHeader > data.len) return null;
|
||||
const keyLen = encode.readU32(data, p + 9);
|
||||
const valLen = encode.readU32(data, p + 13);
|
||||
const end = p + opHeader + keyLen + valLen;
|
||||
if (end > data.len) return null;
|
||||
p = end;
|
||||
}
|
||||
return p - bodyStart;
|
||||
}
|
||||
if (bodyStart + headerLen > data.len) return null;
|
||||
const keyLen = encode.readU32(data, bodyStart + 9);
|
||||
const valLen = encode.readU32(data, bodyStart + 13);
|
||||
return headerLen + keyLen + valLen;
|
||||
}
|
||||
|
||||
fn decodeFrame(allocator: Allocator, body: []const u8, out: *std.ArrayList(OwnedRecord)) !void {
|
||||
if (body[8] == batchMarker) {
|
||||
const opCount = encode.readU32(body, 9);
|
||||
var p: usize = batchHeader;
|
||||
var n: u32 = 0;
|
||||
while (n < opCount) : (n += 1) {
|
||||
const kind: Kind = @enumFromInt(body[p]);
|
||||
const seq = encode.readU64(body, p + 1);
|
||||
const keyLen = encode.readU32(body, p + 9);
|
||||
const valLen = encode.readU32(body, p + 13);
|
||||
const keyStart = p + opHeader;
|
||||
const key = try allocator.dupe(u8, body[keyStart .. keyStart + keyLen]);
|
||||
errdefer allocator.free(key);
|
||||
const value = try allocator.dupe(u8, body[keyStart + keyLen .. keyStart + keyLen + valLen]);
|
||||
try out.append(.{ .kind = kind, .seq = seq, .key = key, .value = value });
|
||||
p = keyStart + keyLen + valLen;
|
||||
}
|
||||
return;
|
||||
}
|
||||
const seq = encode.readU64(body, 0);
|
||||
const kind: Kind = @enumFromInt(body[8]);
|
||||
const keyLen = encode.readU32(body, 9);
|
||||
const valLen = encode.readU32(body, 13);
|
||||
const key = try allocator.dupe(u8, body[headerLen .. headerLen + keyLen]);
|
||||
errdefer allocator.free(key);
|
||||
const value = try allocator.dupe(u8, body[headerLen + keyLen .. headerLen + keyLen + valLen]);
|
||||
try out.append(.{ .kind = kind, .seq = seq, .key = key, .value = value });
|
||||
}
|
||||
|
||||
/// Ties a WalTailer to a target Db, applying newly tailed records as a single
|
||||
/// WriteBatch per syncOnce(). Replicated writes get FRESH sequence numbers on
|
||||
/// the target (Db.apply always assigns from its own nextSeq), so the replica
|
||||
/// mirrors primary key/value state but not primary's original seq history -
|
||||
/// same documented tradeoff as backup.zig's restore.
|
||||
pub const Replicator = struct {
|
||||
allocator: Allocator,
|
||||
tailer: WalTailer,
|
||||
target: *Db,
|
||||
|
||||
pub fn init(allocator: Allocator, primaryDirPath: []const u8, target: *Db) !Replicator {
|
||||
const gen = try readGenerationByPath(primaryDirPath);
|
||||
return initFromOffset(allocator, primaryDirPath, target, 0, gen);
|
||||
}
|
||||
|
||||
pub fn initFromOffset(allocator: Allocator, primaryDirPath: []const u8, target: *Db, startOffset: u64, startGeneration: u64) !Replicator {
|
||||
const tailer = try WalTailer.open(allocator, primaryDirPath, startOffset, startGeneration);
|
||||
return .{ .allocator = allocator, .tailer = tailer, .target = target };
|
||||
}
|
||||
|
||||
pub fn deinit(self: *Replicator) void {
|
||||
self.tailer.deinit();
|
||||
}
|
||||
|
||||
pub fn offset(self: *const Replicator) u64 {
|
||||
return self.tailer.offset;
|
||||
}
|
||||
|
||||
pub fn generation(self: *const Replicator) u64 {
|
||||
return self.tailer.generation;
|
||||
}
|
||||
|
||||
/// Applies every record newly visible in the primary's WAL since the last
|
||||
/// call. Returns the number of records applied (0 if nothing new).
|
||||
/// Propagates error.WalTruncatedAheadOfTailer if the primary flushed
|
||||
/// (truncating its WAL) before this replicator consumed everything - the
|
||||
/// caller must recover via fullResync().
|
||||
pub fn syncOnce(self: *Replicator) !usize {
|
||||
const records = try self.tailer.poll();
|
||||
defer freeAll(self.allocator, records);
|
||||
if (records.len == 0) return 0;
|
||||
|
||||
var batch = WriteBatch.init(self.allocator);
|
||||
defer batch.deinit();
|
||||
for (records) |record| {
|
||||
switch (record.kind) {
|
||||
.value => try batch.put(record.key, record.value),
|
||||
.tombstone => try batch.delete(record.key),
|
||||
}
|
||||
}
|
||||
try self.target.apply(&batch);
|
||||
return records.len;
|
||||
}
|
||||
|
||||
/// Recovers from a fallen-behind tailer (or bootstraps a brand new
|
||||
/// replica) by copying every live key/value from a fresh primary
|
||||
/// Snapshot into the target in chunks, then re-bases the tailer at the
|
||||
/// primary's current WAL offset. Snapshot() flushes the primary first
|
||||
/// (see db.zig), so immediately after this call the primary's wal.log is
|
||||
/// known to be empty - offset 0 is therefore always the correct restart
|
||||
/// point. Requires no concurrent writer on the primary during the call
|
||||
/// (same single-writer assumption as the rest of this project - Db has
|
||||
/// no internal cross-thread synchronization).
|
||||
pub fn fullResync(self: *Replicator, primary: *Db, chunkSize: usize) !usize {
|
||||
const snap = try primary.snapshot();
|
||||
var iter = try snap.scan(self.allocator, null, null);
|
||||
defer iter.deinit();
|
||||
|
||||
var batch = WriteBatch.init(self.allocator);
|
||||
defer batch.deinit();
|
||||
var applied: usize = 0;
|
||||
var staged: usize = 0;
|
||||
while (try iter.next()) |entry| {
|
||||
try batch.put(entry.key, entry.value);
|
||||
staged += 1;
|
||||
applied += 1;
|
||||
if (staged >= chunkSize) {
|
||||
try self.target.apply(&batch);
|
||||
batch.deinit();
|
||||
batch = WriteBatch.init(self.allocator);
|
||||
staged = 0;
|
||||
}
|
||||
}
|
||||
if (staged > 0) try self.target.apply(&batch);
|
||||
|
||||
self.tailer.offset = 0;
|
||||
// snapshot() above flushes the primary if it had unflushed writes,
|
||||
// which bumps its WAL generation (see wal.zig) - re-baseline here so
|
||||
// the very next poll() does not immediately see a "changed"
|
||||
// generation caused by this resync itself.
|
||||
self.tailer.generation = try walmod.currentGeneration(primary.dir);
|
||||
return applied;
|
||||
}
|
||||
};
|
||||
|
||||
const stateMagic = "CRPS";
|
||||
const stateFormatVersion: u32 = 2;
|
||||
|
||||
/// A tail offset paired with the WAL generation it was captured against.
|
||||
/// Persisting both (not just the offset) is what lets a brand-new process -
|
||||
/// which never saw the in-memory generation check in WalTailer.poll() - tell
|
||||
/// whether the primary truncated its WAL while this process was not running
|
||||
/// at all, the same blind spot the in-memory check closes for a live one.
|
||||
pub const PersistedState = struct {
|
||||
offset: u64,
|
||||
generation: u64,
|
||||
};
|
||||
const stateFileName = "replicate.state";
|
||||
const stateTmpName = "replicate.state.tmp";
|
||||
const maxStateRead: usize = 64 * 1024;
|
||||
|
||||
pub const StateError = error{
|
||||
BadMagic,
|
||||
UnsupportedVersion,
|
||||
Truncated,
|
||||
Corrupt,
|
||||
PrimaryMismatch,
|
||||
};
|
||||
|
||||
/// Atomically persists a replicator's tail offset next to the replica's own
|
||||
/// data (in replicaDirPath), so a later, brand-new `cobalt replicate --once`
|
||||
/// process - which has no in-memory Replicator to resume from - can pick up
|
||||
/// tailing exactly where the previous invocation left off instead of paying
|
||||
/// for a full snapshot fullResync() every single time it starts. Also stores
|
||||
/// primaryDirPath so a mismatched later call (pointed at a different primary
|
||||
/// by mistake) is caught by loadOffset() rather than silently tailing the
|
||||
/// wrong WAL from a stale offset. Written via tmp-file + rename (same pattern
|
||||
/// as version.zig's writeCurrent), so a crash mid-write can never leave a
|
||||
/// half-written state file where a good one used to be.
|
||||
pub fn saveState(allocator: Allocator, replicaDirPath: []const u8, primaryDirPath: []const u8, offsetValue: u64, generationValue: u64) !void {
|
||||
var dir = try std.fs.cwd().openDir(replicaDirPath, .{});
|
||||
defer dir.close();
|
||||
|
||||
var body = std.ArrayList(u8).init(allocator);
|
||||
defer body.deinit();
|
||||
const writer = encode.ByteWriter.init(&body);
|
||||
try writer.slice(stateMagic);
|
||||
try writer.putU32(stateFormatVersion);
|
||||
try writer.putU64(offsetValue);
|
||||
try writer.putU64(generationValue);
|
||||
try writer.putU32(@intCast(primaryDirPath.len));
|
||||
try writer.slice(primaryDirPath);
|
||||
try writer.putU32(encode.crc(body.items));
|
||||
|
||||
{
|
||||
const file = try dir.createFile(stateTmpName, .{ .truncate = true });
|
||||
defer file.close();
|
||||
try file.writeAll(body.items);
|
||||
try file.sync();
|
||||
}
|
||||
try dir.rename(stateTmpName, stateFileName);
|
||||
}
|
||||
|
||||
/// Loads a previously persisted tail offset for a replica directory. Returns
|
||||
/// null (not an error) when no state file has ever been written there yet -
|
||||
/// callers treat that as "fresh replica, bootstrap via fullResync()" exactly
|
||||
/// like before this existed. Returns StateError.PrimaryMismatch if the file
|
||||
/// was written while tailing a different primary directory, so a replica
|
||||
/// can never silently resume against the wrong primary from a stale offset.
|
||||
pub fn loadOffset(allocator: Allocator, replicaDirPath: []const u8, primaryDirPath: []const u8) !?PersistedState {
|
||||
var dir = try std.fs.cwd().openDir(replicaDirPath, .{});
|
||||
defer dir.close();
|
||||
|
||||
const data = dir.readFileAlloc(allocator, stateFileName, maxStateRead) catch |err| switch (err) {
|
||||
error.FileNotFound => return null,
|
||||
else => return err,
|
||||
};
|
||||
defer allocator.free(data);
|
||||
|
||||
const fixedLen = 4 + 4 + 8 + 8 + 4;
|
||||
if (data.len < fixedLen + 4) return StateError.Truncated;
|
||||
if (!std.mem.eql(u8, data[0..4], stateMagic)) return StateError.BadMagic;
|
||||
const version = encode.readU32(data, 4);
|
||||
if (version != stateFormatVersion) return StateError.UnsupportedVersion;
|
||||
const offsetValue = encode.readU64(data, 8);
|
||||
const generationValue = encode.readU64(data, 16);
|
||||
const pathLen = encode.readU32(data, 24);
|
||||
const pathStart = fixedLen;
|
||||
if (data.len != pathStart + pathLen + 4) return StateError.Truncated;
|
||||
const storedPath = data[pathStart .. pathStart + pathLen];
|
||||
const storedCrc = encode.readU32(data, pathStart + pathLen);
|
||||
if (encode.crc(data[0 .. pathStart + pathLen]) != storedCrc) return StateError.Corrupt;
|
||||
if (!std.mem.eql(u8, storedPath, primaryDirPath)) return StateError.PrimaryMismatch;
|
||||
|
||||
return .{ .offset = offsetValue, .generation = generationValue };
|
||||
}
|
||||
|
||||
test "replicator mirrors puts and deletes from a live primary" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
const primaryDir = try std.fs.path.join(allocator, &.{ base, "primary" });
|
||||
defer allocator.free(primaryDir);
|
||||
const replicaDir = try std.fs.path.join(allocator, &.{ base, "replica" });
|
||||
defer allocator.free(replicaDir);
|
||||
|
||||
var primary = try Db.open(allocator, primaryDir, .{});
|
||||
defer primary.close();
|
||||
var replica = try Db.open(allocator, replicaDir, .{});
|
||||
defer replica.close();
|
||||
|
||||
try primary.put("alpha", "one");
|
||||
try primary.put("bravo", "two");
|
||||
|
||||
var rep = try Replicator.init(allocator, primaryDir, &replica);
|
||||
defer rep.deinit();
|
||||
|
||||
const applied = try rep.syncOnce();
|
||||
try std.testing.expectEqual(@as(usize, 2), applied);
|
||||
|
||||
const a = try replica.get(allocator, "alpha");
|
||||
defer if (a) |v| allocator.free(v);
|
||||
try std.testing.expectEqualStrings("one", a.?);
|
||||
const b = try replica.get(allocator, "bravo");
|
||||
defer if (b) |v| allocator.free(v);
|
||||
try std.testing.expectEqualStrings("two", b.?);
|
||||
|
||||
try primary.delete("alpha");
|
||||
const applied2 = try rep.syncOnce();
|
||||
try std.testing.expectEqual(@as(usize, 1), applied2);
|
||||
|
||||
const a2 = try replica.get(allocator, "alpha");
|
||||
try std.testing.expectEqual(@as(?[]u8, null), a2);
|
||||
}
|
||||
|
||||
test "replicator syncOnce is idempotent between polls" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
const primaryDir = try std.fs.path.join(allocator, &.{ base, "primary" });
|
||||
defer allocator.free(primaryDir);
|
||||
const replicaDir = try std.fs.path.join(allocator, &.{ base, "replica" });
|
||||
defer allocator.free(replicaDir);
|
||||
|
||||
var primary = try Db.open(allocator, primaryDir, .{});
|
||||
defer primary.close();
|
||||
var replica = try Db.open(allocator, replicaDir, .{});
|
||||
defer replica.close();
|
||||
|
||||
var rep = try Replicator.init(allocator, primaryDir, &replica);
|
||||
defer rep.deinit();
|
||||
|
||||
try primary.put("k1", "v1");
|
||||
try std.testing.expectEqual(@as(usize, 1), try rep.syncOnce());
|
||||
try std.testing.expectEqual(@as(usize, 0), try rep.syncOnce());
|
||||
|
||||
try primary.put("k2", "v2");
|
||||
try std.testing.expectEqual(@as(usize, 1), try rep.syncOnce());
|
||||
try std.testing.expectEqual(@as(usize, 0), try rep.syncOnce());
|
||||
|
||||
var count: usize = 0;
|
||||
var iter = try replica.scan(allocator, null, null);
|
||||
defer iter.deinit();
|
||||
while (try iter.next()) |_| count += 1;
|
||||
try std.testing.expectEqual(@as(usize, 2), count);
|
||||
}
|
||||
|
||||
test "replicator applies batched writes from Db.apply as one poll" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
const primaryDir = try std.fs.path.join(allocator, &.{ base, "primary" });
|
||||
defer allocator.free(primaryDir);
|
||||
const replicaDir = try std.fs.path.join(allocator, &.{ base, "replica" });
|
||||
defer allocator.free(replicaDir);
|
||||
|
||||
var primary = try Db.open(allocator, primaryDir, .{});
|
||||
defer primary.close();
|
||||
var replica = try Db.open(allocator, replicaDir, .{});
|
||||
defer replica.close();
|
||||
|
||||
var writeBatch = WriteBatch.init(allocator);
|
||||
defer writeBatch.deinit();
|
||||
try writeBatch.put("x", "1");
|
||||
try writeBatch.put("y", "2");
|
||||
try writeBatch.delete("z");
|
||||
try primary.apply(&writeBatch);
|
||||
|
||||
var rep = try Replicator.init(allocator, primaryDir, &replica);
|
||||
defer rep.deinit();
|
||||
const applied = try rep.syncOnce();
|
||||
try std.testing.expectEqual(@as(usize, 3), applied);
|
||||
|
||||
const x = try replica.get(allocator, "x");
|
||||
defer if (x) |v| allocator.free(v);
|
||||
try std.testing.expectEqualStrings("1", x.?);
|
||||
const z = try replica.get(allocator, "z");
|
||||
try std.testing.expectEqual(@as(?[]u8, null), z);
|
||||
}
|
||||
|
||||
test "replicator poll does not consume a torn tail frame" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
const primaryDir = try std.fs.path.join(allocator, &.{ base, "primary" });
|
||||
defer allocator.free(primaryDir);
|
||||
const replicaDir = try std.fs.path.join(allocator, &.{ base, "replica" });
|
||||
defer allocator.free(replicaDir);
|
||||
|
||||
var primary = try Db.open(allocator, primaryDir, .{});
|
||||
defer primary.close();
|
||||
var replica = try Db.open(allocator, replicaDir, .{});
|
||||
defer replica.close();
|
||||
|
||||
var rep = try Replicator.init(allocator, primaryDir, &replica);
|
||||
defer rep.deinit();
|
||||
|
||||
try primary.put("good", "data");
|
||||
try std.testing.expectEqual(@as(usize, 1), try rep.syncOnce());
|
||||
|
||||
// Simulate a writer mid-append: raw bytes appended after the good frame
|
||||
// that never form a complete, crc-verified frame.
|
||||
try primary.wal.file.seekFromEnd(0);
|
||||
try primary.wal.file.writeAll(&[_]u8{ 0xAA, 0xBB, 0xCC });
|
||||
try primary.wal.sync();
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 0), try rep.syncOnce());
|
||||
|
||||
// The writer "finishes": a real, complete frame lands after the torn
|
||||
// bytes are naturally never completed (matches wal.zig's own tolerance -
|
||||
// torn tail bytes are simply never a valid prefix of a later frame, so
|
||||
// this models a stalled partial write, not a resumed one).
|
||||
var count: usize = 0;
|
||||
var iter = try replica.scan(allocator, null, null);
|
||||
defer iter.deinit();
|
||||
while (try iter.next()) |_| count += 1;
|
||||
try std.testing.expectEqual(@as(usize, 1), count);
|
||||
}
|
||||
|
||||
test "replicator surfaces truncation when it falls behind a primary flush" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
const primaryDir = try std.fs.path.join(allocator, &.{ base, "primary" });
|
||||
defer allocator.free(primaryDir);
|
||||
const replicaDir = try std.fs.path.join(allocator, &.{ base, "replica" });
|
||||
defer allocator.free(replicaDir);
|
||||
|
||||
var primary = try Db.open(allocator, primaryDir, .{});
|
||||
defer primary.close();
|
||||
var replica = try Db.open(allocator, replicaDir, .{});
|
||||
defer replica.close();
|
||||
|
||||
try primary.put("k", "v");
|
||||
|
||||
var rep = try Replicator.init(allocator, primaryDir, &replica);
|
||||
defer rep.deinit();
|
||||
// Advance the tailer's offset without applying, then flush the primary
|
||||
// (truncates wal.log to 0) so the tailer's remembered offset now points
|
||||
// past the end of the (now-empty) file.
|
||||
const staleRecords = try rep.tailer.poll();
|
||||
freeAll(allocator, staleRecords);
|
||||
try primary.flush();
|
||||
|
||||
// The primary's flush() bumped its WAL generation, so the generation
|
||||
// check in poll() fires before the (now also true, but weaker) raw
|
||||
// byte-offset comparison would have.
|
||||
try std.testing.expectError(ReplicateError.WalGenerationChanged, rep.syncOnce());
|
||||
|
||||
const recovered = try rep.fullResync(&primary, 500);
|
||||
try std.testing.expectEqual(@as(usize, 1), recovered);
|
||||
try std.testing.expectEqual(@as(u64, 0), rep.offset());
|
||||
|
||||
const v = try replica.get(allocator, "k");
|
||||
defer if (v) |val| allocator.free(val);
|
||||
try std.testing.expectEqualStrings("v", v.?);
|
||||
|
||||
// New writes after the resync are picked up normally through the
|
||||
// freshly-restarted tail at offset 0.
|
||||
try primary.put("k2", "v2");
|
||||
try std.testing.expectEqual(@as(usize, 1), try rep.syncOnce());
|
||||
const v2 = try replica.get(allocator, "k2");
|
||||
defer if (v2) |val| allocator.free(val);
|
||||
try std.testing.expectEqualStrings("v2", v2.?);
|
||||
}
|
||||
|
||||
test "replicator fullResync bootstraps a brand new replica" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
const primaryDir = try std.fs.path.join(allocator, &.{ base, "primary" });
|
||||
defer allocator.free(primaryDir);
|
||||
const replicaDir = try std.fs.path.join(allocator, &.{ base, "replica" });
|
||||
defer allocator.free(replicaDir);
|
||||
|
||||
var primary = try Db.open(allocator, primaryDir, .{});
|
||||
defer primary.close();
|
||||
var replica = try Db.open(allocator, replicaDir, .{});
|
||||
defer replica.close();
|
||||
|
||||
var i: usize = 0;
|
||||
while (i < 50) : (i += 1) {
|
||||
var keyBuf: [16]u8 = undefined;
|
||||
var valBuf: [16]u8 = undefined;
|
||||
const key = try std.fmt.bufPrint(&keyBuf, "key{d}", .{i});
|
||||
const val = try std.fmt.bufPrint(&valBuf, "val{d}", .{i});
|
||||
try primary.put(key, val);
|
||||
}
|
||||
|
||||
var rep = try Replicator.init(allocator, primaryDir, &replica);
|
||||
defer rep.deinit();
|
||||
const applied = try rep.fullResync(&primary, 7);
|
||||
try std.testing.expectEqual(@as(usize, 50), applied);
|
||||
|
||||
var count: usize = 0;
|
||||
var iter = try replica.scan(allocator, null, null);
|
||||
defer iter.deinit();
|
||||
while (try iter.next()) |_| count += 1;
|
||||
try std.testing.expectEqual(@as(usize, 50), count);
|
||||
|
||||
// Fresh writes after the bootstrap keep tailing correctly from offset 0.
|
||||
try primary.put("late", "arrival");
|
||||
try std.testing.expectEqual(@as(usize, 1), try rep.syncOnce());
|
||||
const v = try replica.get(allocator, "late");
|
||||
defer if (v) |val| allocator.free(val);
|
||||
try std.testing.expectEqualStrings("arrival", v.?);
|
||||
}
|
||||
|
||||
test "walTailer rejects a nonexistent primary directory" {
|
||||
const allocator = std.testing.allocator;
|
||||
try std.testing.expectError(error.FileNotFound, WalTailer.open(allocator, "/nonexistent", 0, 0));
|
||||
}
|
||||
|
||||
test "loadOffset returns null for a replica directory that never saved state" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
|
||||
const loaded = try loadOffset(allocator, base, "/some/primary");
|
||||
try std.testing.expectEqual(@as(?PersistedState, null), loaded);
|
||||
}
|
||||
|
||||
test "saveState then loadOffset round-trips the offset, generation, and primary path" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
|
||||
try saveState(allocator, base, "/data/primary", 4096, 7);
|
||||
const loaded = try loadOffset(allocator, base, "/data/primary");
|
||||
try std.testing.expectEqual(@as(u64, 4096), loaded.?.offset);
|
||||
try std.testing.expectEqual(@as(u64, 7), loaded.?.generation);
|
||||
}
|
||||
|
||||
test "saveState overwrites a prior state file with the newest offset and generation" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
|
||||
try saveState(allocator, base, "/data/primary", 10, 0);
|
||||
try saveState(allocator, base, "/data/primary", 20, 1);
|
||||
try saveState(allocator, base, "/data/primary", 30, 2);
|
||||
const loaded = try loadOffset(allocator, base, "/data/primary");
|
||||
try std.testing.expectEqual(@as(u64, 30), loaded.?.offset);
|
||||
try std.testing.expectEqual(@as(u64, 2), loaded.?.generation);
|
||||
}
|
||||
|
||||
test "loadOffset rejects a state file saved for a different primary" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
|
||||
try saveState(allocator, base, "/data/primary-a", 123, 0);
|
||||
try std.testing.expectError(StateError.PrimaryMismatch, loadOffset(allocator, base, "/data/primary-b"));
|
||||
}
|
||||
|
||||
test "loadOffset rejects a corrupted state file without crashing" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
|
||||
try saveState(allocator, base, "/data/primary", 55, 3);
|
||||
|
||||
{
|
||||
const file = try tmp.dir.openFile(stateFileName, .{ .mode = .read_write });
|
||||
defer file.close();
|
||||
try file.seekTo(9);
|
||||
try file.writeAll(&[_]u8{0xFF});
|
||||
}
|
||||
|
||||
try std.testing.expectError(StateError.Corrupt, loadOffset(allocator, base, "/data/primary"));
|
||||
}
|
||||
|
||||
test "loadOffset rejects a truncated state file" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
|
||||
try saveState(allocator, base, "/data/primary", 55, 3);
|
||||
|
||||
{
|
||||
const file = try tmp.dir.openFile(stateFileName, .{ .mode = .read_write });
|
||||
defer file.close();
|
||||
try file.setEndPos(6);
|
||||
}
|
||||
|
||||
try std.testing.expectError(StateError.Truncated, loadOffset(allocator, base, "/data/primary"));
|
||||
}
|
||||
|
||||
test "loadOffset rejects a bad magic header" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
|
||||
try saveState(allocator, base, "/data/primary", 55, 3);
|
||||
|
||||
{
|
||||
const file = try tmp.dir.openFile(stateFileName, .{ .mode = .read_write });
|
||||
defer file.close();
|
||||
try file.seekTo(0);
|
||||
try file.writeAll("XXXX");
|
||||
}
|
||||
|
||||
try std.testing.expectError(StateError.BadMagic, loadOffset(allocator, base, "/data/primary"));
|
||||
}
|
||||
|
||||
test "loadOffset rejects a state file written in the pre-generation v1 format" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
|
||||
var body = std.ArrayList(u8).init(allocator);
|
||||
defer body.deinit();
|
||||
const writer = encode.ByteWriter.init(&body);
|
||||
try writer.slice("CRPS");
|
||||
try writer.putU32(1);
|
||||
try writer.putU64(999);
|
||||
const primaryPath = "/data/primary";
|
||||
try writer.putU32(@intCast(primaryPath.len));
|
||||
try writer.slice(primaryPath);
|
||||
try writer.putU32(encode.crc(body.items));
|
||||
|
||||
const file = try tmp.dir.createFile(stateFileName, .{ .truncate = true });
|
||||
defer file.close();
|
||||
try file.writeAll(body.items);
|
||||
|
||||
try std.testing.expectError(StateError.UnsupportedVersion, loadOffset(allocator, base, "/data/primary"));
|
||||
}
|
||||
|
||||
test "saveState and loadOffset round-trip through a real replicate directory end to end" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
const base = try tmp.dir.realpathAlloc(allocator, ".");
|
||||
defer allocator.free(base);
|
||||
const primaryDir = try std.fs.path.join(allocator, &.{ base, "primary" });
|
||||
defer allocator.free(primaryDir);
|
||||
const replicaDir = try std.fs.path.join(allocator, &.{ base, "replica" });
|
||||
defer allocator.free(replicaDir);
|
||||
|
||||
var primary = try Db.open(allocator, primaryDir, .{});
|
||||
defer primary.close();
|
||||
var replica = try Db.open(allocator, replicaDir, .{});
|
||||
defer replica.close();
|
||||
|
||||
try primary.put("a", "1");
|
||||
try primary.put("b", "2");
|
||||
|
||||
var rep = try Replicator.init(allocator, primaryDir, &replica);
|
||||
defer rep.deinit();
|
||||
_ = try rep.syncOnce();
|
||||
try saveState(allocator, replicaDir, primaryDir, rep.offset(), rep.generation());
|
||||
const savedOffset = rep.offset();
|
||||
try std.testing.expect(savedOffset > 0);
|
||||
|
||||
// A brand new process would reopen the replica Db and reconstruct a
|
||||
// Replicator from the persisted offset instead of offset 0.
|
||||
var replica2 = try Db.open(allocator, replicaDir, .{});
|
||||
defer replica2.close();
|
||||
const resumed = (try loadOffset(allocator, replicaDir, primaryDir)).?;
|
||||
try std.testing.expectEqual(savedOffset, resumed.offset);
|
||||
try std.testing.expectEqual(rep.generation(), resumed.generation);
|
||||
|
||||
var rep2 = try Replicator.initFromOffset(allocator, primaryDir, &replica2, resumed.offset, resumed.generation);
|
||||
defer rep2.deinit();
|
||||
try primary.put("c", "3");
|
||||
const applied = try rep2.syncOnce();
|
||||
try std.testing.expectEqual(@as(usize, 1), applied);
|
||||
|
||||
const c = try replica2.get(allocator, "c");
|
||||
defer if (c) |v| allocator.free(v);
|
||||
try std.testing.expectEqualStrings("3", c.?);
|
||||
// The earlier keys were never re-applied through rep2 (it resumed past
|
||||
// them), yet they are still visible because replica2 is the same
|
||||
// on-disk replica Db that rep already wrote them into.
|
||||
const a = try replica2.get(allocator, "a");
|
||||
defer if (a) |v| allocator.free(v);
|
||||
try std.testing.expectEqualStrings("1", a.?);
|
||||
}
|
||||
+100
@@ -17,6 +17,80 @@ pub const BlockCache = @import("cache.zig").BlockCache;
|
||||
pub const compress = @import("compress.zig");
|
||||
pub const Transaction = @import("txn.zig").Transaction;
|
||||
pub const Stats = @import("stats.zig").Stats;
|
||||
pub const protocol = @import("protocol.zig");
|
||||
pub const Server = @import("server.zig").Server;
|
||||
pub const Client = @import("client.zig").Client;
|
||||
pub const row = @import("row.zig");
|
||||
pub const index = @import("index.zig");
|
||||
pub const query = @import("query.zig");
|
||||
pub const Table = query.Table;
|
||||
pub const Predicate = query.Predicate;
|
||||
pub const Page = query.Page;
|
||||
pub const expr = @import("expr.zig");
|
||||
pub const join = @import("join.zig");
|
||||
pub const JoinedRow = join.JoinedRow;
|
||||
pub const OuterJoinedRow = join.OuterJoinedRow;
|
||||
pub const RightOuterJoinedRow = join.RightOuterJoinedRow;
|
||||
pub const FullOuterJoinedRow = join.FullOuterJoinedRow;
|
||||
pub const txnTable = @import("txn_table.zig");
|
||||
pub const TableTxn = txnTable.TableTxn;
|
||||
pub const collections = @import("collections.zig");
|
||||
pub const Counter = collections.Counter;
|
||||
pub const Hash = collections.Hash;
|
||||
pub const Set = collections.Set;
|
||||
pub const List = collections.List;
|
||||
pub const ZSet = collections.ZSet;
|
||||
pub const backup = @import("backup.zig");
|
||||
pub const BackupError = backup.BackupError;
|
||||
pub const BackupInfo = backup.BackupInfo;
|
||||
pub const metrics = @import("metrics.zig");
|
||||
pub const MetricCounter = metrics.MetricCounter;
|
||||
pub const Gauge = metrics.Gauge;
|
||||
pub const Histogram = metrics.Histogram;
|
||||
pub const TokenBucket = metrics.TokenBucket;
|
||||
pub const SlidingWindowLimiter = metrics.SlidingWindowLimiter;
|
||||
pub const MetricsRegistry = metrics.Registry;
|
||||
pub const replicate = @import("replicate.zig");
|
||||
pub const Replicator = replicate.Replicator;
|
||||
pub const WalTailer = replicate.WalTailer;
|
||||
pub const ReplicateError = replicate.ReplicateError;
|
||||
pub const heap = @import("heap.zig");
|
||||
pub const MinHeap = heap.MinHeap;
|
||||
pub const version = @import("version.zig");
|
||||
pub const VersionSet = version.VersionSet;
|
||||
pub const VersionEdit = version.VersionEdit;
|
||||
pub const auth = @import("auth.zig");
|
||||
pub const AuthStore = auth.Store;
|
||||
pub const AuthRole = auth.Role;
|
||||
pub const AuthError = auth.AuthError;
|
||||
pub const session = @import("session.zig");
|
||||
pub const SessionStore = session.SessionStore;
|
||||
pub const ttl = @import("ttl.zig");
|
||||
pub const TtlResult = ttl.TtlResult;
|
||||
pub const collttl = @import("collttl.zig");
|
||||
pub const CollectionKind = collections.Kind;
|
||||
pub const concurrent = @import("concurrent.zig");
|
||||
pub const GuardedDb = concurrent.GuardedDb;
|
||||
pub const keyspace = @import("keyspace.zig");
|
||||
pub const repl = @import("repl.zig");
|
||||
pub const Repl = repl.Repl;
|
||||
pub const ReplError = repl.ReplError;
|
||||
pub const dump = @import("dump.zig");
|
||||
pub const aggregate = @import("aggregate.zig");
|
||||
pub const AggKind = aggregate.AggKind;
|
||||
pub const AggSpec = aggregate.AggSpec;
|
||||
pub const GroupBy = aggregate.GroupBy;
|
||||
pub const GroupRow = aggregate.GroupRow;
|
||||
pub const AggregateError = aggregate.AggregateError;
|
||||
pub const window = @import("window.zig");
|
||||
pub const WindowKind = window.WindowKind;
|
||||
pub const WindowSpec = window.WindowSpec;
|
||||
pub const WindowError = window.WindowError;
|
||||
pub const WindowFrameMode = window.FrameMode;
|
||||
pub const pubsub = @import("pubsub.zig");
|
||||
pub const PubSubHub = pubsub.Hub;
|
||||
pub const PubSubSubscriber = pubsub.Subscriber;
|
||||
pub const PubSubMessage = pubsub.Message;
|
||||
|
||||
test {
|
||||
_ = @import("encode.zig");
|
||||
@@ -30,4 +104,30 @@ test {
|
||||
_ = @import("txn.zig");
|
||||
_ = @import("stats.zig");
|
||||
_ = @import("db.zig");
|
||||
_ = @import("protocol.zig");
|
||||
_ = @import("server.zig");
|
||||
_ = @import("client.zig");
|
||||
_ = @import("row.zig");
|
||||
_ = @import("index.zig");
|
||||
_ = @import("query.zig");
|
||||
_ = @import("expr.zig");
|
||||
_ = @import("join.zig");
|
||||
_ = @import("txn_table.zig");
|
||||
_ = @import("collections.zig");
|
||||
_ = @import("backup.zig");
|
||||
_ = @import("metrics.zig");
|
||||
_ = @import("replicate.zig");
|
||||
_ = @import("heap.zig");
|
||||
_ = @import("version.zig");
|
||||
_ = @import("auth.zig");
|
||||
_ = @import("session.zig");
|
||||
_ = @import("ttl.zig");
|
||||
_ = @import("collttl.zig");
|
||||
_ = @import("concurrent.zig");
|
||||
_ = @import("keyspace.zig");
|
||||
_ = @import("repl.zig");
|
||||
_ = @import("aggregate.zig");
|
||||
_ = @import("dump.zig");
|
||||
_ = @import("window.zig");
|
||||
_ = @import("pubsub.zig");
|
||||
}
|
||||
+487
@@ -0,0 +1,487 @@
|
||||
const std = @import("std");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
|
||||
/// Dynamically typed value stored inside a row field, an index key, or a
|
||||
/// primary key. Kept small and copyable by tag; owned byte payloads (text,
|
||||
/// bytes) are heap-allocated and must be freed by whoever owns the Row/Value.
|
||||
pub const ValueTag = enum(u8) {
|
||||
none = 0,
|
||||
boolean = 1,
|
||||
int = 2,
|
||||
float = 3,
|
||||
text = 4,
|
||||
bytes = 5,
|
||||
};
|
||||
|
||||
pub const Value = union(ValueTag) {
|
||||
none: void,
|
||||
boolean: bool,
|
||||
int: i64,
|
||||
float: f64,
|
||||
text: []const u8,
|
||||
bytes: []const u8,
|
||||
|
||||
pub fn dupe(self: Value, allocator: Allocator) !Value {
|
||||
return switch (self) {
|
||||
.text => |s| .{ .text = try allocator.dupe(u8, s) },
|
||||
.bytes => |s| .{ .bytes = try allocator.dupe(u8, s) },
|
||||
else => self,
|
||||
};
|
||||
}
|
||||
|
||||
pub fn free(self: Value, allocator: Allocator) void {
|
||||
switch (self) {
|
||||
.text => |s| allocator.free(s),
|
||||
.bytes => |s| allocator.free(s),
|
||||
else => {},
|
||||
}
|
||||
}
|
||||
|
||||
pub fn eql(a: Value, b: Value) bool {
|
||||
return compare(a, b) == .eq;
|
||||
}
|
||||
|
||||
/// Parses a single CLI/query-language token into a typed Value, using the
|
||||
/// same `i:`/`f:`/`b:`/`s:` type-prefix convention throughout the CLI
|
||||
/// (unprefixed tokens default to text). The returned `.text`/`.bytes`
|
||||
/// payload (when present) is a slice into `text` - callers that need it
|
||||
/// to outlive `text` must `dupe` the result.
|
||||
pub fn parseCli(text: []const u8) error{InvalidValue}!Value {
|
||||
if (std.mem.startsWith(u8, text, "i:")) {
|
||||
const n = std.fmt.parseInt(i64, text[2..], 10) catch return error.InvalidValue;
|
||||
return .{ .int = n };
|
||||
} else if (std.mem.startsWith(u8, text, "f:")) {
|
||||
const f = std.fmt.parseFloat(f64, text[2..]) catch return error.InvalidValue;
|
||||
return .{ .float = f };
|
||||
} else if (std.mem.startsWith(u8, text, "b:")) {
|
||||
if (std.mem.eql(u8, text[2..], "true")) return .{ .boolean = true };
|
||||
if (std.mem.eql(u8, text[2..], "false")) return .{ .boolean = false };
|
||||
return error.InvalidValue;
|
||||
} else if (std.mem.startsWith(u8, text, "s:")) {
|
||||
return .{ .text = text[2..] };
|
||||
} else {
|
||||
return .{ .text = text };
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/// Total order across all values: first by type tag, then by payload. This
|
||||
/// keeps a well-defined ordering even when a field holds mixed types across
|
||||
/// rows, which the ordered byte encoding below must mirror exactly.
|
||||
pub fn compare(a: Value, b: Value) std.math.Order {
|
||||
const ta = @intFromEnum(@as(ValueTag, a));
|
||||
const tb = @intFromEnum(@as(ValueTag, b));
|
||||
if (ta != tb) return std.math.order(ta, tb);
|
||||
return switch (a) {
|
||||
.none => .eq,
|
||||
.boolean => |av| std.math.order(@intFromBool(av), @intFromBool(b.boolean)),
|
||||
.int => |av| std.math.order(av, b.int),
|
||||
.float => |av| std.math.order(av, b.float),
|
||||
.text => |av| std.mem.order(u8, av, b.text),
|
||||
.bytes => |av| std.mem.order(u8, av, b.bytes),
|
||||
};
|
||||
}
|
||||
|
||||
pub const Field = struct {
|
||||
name: []const u8,
|
||||
value: Value,
|
||||
};
|
||||
|
||||
/// Self-describing, schema-less row: an ordered list of named fields. Rows
|
||||
/// are the unit stored under a table's primary-key row keys and are what
|
||||
/// query.zig hands back from scans and index lookups.
|
||||
pub const Row = struct {
|
||||
allocator: Allocator,
|
||||
fields: std.ArrayList(Field),
|
||||
|
||||
pub fn init(allocator: Allocator) Row {
|
||||
return .{ .allocator = allocator, .fields = std.ArrayList(Field).init(allocator) };
|
||||
}
|
||||
|
||||
pub fn deinit(self: *Row) void {
|
||||
for (self.fields.items) |f| {
|
||||
self.allocator.free(f.name);
|
||||
f.value.free(self.allocator);
|
||||
}
|
||||
self.fields.deinit();
|
||||
}
|
||||
|
||||
/// Sets a field, replacing any existing field with the same name. Copies
|
||||
/// both the name and any owned payload of `value` into the row.
|
||||
pub fn set(self: *Row, name: []const u8, value: Value) !void {
|
||||
const owned = try value.dupe(self.allocator);
|
||||
errdefer owned.free(self.allocator);
|
||||
for (self.fields.items) |*f| {
|
||||
if (std.mem.eql(u8, f.name, name)) {
|
||||
f.value.free(self.allocator);
|
||||
f.value = owned;
|
||||
return;
|
||||
}
|
||||
}
|
||||
const nameCopy = try self.allocator.dupe(u8, name);
|
||||
try self.fields.append(.{ .name = nameCopy, .value = owned });
|
||||
}
|
||||
|
||||
pub fn get(self: *const Row, name: []const u8) ?Value {
|
||||
for (self.fields.items) |f| {
|
||||
if (std.mem.eql(u8, f.name, name)) return f.value;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
pub fn count(self: *const Row) usize {
|
||||
return self.fields.items.len;
|
||||
}
|
||||
|
||||
/// Serializes the row to a self-contained byte blob suitable as a Db
|
||||
/// value: [u32 fieldCount] { [u16 nameLen][name][u8 tag][payload] }*.
|
||||
pub fn encode(self: *const Row, allocator: Allocator) ![]u8 {
|
||||
var out = std.ArrayList(u8).init(allocator);
|
||||
errdefer out.deinit();
|
||||
try putU32(&out, @intCast(self.fields.items.len));
|
||||
for (self.fields.items) |f| {
|
||||
try putU16(&out, @intCast(f.name.len));
|
||||
try out.appendSlice(f.name);
|
||||
try out.append(@intFromEnum(@as(ValueTag, f.value)));
|
||||
switch (f.value) {
|
||||
.none => {},
|
||||
.boolean => |v| try out.append(if (v) 1 else 0),
|
||||
.int => |v| try putU64(&out, @bitCast(v)),
|
||||
.float => |v| try putU64(&out, @bitCast(v)),
|
||||
.text, .bytes => |s| {
|
||||
try putU32(&out, @intCast(s.len));
|
||||
try out.appendSlice(s);
|
||||
},
|
||||
}
|
||||
}
|
||||
return out.toOwnedSlice();
|
||||
}
|
||||
|
||||
pub fn decode(allocator: Allocator, data: []const u8) !Row {
|
||||
var row = Row.init(allocator);
|
||||
errdefer row.deinit();
|
||||
var pos: usize = 0;
|
||||
const fieldCount = try readU32(data, &pos);
|
||||
var i: u32 = 0;
|
||||
while (i < fieldCount) : (i += 1) {
|
||||
const nameLen = try readU16(data, &pos);
|
||||
if (pos + nameLen > data.len) return error.Malformed;
|
||||
const name = data[pos .. pos + nameLen];
|
||||
pos += nameLen;
|
||||
if (pos >= data.len) return error.Malformed;
|
||||
const tag: ValueTag = std.meta.intToEnum(ValueTag, data[pos]) catch return error.Malformed;
|
||||
pos += 1;
|
||||
const value: Value = switch (tag) {
|
||||
.none => .{ .none = {} },
|
||||
.boolean => blk: {
|
||||
if (pos >= data.len) return error.Malformed;
|
||||
const b = data[pos];
|
||||
pos += 1;
|
||||
break :blk .{ .boolean = b != 0 };
|
||||
},
|
||||
.int => .{ .int = @bitCast(try readU64(data, &pos)) },
|
||||
.float => .{ .float = @bitCast(try readU64(data, &pos)) },
|
||||
.text => blk: {
|
||||
const len = try readU32(data, &pos);
|
||||
if (pos + len > data.len) return error.Malformed;
|
||||
const s = data[pos .. pos + len];
|
||||
pos += len;
|
||||
break :blk .{ .text = s };
|
||||
},
|
||||
.bytes => blk: {
|
||||
const len = try readU32(data, &pos);
|
||||
if (pos + len > data.len) return error.Malformed;
|
||||
const s = data[pos .. pos + len];
|
||||
pos += len;
|
||||
break :blk .{ .bytes = s };
|
||||
},
|
||||
};
|
||||
try row.set(name, value);
|
||||
}
|
||||
if (pos != data.len) return error.Malformed;
|
||||
return row;
|
||||
}
|
||||
};
|
||||
|
||||
fn putU16(out: *std.ArrayList(u8), v: u16) !void {
|
||||
var b: [2]u8 = undefined;
|
||||
std.mem.writeInt(u16, &b, v, .little);
|
||||
try out.appendSlice(&b);
|
||||
}
|
||||
|
||||
fn putU32(out: *std.ArrayList(u8), v: u32) !void {
|
||||
var b: [4]u8 = undefined;
|
||||
std.mem.writeInt(u32, &b, v, .little);
|
||||
try out.appendSlice(&b);
|
||||
}
|
||||
|
||||
fn putU64(out: *std.ArrayList(u8), v: u64) !void {
|
||||
var b: [8]u8 = undefined;
|
||||
std.mem.writeInt(u64, &b, v, .little);
|
||||
try out.appendSlice(&b);
|
||||
}
|
||||
|
||||
fn readU16(data: []const u8, pos: *usize) !u16 {
|
||||
if (pos.* + 2 > data.len) return error.Malformed;
|
||||
const v = std.mem.readInt(u16, data[pos.*..][0..2], .little);
|
||||
pos.* += 2;
|
||||
return v;
|
||||
}
|
||||
|
||||
fn readU32(data: []const u8, pos: *usize) !u32 {
|
||||
if (pos.* + 4 > data.len) return error.Malformed;
|
||||
const v = std.mem.readInt(u32, data[pos.*..][0..4], .little);
|
||||
pos.* += 4;
|
||||
return v;
|
||||
}
|
||||
|
||||
fn readU64(data: []const u8, pos: *usize) !u64 {
|
||||
if (pos.* + 8 > data.len) return error.Malformed;
|
||||
const v = std.mem.readInt(u64, data[pos.*..][0..8], .little);
|
||||
pos.* += 8;
|
||||
return v;
|
||||
}
|
||||
|
||||
/// Appends an order-preserving byte encoding of `value` to `out`, such that
|
||||
/// lexicographic byte comparison of two encodings matches `compare()` of the
|
||||
/// original values. Used to build index keys and primary-key suffixes that
|
||||
/// must sort correctly inside the LSM tree's byte-ordered keyspace.
|
||||
///
|
||||
/// Text/bytes payloads are escaped (0x00 -> 0x00 0xFF) and terminated with
|
||||
/// 0x00 0x00 so that concatenating several encoded values back to back stays
|
||||
/// unambiguous and order-preserving (no payload can be a byte-for-byte prefix
|
||||
/// of another once escaped and terminated).
|
||||
pub fn appendOrdered(out: *std.ArrayList(u8), value: Value) !void {
|
||||
try out.append(@intFromEnum(@as(ValueTag, value)));
|
||||
switch (value) {
|
||||
.none => {},
|
||||
.boolean => |v| try out.append(if (v) 1 else 0),
|
||||
.int => |v| {
|
||||
const flipped: u64 = @as(u64, @bitCast(v)) ^ 0x8000000000000000;
|
||||
var b: [8]u8 = undefined;
|
||||
std.mem.writeInt(u64, &b, flipped, .big);
|
||||
try out.appendSlice(&b);
|
||||
},
|
||||
.float => |v| {
|
||||
var bits: u64 = @bitCast(v);
|
||||
if (bits & 0x8000000000000000 != 0) {
|
||||
bits = ~bits;
|
||||
} else {
|
||||
bits |= 0x8000000000000000;
|
||||
}
|
||||
var b: [8]u8 = undefined;
|
||||
std.mem.writeInt(u64, &b, bits, .big);
|
||||
try out.appendSlice(&b);
|
||||
},
|
||||
.text, .bytes => |s| {
|
||||
for (s) |byte| {
|
||||
if (byte == 0x00) {
|
||||
try out.appendSlice(&.{ 0x00, 0xFF });
|
||||
} else {
|
||||
try out.append(byte);
|
||||
}
|
||||
}
|
||||
try out.appendSlice(&.{ 0x00, 0x00 });
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub fn encodeOrdered(allocator: Allocator, value: Value) ![]u8 {
|
||||
var out = std.ArrayList(u8).init(allocator);
|
||||
errdefer out.deinit();
|
||||
try appendOrdered(&out, value);
|
||||
return out.toOwnedSlice();
|
||||
}
|
||||
|
||||
/// Encodes a single standalone Value using the same (non order-preserving)
|
||||
/// tag+payload format as a Row field, minus the field name. Used to store
|
||||
/// primary keys as index-entry values, where later exact reconstruction
|
||||
/// matters more than byte ordering.
|
||||
pub fn encodeValue(allocator: Allocator, value: Value) ![]u8 {
|
||||
var out = std.ArrayList(u8).init(allocator);
|
||||
errdefer out.deinit();
|
||||
try out.append(@intFromEnum(@as(ValueTag, value)));
|
||||
switch (value) {
|
||||
.none => {},
|
||||
.boolean => |v| try out.append(if (v) 1 else 0),
|
||||
.int => |v| try putU64(&out, @bitCast(v)),
|
||||
.float => |v| try putU64(&out, @bitCast(v)),
|
||||
.text, .bytes => |s| {
|
||||
try putU32(&out, @intCast(s.len));
|
||||
try out.appendSlice(s);
|
||||
},
|
||||
}
|
||||
return out.toOwnedSlice();
|
||||
}
|
||||
|
||||
pub fn decodeValue(allocator: Allocator, data: []const u8) !Value {
|
||||
var pos: usize = 0;
|
||||
if (pos >= data.len) return error.Malformed;
|
||||
const tag: ValueTag = std.meta.intToEnum(ValueTag, data[pos]) catch return error.Malformed;
|
||||
pos += 1;
|
||||
return switch (tag) {
|
||||
.none => .{ .none = {} },
|
||||
.boolean => blk: {
|
||||
if (pos >= data.len) return error.Malformed;
|
||||
break :blk .{ .boolean = data[pos] != 0 };
|
||||
},
|
||||
.int => .{ .int = @bitCast(try readU64(data, &pos)) },
|
||||
.float => .{ .float = @bitCast(try readU64(data, &pos)) },
|
||||
.text => blk: {
|
||||
const len = try readU32(data, &pos);
|
||||
if (pos + len > data.len) return error.Malformed;
|
||||
break :blk .{ .text = try allocator.dupe(u8, data[pos .. pos + len]) };
|
||||
},
|
||||
.bytes => blk: {
|
||||
const len = try readU32(data, &pos);
|
||||
if (pos + len > data.len) return error.Malformed;
|
||||
break :blk .{ .bytes = try allocator.dupe(u8, data[pos .. pos + len]) };
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
test "row set get overwrite and free" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = Row.init(allocator);
|
||||
defer row.deinit();
|
||||
|
||||
try row.set("name", .{ .text = "ada" });
|
||||
try row.set("age", .{ .int = 30 });
|
||||
try row.set("name", .{ .text = "ada lovelace" });
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 2), row.count());
|
||||
try std.testing.expectEqualStrings("ada lovelace", row.get("name").?.text);
|
||||
try std.testing.expectEqual(@as(i64, 30), row.get("age").?.int);
|
||||
try std.testing.expect(row.get("missing") == null);
|
||||
}
|
||||
|
||||
test "row encode decode round trip all types" {
|
||||
const allocator = std.testing.allocator;
|
||||
var row = Row.init(allocator);
|
||||
defer row.deinit();
|
||||
|
||||
try row.set("id", .{ .int = -42 });
|
||||
try row.set("score", .{ .float = 3.5 });
|
||||
try row.set("active", .{ .boolean = true });
|
||||
try row.set("bio", .{ .text = "hello world" });
|
||||
try row.set("blob", .{ .bytes = &.{ 0, 1, 2, 255 } });
|
||||
try row.set("nothing", .{ .none = {} });
|
||||
|
||||
const encoded = try row.encode(allocator);
|
||||
defer allocator.free(encoded);
|
||||
|
||||
var decoded = try Row.decode(allocator, encoded);
|
||||
defer decoded.deinit();
|
||||
|
||||
try std.testing.expectEqual(@as(i64, -42), decoded.get("id").?.int);
|
||||
try std.testing.expectEqual(@as(f64, 3.5), decoded.get("score").?.float);
|
||||
try std.testing.expectEqual(true, decoded.get("active").?.boolean);
|
||||
try std.testing.expectEqualStrings("hello world", decoded.get("bio").?.text);
|
||||
try std.testing.expectEqualSlices(u8, &.{ 0, 1, 2, 255 }, decoded.get("blob").?.bytes);
|
||||
try std.testing.expectEqual(ValueTag.none, @as(ValueTag, decoded.get("nothing").?));
|
||||
}
|
||||
|
||||
test "row decode rejects malformed data" {
|
||||
const allocator = std.testing.allocator;
|
||||
try std.testing.expectError(error.Malformed, Row.decode(allocator, &.{ 1, 0, 0, 0 }));
|
||||
try std.testing.expectError(error.Malformed, Row.decode(allocator, &.{ 1, 0, 0, 0, 5, 0, 'a' }));
|
||||
}
|
||||
|
||||
test "compare orders by type tag then payload" {
|
||||
try std.testing.expectEqual(std.math.Order.lt, compare(.{ .none = {} }, .{ .boolean = false }));
|
||||
try std.testing.expectEqual(std.math.Order.lt, compare(.{ .int = 5 }, .{ .float = 0.0 }));
|
||||
try std.testing.expectEqual(std.math.Order.lt, compare(.{ .int = -5 }, .{ .int = 3 }));
|
||||
try std.testing.expectEqual(std.math.Order.gt, compare(.{ .text = "b" }, .{ .text = "a" }));
|
||||
try std.testing.expect(Value.eql(.{ .text = "x" }, .{ .text = "x" }));
|
||||
}
|
||||
|
||||
test "appendOrdered preserves int order across sign boundary" {
|
||||
const allocator = std.testing.allocator;
|
||||
const values = [_]i64{ std.math.minInt(i64), -1000, -1, 0, 1, 1000, std.math.maxInt(i64) };
|
||||
var prev: ?[]u8 = null;
|
||||
defer if (prev) |p| allocator.free(p);
|
||||
for (values) |v| {
|
||||
const encoded = try encodeOrdered(allocator, .{ .int = v });
|
||||
if (prev) |p| {
|
||||
try std.testing.expect(std.mem.order(u8, p, encoded) == .lt);
|
||||
allocator.free(p);
|
||||
}
|
||||
prev = encoded;
|
||||
}
|
||||
}
|
||||
|
||||
test "appendOrdered preserves float order across sign and magnitude" {
|
||||
const allocator = std.testing.allocator;
|
||||
const values = [_]f64{ -100.0, -1.5, -0.001, 0.0, 0.001, 1.5, 100.0 };
|
||||
var prev: ?[]u8 = null;
|
||||
defer if (prev) |p| allocator.free(p);
|
||||
for (values) |v| {
|
||||
const encoded = try encodeOrdered(allocator, .{ .float = v });
|
||||
if (prev) |p| {
|
||||
try std.testing.expect(std.mem.order(u8, p, encoded) == .lt);
|
||||
allocator.free(p);
|
||||
}
|
||||
prev = encoded;
|
||||
}
|
||||
}
|
||||
|
||||
test "appendOrdered escapes embedded zero bytes unambiguously" {
|
||||
const allocator = std.testing.allocator;
|
||||
const a = try encodeOrdered(allocator, .{ .text = "ab" });
|
||||
defer allocator.free(a);
|
||||
const b = try encodeOrdered(allocator, .{ .text = "ab\x00c" });
|
||||
defer allocator.free(b);
|
||||
const c = try encodeOrdered(allocator, .{ .text = "ab" ++ "c"[0..0] });
|
||||
defer allocator.free(c);
|
||||
// "ab" must sort before "ab\x00c" (shorter, unescaped terminator wins),
|
||||
// and the escaped embedded zero must not corrupt ordering.
|
||||
try std.testing.expect(std.mem.order(u8, a, b) == .lt);
|
||||
try std.testing.expect(std.mem.eql(u8, a, c));
|
||||
}
|
||||
|
||||
test "encodeValue decodeValue round trip" {
|
||||
const allocator = std.testing.allocator;
|
||||
const original = Value{ .text = "primary-key-42" };
|
||||
const encoded = try encodeValue(allocator, original);
|
||||
defer allocator.free(encoded);
|
||||
const decoded = try decodeValue(allocator, encoded);
|
||||
defer decoded.free(allocator);
|
||||
try std.testing.expectEqualStrings("primary-key-42", decoded.text);
|
||||
|
||||
const intEncoded = try encodeValue(allocator, .{ .int = -7 });
|
||||
defer allocator.free(intEncoded);
|
||||
const intDecoded = try decodeValue(allocator, intEncoded);
|
||||
try std.testing.expectEqual(@as(i64, -7), intDecoded.int);
|
||||
}
|
||||
|
||||
test "appendOrdered text ordering matches string ordering" {
|
||||
const allocator = std.testing.allocator;
|
||||
const words = [_][]const u8{ "apple", "banana", "banjo", "zebra" };
|
||||
var prev: ?[]u8 = null;
|
||||
defer if (prev) |p| allocator.free(p);
|
||||
for (words) |w| {
|
||||
const encoded = try encodeOrdered(allocator, .{ .text = w });
|
||||
if (prev) |p| {
|
||||
try std.testing.expect(std.mem.order(u8, p, encoded) == .lt);
|
||||
allocator.free(p);
|
||||
}
|
||||
prev = encoded;
|
||||
}
|
||||
}
|
||||
|
||||
test "parseCli honors type prefixes and defaults to text" {
|
||||
try std.testing.expectEqual(@as(i64, 42), (try Value.parseCli("i:42")).int);
|
||||
try std.testing.expectEqual(@as(i64, -3), (try Value.parseCli("i:-3")).int);
|
||||
try std.testing.expectEqual(@as(f64, 1.5), (try Value.parseCli("f:1.5")).float);
|
||||
try std.testing.expectEqual(true, (try Value.parseCli("b:true")).boolean);
|
||||
try std.testing.expectEqual(false, (try Value.parseCli("b:false")).boolean);
|
||||
try std.testing.expectEqualStrings("hi", (try Value.parseCli("s:hi")).text);
|
||||
try std.testing.expectEqualStrings("plain", (try Value.parseCli("plain")).text);
|
||||
}
|
||||
|
||||
test "parseCli rejects malformed prefixed values" {
|
||||
try std.testing.expectError(error.InvalidValue, Value.parseCli("i:notanumber"));
|
||||
try std.testing.expectError(error.InvalidValue, Value.parseCli("f:notafloat"));
|
||||
try std.testing.expectError(error.InvalidValue, Value.parseCli("b:maybe"));
|
||||
}
|
||||
+8145
File diff suppressed because it is too large.
Load diff
+213
@@ -0,0 +1,213 @@
|
||||
const std = @import("std");
|
||||
const authmod = @import("auth.zig");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
|
||||
/// Random session tokens with expiration, independent of any one TCP
|
||||
/// connection - the layer that replaces "authenticated for as long as this
|
||||
/// socket stays open" (server.zig's `connRole`) with "authenticated until
|
||||
/// this token expires or is revoked". A client that authenticates once,
|
||||
/// gets a token back, and disconnects can reconnect later and resume the
|
||||
/// same role by presenting the token (`auth_token` opcode in protocol.zig)
|
||||
/// instead of the password again - and the token stops working on its own
|
||||
/// after its TTL elapses even if the same connection stays open the whole
|
||||
/// time, unlike plain password auth which never expires mid-connection.
|
||||
pub const tokenBytes = 24;
|
||||
pub const tokenHexLen = tokenBytes * 2;
|
||||
|
||||
fn generateTokenHex() [tokenHexLen]u8 {
|
||||
var raw: [tokenBytes]u8 = undefined;
|
||||
std.crypto.random.bytes(&raw);
|
||||
return std.fmt.bytesToHex(raw, .lower);
|
||||
}
|
||||
|
||||
pub const Session = struct {
|
||||
role: authmod.Role,
|
||||
expiresAtNanos: i128,
|
||||
};
|
||||
|
||||
/// In-memory only (never persisted to Db - a session is a login artifact,
|
||||
/// not application data): thread-safe via its own mutex since `server.zig`
|
||||
/// dispatches requests from one worker thread per connection and any of
|
||||
/// them may issue, validate, or revoke a token concurrently.
|
||||
pub const SessionStore = struct {
|
||||
allocator: Allocator,
|
||||
mutex: std.Thread.Mutex = .{},
|
||||
sessions: std.StringHashMapUnmanaged(Session) = .{},
|
||||
|
||||
pub fn init(allocator: Allocator) SessionStore {
|
||||
return .{ .allocator = allocator };
|
||||
}
|
||||
|
||||
pub fn deinit(self: *SessionStore) void {
|
||||
var it = self.sessions.iterator();
|
||||
while (it.next()) |entry| self.allocator.free(entry.key_ptr.*);
|
||||
self.sessions.deinit(self.allocator);
|
||||
}
|
||||
|
||||
/// Issues a fresh random token bound to `role`, valid until
|
||||
/// `nowNanos + ttlNanos`. Caller owns the returned slice.
|
||||
pub fn issue(self: *SessionStore, role: authmod.Role, ttlNanos: i128, nowNanos: i128) ![]u8 {
|
||||
const hex = generateTokenHex();
|
||||
const keyOwned = try self.allocator.dupe(u8, &hex);
|
||||
errdefer self.allocator.free(keyOwned);
|
||||
const returned = try self.allocator.dupe(u8, &hex);
|
||||
errdefer self.allocator.free(returned);
|
||||
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
try self.sessions.put(self.allocator, keyOwned, .{ .role = role, .expiresAtNanos = nowNanos + ttlNanos });
|
||||
return returned;
|
||||
}
|
||||
|
||||
/// Returns the session bound to `token` (role and expiry) if it exists
|
||||
/// and has not expired as of `nowNanos`. An expired session is deleted
|
||||
/// right here (lazy cleanup) rather than left around for
|
||||
/// `sweepExpired` to eventually find - a token that has expired must
|
||||
/// never again validate as live, so the deletion happens on the same
|
||||
/// lookup that discovers it. Returning the expiry (not just the role)
|
||||
/// lets a caller such as `server.zig` re-check the same session later
|
||||
/// against a fresh clock without a second store lookup.
|
||||
pub fn validate(self: *SessionStore, token: []const u8, nowNanos: i128) ?Session {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
const entry = self.sessions.getEntry(token) orelse return null;
|
||||
if (entry.value_ptr.expiresAtNanos <= nowNanos) {
|
||||
const key = entry.key_ptr.*;
|
||||
_ = self.sessions.remove(token);
|
||||
self.allocator.free(key);
|
||||
return null;
|
||||
}
|
||||
return entry.value_ptr.*;
|
||||
}
|
||||
|
||||
/// Explicit logout: removes a token before it would naturally expire.
|
||||
/// Returns whether a session was actually present.
|
||||
pub fn revoke(self: *SessionStore, token: []const u8) bool {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
const kv = self.sessions.fetchRemove(token) orelse return false;
|
||||
self.allocator.free(kv.key);
|
||||
return true;
|
||||
}
|
||||
|
||||
/// Removes every session expired as of `nowNanos`, returning how many
|
||||
/// were dropped. `validate` already deletes an expired session the
|
||||
/// moment someone tries to use it, but a token nobody ever presents
|
||||
/// again would otherwise sit in memory forever - this is the
|
||||
/// housekeeping sweep a long-lived server can run periodically.
|
||||
pub fn sweepExpired(self: *SessionStore, nowNanos: i128) usize {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
|
||||
var doomed = std.ArrayList([]const u8).init(self.allocator);
|
||||
defer doomed.deinit();
|
||||
var it = self.sessions.iterator();
|
||||
while (it.next()) |entry| {
|
||||
if (entry.value_ptr.expiresAtNanos <= nowNanos) {
|
||||
doomed.append(entry.key_ptr.*) catch continue;
|
||||
}
|
||||
}
|
||||
for (doomed.items) |key| {
|
||||
_ = self.sessions.remove(key);
|
||||
self.allocator.free(key);
|
||||
}
|
||||
return doomed.items.len;
|
||||
}
|
||||
|
||||
pub fn count(self: *SessionStore) usize {
|
||||
self.mutex.lock();
|
||||
defer self.mutex.unlock();
|
||||
return self.sessions.count();
|
||||
}
|
||||
};
|
||||
|
||||
test "issue then validate returns the bound role before expiry" {
|
||||
var store = SessionStore.init(std.testing.allocator);
|
||||
defer store.deinit();
|
||||
|
||||
const token = try store.issue(.readwrite, 1_000_000_000, 0);
|
||||
defer std.testing.allocator.free(token);
|
||||
|
||||
try std.testing.expectEqual(@as(usize, tokenHexLen), token.len);
|
||||
try std.testing.expectEqual(authmod.Role.readwrite, store.validate(token, 500_000_000).?.role);
|
||||
try std.testing.expectEqual(@as(usize, 1), store.count());
|
||||
}
|
||||
|
||||
test "validate returns null once past expiry and removes the session" {
|
||||
var store = SessionStore.init(std.testing.allocator);
|
||||
defer store.deinit();
|
||||
|
||||
const token = try store.issue(.admin, 1_000, 0);
|
||||
defer std.testing.allocator.free(token);
|
||||
|
||||
try std.testing.expect(store.validate(token, 1_000) == null);
|
||||
try std.testing.expectEqual(@as(usize, 0), store.count());
|
||||
// A second lookup of the same (already-removed) token still misses cleanly.
|
||||
try std.testing.expect(store.validate(token, 1_000) == null);
|
||||
}
|
||||
|
||||
test "unknown token never validates" {
|
||||
var store = SessionStore.init(std.testing.allocator);
|
||||
defer store.deinit();
|
||||
try std.testing.expect(store.validate("does-not-exist", 0) == null);
|
||||
}
|
||||
|
||||
test "revoke removes a session before its natural expiry" {
|
||||
var store = SessionStore.init(std.testing.allocator);
|
||||
defer store.deinit();
|
||||
|
||||
const token = try store.issue(.readonly, 1_000_000_000, 0);
|
||||
defer std.testing.allocator.free(token);
|
||||
|
||||
try std.testing.expect(store.revoke(token));
|
||||
try std.testing.expect(store.validate(token, 0) == null);
|
||||
// Revoking an already-gone token is a clean no-op, not an error.
|
||||
try std.testing.expect(!store.revoke(token));
|
||||
}
|
||||
|
||||
test "two tokens for different roles are independent" {
|
||||
var store = SessionStore.init(std.testing.allocator);
|
||||
defer store.deinit();
|
||||
|
||||
const a = try store.issue(.readonly, 1_000_000_000, 0);
|
||||
defer std.testing.allocator.free(a);
|
||||
const b = try store.issue(.admin, 1_000_000_000, 0);
|
||||
defer std.testing.allocator.free(b);
|
||||
|
||||
try std.testing.expect(!std.mem.eql(u8, a, b));
|
||||
try std.testing.expectEqual(authmod.Role.readonly, store.validate(a, 0).?.role);
|
||||
try std.testing.expectEqual(authmod.Role.admin, store.validate(b, 0).?.role);
|
||||
|
||||
try std.testing.expect(store.revoke(a));
|
||||
try std.testing.expect(store.validate(a, 0) == null);
|
||||
// Revoking `a` must never touch `b`'s still-live session.
|
||||
try std.testing.expectEqual(authmod.Role.admin, store.validate(b, 0).?.role);
|
||||
}
|
||||
|
||||
test "sweepExpired drops only expired sessions and reports the count" {
|
||||
var store = SessionStore.init(std.testing.allocator);
|
||||
defer store.deinit();
|
||||
|
||||
const stale1 = try store.issue(.readonly, 100, 0);
|
||||
defer std.testing.allocator.free(stale1);
|
||||
const stale2 = try store.issue(.readwrite, 100, 0);
|
||||
defer std.testing.allocator.free(stale2);
|
||||
const fresh = try store.issue(.admin, 1_000_000_000, 0);
|
||||
defer std.testing.allocator.free(fresh);
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 3), store.count());
|
||||
try std.testing.expectEqual(@as(usize, 2), store.sweepExpired(200));
|
||||
try std.testing.expectEqual(@as(usize, 1), store.count());
|
||||
try std.testing.expectEqual(authmod.Role.admin, store.validate(fresh, 200).?.role);
|
||||
|
||||
// A second sweep at the same time has nothing left to remove.
|
||||
try std.testing.expectEqual(@as(usize, 0), store.sweepExpired(200));
|
||||
}
|
||||
|
||||
test "generateTokenHex produces distinct, fully hex-encoded tokens" {
|
||||
const a = generateTokenHex();
|
||||
const b = generateTokenHex();
|
||||
try std.testing.expect(!std.mem.eql(u8, &a, &b));
|
||||
for (a) |c| try std.testing.expect(std.ascii.isHex(c));
|
||||
}
|
||||
+471
-9
@@ -24,11 +24,26 @@ fn order(a: []const u8, b: []const u8) std.math.Order {
|
||||
}
|
||||
|
||||
const magic: u32 = 0xC0BA1701;
|
||||
const version: u32 = 1;
|
||||
const version: u32 = 2;
|
||||
const footerLen = 4 + 4 + 8 + 4 + 8 + 4 + 8 + 8;
|
||||
const recordHeader = 1 + 8 + 4 + 4;
|
||||
const defaultBlockTarget = 4096;
|
||||
|
||||
/// Every Nth record in a data block is a "restart point": stored with its
|
||||
/// full key (shared == 0) instead of a shared-prefix delta against the
|
||||
/// previous key. This bounds how many deltas must be re-applied to recover
|
||||
/// any one key (never more than `restartInterval - 1`) and, more
|
||||
/// importantly here, gives block expansion a cheap, structural corruption
|
||||
/// check: a real restart offset must always land on a shared == 0 record.
|
||||
const restartInterval: usize = 16;
|
||||
|
||||
fn commonPrefixLen(a: []const u8, b: []const u8) usize {
|
||||
const n = @min(a.len, b.len);
|
||||
var i: usize = 0;
|
||||
while (i < n and a[i] == b[i]) : (i += 1) {}
|
||||
return i;
|
||||
}
|
||||
|
||||
pub const Builder = struct {
|
||||
allocator: Allocator,
|
||||
file: std.fs.File,
|
||||
@@ -36,6 +51,9 @@ pub const Builder = struct {
|
||||
blockBuf: std.ArrayList(u8),
|
||||
blockFirstKey: std.ArrayList(u8),
|
||||
indexBuf: std.ArrayList(u8),
|
||||
lastKey: std.ArrayList(u8),
|
||||
blockRestarts: std.ArrayList(u32),
|
||||
blockRecordCount: usize,
|
||||
numEntries: u64,
|
||||
logicalBytes: u64,
|
||||
maxSeq: u64,
|
||||
@@ -51,6 +69,9 @@ pub const Builder = struct {
|
||||
.blockBuf = std.ArrayList(u8).init(allocator),
|
||||
.blockFirstKey = std.ArrayList(u8).init(allocator),
|
||||
.indexBuf = std.ArrayList(u8).init(allocator),
|
||||
.lastKey = std.ArrayList(u8).init(allocator),
|
||||
.blockRestarts = std.ArrayList(u32).init(allocator),
|
||||
.blockRecordCount = 0,
|
||||
.numEntries = 0,
|
||||
.logicalBytes = 0,
|
||||
.maxSeq = 0,
|
||||
@@ -63,6 +84,8 @@ pub const Builder = struct {
|
||||
self.blockBuf.deinit();
|
||||
self.blockFirstKey.deinit();
|
||||
self.indexBuf.deinit();
|
||||
self.lastKey.deinit();
|
||||
self.blockRestarts.deinit();
|
||||
self.bloom.deinit();
|
||||
self.file.close();
|
||||
}
|
||||
@@ -72,14 +95,25 @@ pub const Builder = struct {
|
||||
self.blockFirstKey.clearRetainingCapacity();
|
||||
try self.blockFirstKey.appendSlice(record.key);
|
||||
}
|
||||
|
||||
const isRestart = self.blockRecordCount % restartInterval == 0;
|
||||
const shared: u32 = if (isRestart) 0 else @intCast(commonPrefixLen(self.lastKey.items, record.key));
|
||||
const unshared = record.key[shared..];
|
||||
if (isRestart) try self.blockRestarts.append(@intCast(self.blockBuf.items.len));
|
||||
|
||||
const writer = encode.ByteWriter.init(&self.blockBuf);
|
||||
try writer.byte(@intFromEnum(record.kind));
|
||||
try writer.putU64(record.seq);
|
||||
try writer.putU32(@intCast(record.key.len));
|
||||
try writer.putU32(shared);
|
||||
try writer.putU32(@intCast(unshared.len));
|
||||
try writer.putU32(@intCast(record.value.len));
|
||||
try writer.slice(record.key);
|
||||
try writer.slice(unshared);
|
||||
try writer.slice(record.value);
|
||||
|
||||
self.lastKey.clearRetainingCapacity();
|
||||
try self.lastKey.appendSlice(record.key);
|
||||
self.blockRecordCount += 1;
|
||||
|
||||
self.bloom.add(record.key);
|
||||
self.numEntries += 1;
|
||||
self.logicalBytes += recordHeader + record.key.len + record.value.len;
|
||||
@@ -90,6 +124,15 @@ pub const Builder = struct {
|
||||
|
||||
fn flushBlock(self: *Builder) !void {
|
||||
if (self.blockBuf.items.len == 0) return;
|
||||
|
||||
// The restart-offset trailer rides inside the same logical block
|
||||
// bytes as the records themselves, so it survives compression and
|
||||
// is immediately available to `expandBlock` on the read side
|
||||
// without a separate on-disk section or footer entry.
|
||||
const trailerWriter = encode.ByteWriter.init(&self.blockBuf);
|
||||
for (self.blockRestarts.items) |restartOffset| try trailerWriter.putU32(restartOffset);
|
||||
try trailerWriter.putU32(@intCast(self.blockRestarts.items.len));
|
||||
|
||||
const originalLen: u32 = @intCast(self.blockBuf.items.len);
|
||||
|
||||
const encoded = try compress.encode(self.allocator, self.blockBuf.items);
|
||||
@@ -116,6 +159,9 @@ pub const Builder = struct {
|
||||
|
||||
self.offset += physLen + 4;
|
||||
self.blockBuf.clearRetainingCapacity();
|
||||
self.blockRestarts.clearRetainingCapacity();
|
||||
self.blockRecordCount = 0;
|
||||
self.lastKey.clearRetainingCapacity();
|
||||
}
|
||||
|
||||
pub fn finish(self: *Builder) !void {
|
||||
@@ -160,6 +206,13 @@ const IndexEntry = struct {
|
||||
offset: u64,
|
||||
blockLen: u32,
|
||||
firstKey: []const u8,
|
||||
// Offsets (relative to this block's own bytes, i.e. usable directly
|
||||
// against the slice `blockBytes()` returns) of every restart record in
|
||||
// the flattened block, in ascending order. restarts[0] is always 0.
|
||||
// Lets point lookups binary-search straight to the <= restartInterval
|
||||
// records that could hold the key instead of linear-scanning the whole
|
||||
// block from position 0.
|
||||
restarts: []u32,
|
||||
};
|
||||
|
||||
const blockPrefix = 1 + 4;
|
||||
@@ -188,6 +241,7 @@ pub const Reader = struct {
|
||||
if (buffer.len < footerLen) return error.CorruptSstable;
|
||||
const footerStart = buffer.len - footerLen;
|
||||
if (readU32(buffer, footerStart) != magic) return error.CorruptSstable;
|
||||
if (readU32(buffer, footerStart + 4) != version) return error.UnsupportedSstableVersion;
|
||||
const indexOffset = readU64(buffer, footerStart + 8);
|
||||
const indexLen = readU32(buffer, footerStart + 16);
|
||||
const bloomOffset = readU64(buffer, footerStart + 20);
|
||||
@@ -204,7 +258,10 @@ pub const Reader = struct {
|
||||
var decoded = std.ArrayList(u8).init(allocator);
|
||||
errdefer decoded.deinit();
|
||||
var entries = std.ArrayList(IndexEntry).init(allocator);
|
||||
errdefer entries.deinit();
|
||||
errdefer {
|
||||
for (entries.items) |e| allocator.free(e.restarts);
|
||||
entries.deinit();
|
||||
}
|
||||
|
||||
var pos: usize = 0;
|
||||
while (pos < indexData.len) {
|
||||
@@ -221,11 +278,19 @@ pub const Reader = struct {
|
||||
const payload = physical[blockPrefix..];
|
||||
const plainBlock = try compress.decode(allocator, codec, payload, originalLen);
|
||||
defer allocator.free(plainBlock);
|
||||
const expanded = try expandBlock(allocator, plainBlock);
|
||||
defer allocator.free(expanded.bytes);
|
||||
errdefer allocator.free(expanded.restarts);
|
||||
|
||||
const decOffset = decoded.items.len;
|
||||
try decoded.appendSlice(plainBlock);
|
||||
try decoded.appendSlice(expanded.bytes);
|
||||
|
||||
try entries.append(.{ .offset = decOffset, .blockLen = @intCast(plainBlock.len), .firstKey = firstKey });
|
||||
try entries.append(.{
|
||||
.offset = decOffset,
|
||||
.blockLen = @intCast(expanded.bytes.len),
|
||||
.firstKey = firstKey,
|
||||
.restarts = expanded.restarts,
|
||||
});
|
||||
pos = keyStart + firstKeyLen;
|
||||
}
|
||||
|
||||
@@ -265,6 +330,7 @@ pub const Reader = struct {
|
||||
}
|
||||
|
||||
pub fn deinit(self: *Reader) void {
|
||||
for (self.index) |e| self.allocator.free(e.restarts);
|
||||
self.allocator.free(self.index);
|
||||
self.bloom.deinit();
|
||||
self.allocator.free(self.dataBuf);
|
||||
@@ -279,7 +345,7 @@ pub const Reader = struct {
|
||||
if (!self.bloom.mayContain(key)) return null;
|
||||
const blockIdx = self.findBlock(key) orelse return null;
|
||||
const block = self.blockBytes(blockIdx);
|
||||
var pos: usize = 0;
|
||||
var pos: usize = Reader.seekInBlock(block, self.index[blockIdx].restarts, key);
|
||||
while (pos < block.len) {
|
||||
const decoded = decodeRecord(block, pos);
|
||||
switch (order(decoded.record.key, key)) {
|
||||
@@ -307,7 +373,7 @@ pub const Reader = struct {
|
||||
const block = cached orelse self.dataBuf[entry.offset .. entry.offset + entry.blockLen];
|
||||
if (cached == null) try blockCache.put(blockKey, block);
|
||||
|
||||
var pos: usize = 0;
|
||||
var pos: usize = Reader.seekInBlock(block, entry.restarts, key);
|
||||
while (pos < block.len) {
|
||||
const decoded = decodeRecord(block, pos);
|
||||
switch (order(decoded.record.key, key)) {
|
||||
@@ -329,7 +395,10 @@ pub const Reader = struct {
|
||||
pub fn cursor(self: *const Reader, start: ?[]const u8) Cursor {
|
||||
var c = Cursor{ .reader = self, .blockIdx = 0, .pos = 0, .nextPos = 0, .cur = null };
|
||||
if (self.index.len == 0) return c;
|
||||
c.blockIdx = if (start) |s| (self.findBlock(s) orelse 0) else 0;
|
||||
if (start) |s| {
|
||||
c.blockIdx = self.findBlock(s) orelse 0;
|
||||
c.pos = Reader.seekInBlock(self.blockBytes(c.blockIdx), self.index[c.blockIdx].restarts, s);
|
||||
}
|
||||
c.refresh();
|
||||
if (start) |s| {
|
||||
while (c.valid() and order(c.cur.?.key, s) == .lt) c.advance();
|
||||
@@ -400,6 +469,30 @@ pub const Reader = struct {
|
||||
return ans;
|
||||
}
|
||||
|
||||
/// Binary-searches `restarts` (each entry the flat-block offset of a
|
||||
/// restart record, ascending by key since keys are sorted within a
|
||||
/// block) for the rightmost restart whose key is <= `key`, and returns
|
||||
/// its offset. Callers then linear-scan forward from there instead of
|
||||
/// from position 0, bounding the scan to at most `restartInterval`
|
||||
/// records regardless of how many records the block holds. Returns 0
|
||||
/// (restarts[0], always the block's first record) when `key` sorts
|
||||
/// before every restart.
|
||||
fn seekInBlock(block: []const u8, restarts: []const u32, key: []const u8) usize {
|
||||
if (restarts.len == 0) return 0;
|
||||
var lo: usize = 0;
|
||||
var hi: usize = restarts.len;
|
||||
while (lo < hi) {
|
||||
const mid = lo + (hi - lo) / 2;
|
||||
const rec = decodeRecord(block, restarts[mid]);
|
||||
if (order(rec.record.key, key) != .gt) {
|
||||
lo = mid + 1;
|
||||
} else {
|
||||
hi = mid;
|
||||
}
|
||||
}
|
||||
return if (lo == 0) 0 else restarts[lo - 1];
|
||||
}
|
||||
|
||||
fn blockBytes(self: *const Reader, idx: usize) []const u8 {
|
||||
const entry = self.index[idx];
|
||||
return self.dataBuf[entry.offset .. entry.offset + entry.blockLen];
|
||||
@@ -519,6 +612,113 @@ const Decoded = struct {
|
||||
next: usize,
|
||||
};
|
||||
|
||||
const ExpandedBlock = struct {
|
||||
bytes: []u8,
|
||||
// Flat-block offsets (into `bytes`) of every restart record, ascending.
|
||||
// Consumed by `Reader.seekInBlock` for binary-search point lookups.
|
||||
restarts: []u32,
|
||||
};
|
||||
|
||||
/// Reconstructs a physical block (prefix-compressed records plus its
|
||||
/// restart-offset trailer, as written by `Builder`) back into the flat
|
||||
/// every-record-carries-its-full-key layout that `decodeRecord` and every
|
||||
/// cursor in this file expect. Expansion happens once, right after
|
||||
/// decompression in `Reader.open`, so `dataBuf`/the block cache keep
|
||||
/// storing fully materialized blocks and the rest of the reader never has
|
||||
/// to know keys were delta-encoded on disk - in particular this preserves
|
||||
/// the invariant `db.zig`'s k-way merge relies on, that a cursor's current
|
||||
/// key slice stays valid until that same cursor is advanced again.
|
||||
///
|
||||
/// Also validates the restart trailer against the record stream: every
|
||||
/// declared restart offset must land exactly on a record whose `shared`
|
||||
/// field is 0, and every record with `shared` 0 at that exact interval
|
||||
/// must have a matching restart entry. This catches a truncated or
|
||||
/// shifted restart table that a whole-block CRC would not (the CRC only
|
||||
/// proves the compressed bytes were not corrupted in transit/on disk, not
|
||||
/// that this structural relationship inside them still holds).
|
||||
///
|
||||
/// Also re-derives, in the *flattened* block's own coordinates, the offset
|
||||
/// of each restart record (`ExpandedBlock.restarts`) so the read side can
|
||||
/// still binary-search to a restart-sized (<= restartInterval records)
|
||||
/// linear-scan window even though the on-disk shared/unshared deltas are
|
||||
/// gone post-expansion.
|
||||
fn expandBlock(allocator: Allocator, physical: []const u8) !ExpandedBlock {
|
||||
if (physical.len < 4) return error.CorruptSstable;
|
||||
const numRestarts = readU32(physical, physical.len - 4);
|
||||
const restartTableBytes = @as(usize, numRestarts) * 4;
|
||||
if (restartTableBytes + 4 > physical.len) return error.CorruptSstable;
|
||||
const recordsEnd = physical.len - 4 - restartTableBytes;
|
||||
const restartTable = physical[recordsEnd .. recordsEnd + restartTableBytes];
|
||||
|
||||
var out = std.ArrayList(u8).init(allocator);
|
||||
errdefer out.deinit();
|
||||
|
||||
var flatRestarts = std.ArrayList(u32).init(allocator);
|
||||
errdefer flatRestarts.deinit();
|
||||
|
||||
var prevKey = std.ArrayList(u8).init(allocator);
|
||||
defer prevKey.deinit();
|
||||
|
||||
const perRecordFixed = 1 + 8 + 4 + 4 + 4;
|
||||
var pos: usize = 0;
|
||||
var restartIdx: usize = 0;
|
||||
var recordIdx: usize = 0;
|
||||
while (pos < recordsEnd) {
|
||||
if (pos + perRecordFixed > recordsEnd) return error.CorruptSstable;
|
||||
const kind = physical[pos];
|
||||
const seq = readU64(physical, pos + 1);
|
||||
const shared = readU32(physical, pos + 9);
|
||||
const unshared = readU32(physical, pos + 13);
|
||||
const valLen = readU32(physical, pos + 17);
|
||||
const keyStart = pos + perRecordFixed;
|
||||
if (shared > prevKey.items.len) return error.CorruptSstable;
|
||||
if (keyStart + @as(usize, unshared) + @as(usize, valLen) > recordsEnd) return error.CorruptSstable;
|
||||
|
||||
// A genuine block (as `Builder` writes it) marks exactly every
|
||||
// `restartInterval`-th record (by position within the block) as a
|
||||
// restart, in order, with shared == 0. Checking both directions -
|
||||
// a restart-interval position must be declared, and a declared
|
||||
// offset must fall on a restart-interval position - catches a
|
||||
// restart table that is present but shifted, truncated, or
|
||||
// fabricated, not just one that is structurally impossible.
|
||||
const shouldBeRestart = recordIdx % restartInterval == 0;
|
||||
const isDeclaredRestart = restartIdx < numRestarts and readU32(restartTable, restartIdx * 4) == pos;
|
||||
if (shouldBeRestart != isDeclaredRestart) return error.CorruptSstable;
|
||||
if (shouldBeRestart) {
|
||||
if (shared != 0) return error.CorruptSstable;
|
||||
restartIdx += 1;
|
||||
try flatRestarts.append(@intCast(out.items.len));
|
||||
}
|
||||
|
||||
const unsharedBytes = physical[keyStart .. keyStart + unshared];
|
||||
const valStart = keyStart + unshared;
|
||||
const value = physical[valStart .. valStart + valLen];
|
||||
|
||||
const fullKey = try allocator.alloc(u8, shared + unshared);
|
||||
defer allocator.free(fullKey);
|
||||
@memcpy(fullKey[0..shared], prevKey.items[0..shared]);
|
||||
@memcpy(fullKey[shared..], unsharedBytes);
|
||||
|
||||
const writer = encode.ByteWriter.init(&out);
|
||||
try writer.byte(kind);
|
||||
try writer.putU64(seq);
|
||||
try writer.putU32(@intCast(fullKey.len));
|
||||
try writer.putU32(valLen);
|
||||
try writer.slice(fullKey);
|
||||
try writer.slice(value);
|
||||
|
||||
prevKey.clearRetainingCapacity();
|
||||
try prevKey.appendSlice(fullKey);
|
||||
|
||||
pos = valStart + valLen;
|
||||
recordIdx += 1;
|
||||
}
|
||||
if (pos != recordsEnd) return error.CorruptSstable;
|
||||
if (restartIdx != numRestarts) return error.CorruptSstable;
|
||||
|
||||
return .{ .bytes = try out.toOwnedSlice(), .restarts = try flatRestarts.toOwnedSlice() };
|
||||
}
|
||||
|
||||
fn decodeRecord(block: []const u8, pos: usize) Decoded {
|
||||
const kind: Kind = @enumFromInt(block[pos]);
|
||||
const seq = readU64(block, pos + 1);
|
||||
@@ -604,6 +804,73 @@ test "sstable cursor seek" {
|
||||
try std.testing.expect(!beyond.valid());
|
||||
}
|
||||
|
||||
test "restart index lets point lookups skip straight past earlier records in a block" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
try buildFixture(allocator, tmp.dir, "test.sst", 500);
|
||||
var reader = try Reader.open(allocator, tmp.dir, "test.sst", 1);
|
||||
defer reader.deinit();
|
||||
|
||||
// Every block must have carried its restart offsets through expansion:
|
||||
// the first is always the block's own first record, and multi-restart
|
||||
// blocks (any block with more than restartInterval records) must have
|
||||
// strictly increasing keys at each restart, matching overall key order.
|
||||
var sawMultiRestartBlock = false;
|
||||
var blockIdx: usize = 0;
|
||||
while (blockIdx < reader.index.len) : (blockIdx += 1) {
|
||||
const entry = reader.index[blockIdx];
|
||||
try std.testing.expect(entry.restarts.len >= 1);
|
||||
try std.testing.expectEqual(@as(u32, 0), entry.restarts[0]);
|
||||
if (entry.restarts.len > 1) sawMultiRestartBlock = true;
|
||||
|
||||
const block = reader.blockBytes(blockIdx);
|
||||
var prevKey: []const u8 = &.{};
|
||||
for (entry.restarts) |r| {
|
||||
const rec = decodeRecord(block, r);
|
||||
if (prevKey.len > 0) try std.testing.expect(order(prevKey, rec.record.key) == .lt);
|
||||
prevKey = rec.record.key;
|
||||
}
|
||||
}
|
||||
// 500 records at ~87/block (blockTarget 4096 / ~47 bytes/record) and
|
||||
// restartInterval 16 guarantees at least one block with several
|
||||
// restarts - if this ever goes false the fixture stopped exercising
|
||||
// the multi-restart path the rest of this test relies on.
|
||||
try std.testing.expect(sawMultiRestartBlock);
|
||||
|
||||
// A probe exactly equal to a later restart's key must land on that
|
||||
// restart's own offset (the binary search must not undershoot to an
|
||||
// earlier restart when there is an exact match).
|
||||
const entry0 = reader.index[0];
|
||||
try std.testing.expect(entry0.restarts.len > 2);
|
||||
const block0 = reader.blockBytes(0);
|
||||
const secondRestartOffset = entry0.restarts[1];
|
||||
const secondRestartKey = decodeRecord(block0, secondRestartOffset).record.key;
|
||||
try std.testing.expectEqual(secondRestartOffset, Reader.seekInBlock(block0, entry0.restarts, secondRestartKey));
|
||||
|
||||
// A probe strictly between the first and second restart's keys must
|
||||
// stay on the first restart (must not overshoot to the second).
|
||||
const firstRestartKey = decodeRecord(block0, entry0.restarts[0]).record.key;
|
||||
var midBuf: [16]u8 = undefined;
|
||||
@memcpy(midBuf[0..firstRestartKey.len], firstRestartKey);
|
||||
midBuf[firstRestartKey.len - 1] += 1; // one key past first restart's key
|
||||
try std.testing.expect(order(midBuf[0..firstRestartKey.len], secondRestartKey) == .lt);
|
||||
try std.testing.expectEqual(entry0.restarts[0], Reader.seekInBlock(block0, entry0.restarts, midBuf[0..firstRestartKey.len]));
|
||||
|
||||
// A probe past every restart's key must resolve to the last restart.
|
||||
const lastRestartKey = decodeRecord(block0, entry0.restarts[entry0.restarts.len - 1]).record.key;
|
||||
var pastBuf: [16]u8 = undefined;
|
||||
@memcpy(pastBuf[0..lastRestartKey.len], lastRestartKey);
|
||||
pastBuf[lastRestartKey.len - 1] += 1;
|
||||
try std.testing.expectEqual(entry0.restarts[entry0.restarts.len - 1], Reader.seekInBlock(block0, entry0.restarts, pastBuf[0..lastRestartKey.len]));
|
||||
|
||||
// End to end: lookups still resolve correctly for keys sitting right at
|
||||
// and around restart boundaries, now that the scan starts mid-block.
|
||||
try std.testing.expectEqualStrings("payload-value", reader.find(firstRestartKey).?.value);
|
||||
try std.testing.expectEqualStrings("payload-value", reader.find(secondRestartKey).?.value);
|
||||
}
|
||||
|
||||
test "sstable reverse cursor descends fully" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
@@ -648,3 +915,198 @@ test "sstable reverse cursor honors upper bound" {
|
||||
defer below.deinit();
|
||||
try std.testing.expect(!below.valid());
|
||||
}
|
||||
|
||||
fn writeRawRecord(w: encode.ByteWriter, kind: Kind, seq: u64, shared: u32, unshared: []const u8, value: []const u8) !void {
|
||||
try w.byte(@intFromEnum(kind));
|
||||
try w.putU64(seq);
|
||||
try w.putU32(shared);
|
||||
try w.putU32(@intCast(unshared.len));
|
||||
try w.putU32(@intCast(value.len));
|
||||
try w.slice(unshared);
|
||||
try w.slice(value);
|
||||
}
|
||||
|
||||
test "expandBlock reconstructs prefix-compressed records" {
|
||||
const allocator = std.testing.allocator;
|
||||
var buf = std.ArrayList(u8).init(allocator);
|
||||
defer buf.deinit();
|
||||
const w = encode.ByteWriter.init(&buf);
|
||||
|
||||
try writeRawRecord(w, .value, 1, 0, "apple", "1");
|
||||
try writeRawRecord(w, .value, 2, 3, "ly", "2"); // "app" + "ly" = "apply"
|
||||
try writeRawRecord(w, .tombstone, 3, 2, "ricot", ""); // "ap" + "ricot" = "apricot"
|
||||
try w.putU32(0); // one restart, at offset 0
|
||||
try w.putU32(1);
|
||||
|
||||
const expanded = try expandBlock(allocator, buf.items);
|
||||
defer allocator.free(expanded.bytes);
|
||||
defer allocator.free(expanded.restarts);
|
||||
|
||||
const d0 = decodeRecord(expanded.bytes, 0);
|
||||
try std.testing.expectEqualStrings("apple", d0.record.key);
|
||||
try std.testing.expectEqualStrings("1", d0.record.value);
|
||||
try std.testing.expectEqual(Kind.value, d0.record.kind);
|
||||
|
||||
const d1 = decodeRecord(expanded.bytes, d0.next);
|
||||
try std.testing.expectEqualStrings("apply", d1.record.key);
|
||||
try std.testing.expectEqualStrings("2", d1.record.value);
|
||||
|
||||
const d2 = decodeRecord(expanded.bytes, d1.next);
|
||||
try std.testing.expectEqualStrings("apricot", d2.record.key);
|
||||
try std.testing.expectEqual(Kind.tombstone, d2.record.kind);
|
||||
try std.testing.expectEqual(d2.next, expanded.bytes.len);
|
||||
|
||||
// Only the first record is a restart in this fixture (interval=16,
|
||||
// 3 records) - `flatRestarts` must record exactly its flat offset (0).
|
||||
try std.testing.expectEqualSlices(u32, &.{0}, expanded.restarts);
|
||||
}
|
||||
|
||||
test "expandBlock rejects a restart declared at the wrong record" {
|
||||
const allocator = std.testing.allocator;
|
||||
var buf = std.ArrayList(u8).init(allocator);
|
||||
defer buf.deinit();
|
||||
const w = encode.ByteWriter.init(&buf);
|
||||
|
||||
const record0End = 1 + 8 + 4 + 4 + 4 + 5 + 1; // header + "apple" + "1"
|
||||
try writeRawRecord(w, .value, 1, 0, "apple", "1");
|
||||
try writeRawRecord(w, .value, 2, 3, "ly", "2");
|
||||
try w.putU32(@intCast(record0End)); // restart claims record 1's offset, not record 0's
|
||||
try w.putU32(1);
|
||||
|
||||
try std.testing.expectError(error.CorruptSstable, expandBlock(allocator, buf.items));
|
||||
}
|
||||
|
||||
test "expandBlock rejects a missing restart declaration" {
|
||||
const allocator = std.testing.allocator;
|
||||
var buf = std.ArrayList(u8).init(allocator);
|
||||
defer buf.deinit();
|
||||
const w = encode.ByteWriter.init(&buf);
|
||||
|
||||
try writeRawRecord(w, .value, 1, 0, "apple", "1");
|
||||
try w.putU32(0); // trailer present but claims zero restarts
|
||||
|
||||
try std.testing.expectError(error.CorruptSstable, expandBlock(allocator, buf.items));
|
||||
}
|
||||
|
||||
test "expandBlock rejects shared length exceeding the previous key" {
|
||||
const allocator = std.testing.allocator;
|
||||
var buf = std.ArrayList(u8).init(allocator);
|
||||
defer buf.deinit();
|
||||
const w = encode.ByteWriter.init(&buf);
|
||||
|
||||
try writeRawRecord(w, .value, 1, 0, "ab", "1"); // prevKey becomes "ab", len 2
|
||||
try writeRawRecord(w, .value, 2, 5, "xyz", "2"); // shared=5 > len("ab")
|
||||
try w.putU32(0);
|
||||
try w.putU32(1);
|
||||
|
||||
try std.testing.expectError(error.CorruptSstable, expandBlock(allocator, buf.items));
|
||||
}
|
||||
|
||||
test "expandBlock rejects a record stream that overruns its declared lengths" {
|
||||
const allocator = std.testing.allocator;
|
||||
var buf = std.ArrayList(u8).init(allocator);
|
||||
defer buf.deinit();
|
||||
const w = encode.ByteWriter.init(&buf);
|
||||
|
||||
try w.byte(@intFromEnum(Kind.value));
|
||||
try w.putU64(1);
|
||||
try w.putU32(0);
|
||||
try w.putU32(100); // claims 100 unshared key bytes that are never written
|
||||
try w.putU32(0);
|
||||
try w.putU32(0);
|
||||
try w.putU32(1);
|
||||
|
||||
try std.testing.expectError(error.CorruptSstable, expandBlock(allocator, buf.items));
|
||||
}
|
||||
|
||||
test "expandBlock rejects trailing bytes that do not form a full record" {
|
||||
const allocator = std.testing.allocator;
|
||||
var buf = std.ArrayList(u8).init(allocator);
|
||||
defer buf.deinit();
|
||||
const w = encode.ByteWriter.init(&buf);
|
||||
|
||||
try writeRawRecord(w, .value, 1, 0, "apple", "1");
|
||||
try w.byte(0xff); // stray byte before the restart trailer
|
||||
try w.putU32(0);
|
||||
try w.putU32(1);
|
||||
|
||||
try std.testing.expectError(error.CorruptSstable, expandBlock(allocator, buf.items));
|
||||
}
|
||||
|
||||
test "Reader.open rejects an sstable written by an unsupported format version" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
try buildFixture(allocator, tmp.dir, "test.sst", 10);
|
||||
|
||||
{
|
||||
var file = try tmp.dir.openFile("test.sst", .{ .mode = .read_write });
|
||||
defer file.close();
|
||||
const size = try file.getEndPos();
|
||||
// version u32 sits right after the magic u32 at the footer's start.
|
||||
try file.seekTo(size - footerLen + 4);
|
||||
var bogus: [4]u8 = undefined;
|
||||
std.mem.writeInt(u32, &bogus, version + 1, encode.endian);
|
||||
try file.writeAll(&bogus);
|
||||
}
|
||||
|
||||
try std.testing.expectError(error.UnsupportedSstableVersion, Reader.open(allocator, tmp.dir, "test.sst", 1));
|
||||
}
|
||||
|
||||
test "prefix compression measurably shrinks sstables with a shared key prefix" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
const count = 3000;
|
||||
|
||||
{
|
||||
var builder = try Builder.init(allocator, tmp.dir, "shared.sst", count);
|
||||
defer builder.deinit();
|
||||
var buf: [64]u8 = undefined;
|
||||
var i: usize = 0;
|
||||
while (i < count) : (i += 1) {
|
||||
const key = try std.fmt.bufPrint(&buf, "users\x00city\x00warsaw\x00pk{d:0>8}", .{i});
|
||||
try builder.add(.{ .key = key, .value = "v", .seq = i + 1, .kind = .value });
|
||||
}
|
||||
try builder.finish();
|
||||
}
|
||||
{
|
||||
var builder = try Builder.init(allocator, tmp.dir, "distinct.sst", count);
|
||||
defer builder.deinit();
|
||||
var buf: [64]u8 = undefined;
|
||||
var i: usize = 0;
|
||||
while (i < count) : (i += 1) {
|
||||
// Same value payload and roughly the same key length as above,
|
||||
// but scrambled so consecutive keys share no meaningful prefix -
|
||||
// isolates the saving to prefix compression rather than the
|
||||
// block-level RLE codec, which would compress a repeated value
|
||||
// just as well in both files.
|
||||
const scrambled = i *% 2654435761;
|
||||
const key = try std.fmt.bufPrint(&buf, "{d:0>20}", .{scrambled});
|
||||
try builder.add(.{ .key = key, .value = "v", .seq = i + 1, .kind = .value });
|
||||
}
|
||||
try builder.finish();
|
||||
}
|
||||
|
||||
var sharedReader = try Reader.open(allocator, tmp.dir, "shared.sst", 1);
|
||||
defer sharedReader.deinit();
|
||||
var distinctReader = try Reader.open(allocator, tmp.dir, "distinct.sst", 2);
|
||||
defer distinctReader.deinit();
|
||||
|
||||
try std.testing.expectEqual(@as(u64, count), sharedReader.numEntries);
|
||||
try std.testing.expectEqual(@as(u64, count), distinctReader.numEntries);
|
||||
// Prefix compression must shrink the file meaningfully below the
|
||||
// no-shared-prefix baseline (measured ~25% smaller at this fixture
|
||||
// size), not just avoid growing it. Bloom filter and index overhead
|
||||
// is identical in both files, so the margin here comes entirely from
|
||||
// smaller data blocks.
|
||||
try std.testing.expect(sharedReader.physicalBytes * 5 < distinctReader.physicalBytes * 4);
|
||||
|
||||
// Correctness survives the smaller on-disk encoding.
|
||||
try std.testing.expectEqualStrings("v", sharedReader.find("users\x00city\x00warsaw\x00pk00000000").?.value);
|
||||
try std.testing.expectEqualStrings("v", sharedReader.find("users\x00city\x00warsaw\x00pk00000150").?.value);
|
||||
try std.testing.expectEqualStrings("v", sharedReader.find("users\x00city\x00warsaw\x00pk00000299").?.value);
|
||||
try std.testing.expectEqual(@as(?Record, null), sharedReader.find("users\x00city\x00warsaw\x00pk99999999"));
|
||||
}
|
||||
+293
@@ -0,0 +1,293 @@
|
||||
const std = @import("std");
|
||||
const dbmod = @import("db.zig");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
const Db = dbmod.Db;
|
||||
|
||||
/// Per-key expiry markers layered over the flat KV keyspace, sharing one Db
|
||||
/// with tables/collections/auth (row.zig/index.zig use 'R'/'I'/'J'/'M',
|
||||
/// collections.zig uses 'X', auth.zig uses 'A'; every key here starts with
|
||||
/// 'T' so none of the namespaces collide):
|
||||
/// "T" <key> -> raw i128 (expiresAtNanos, big-endian)
|
||||
/// The marker lives in the SAME Db as the key it guards (same WAL, same
|
||||
/// flush/compaction lifecycle, same crash-durability boundary) but expiry
|
||||
/// is enforced lazily: nothing here runs a background clock itself. A
|
||||
/// caller must present nowNanos, either on every read (expireIfDue/getLive)
|
||||
/// or through an explicit sweep (sweepExpired) - matching how session.zig's
|
||||
/// SessionStore already treats "now" as an input rather than ambient state
|
||||
/// (and uses the same i128 type std.time.nanoTimestamp() returns), so both
|
||||
/// are trivially testable without sleeping in unit tests.
|
||||
pub const nsTag: u8 = 'T';
|
||||
|
||||
pub const secondNanos: i128 = 1_000_000_000;
|
||||
|
||||
fn markerKey(allocator: Allocator, key: []const u8) ![]u8 {
|
||||
const out = try allocator.alloc(u8, key.len + 1);
|
||||
out[0] = nsTag;
|
||||
@memcpy(out[1..], key);
|
||||
return out;
|
||||
}
|
||||
|
||||
fn putI128(out: *std.ArrayList(u8), v: i128) !void {
|
||||
var b: [16]u8 = undefined;
|
||||
std.mem.writeInt(i128, &b, v, .big);
|
||||
try out.appendSlice(&b);
|
||||
}
|
||||
|
||||
fn readI128(data: []const u8) !i128 {
|
||||
if (data.len != 16) return error.Malformed;
|
||||
return std.mem.readInt(i128, data[0..16], .big);
|
||||
}
|
||||
|
||||
/// Sets (or replaces) an absolute expiry deadline on `key`. Does not verify
|
||||
/// `key` currently exists in the main keyspace - callers that need
|
||||
/// existence-checked semantics (the CLI `expire` command) check db.get
|
||||
/// first, same as Redis distinguishes EXPIRE from the looser PEXPIREAT.
|
||||
pub fn setExpireAt(db: *Db, allocator: Allocator, key: []const u8, expiresAtNanos: i128) !void {
|
||||
const k = try markerKey(allocator, key);
|
||||
defer allocator.free(k);
|
||||
var out = std.ArrayList(u8).init(allocator);
|
||||
defer out.deinit();
|
||||
try putI128(&out, expiresAtNanos);
|
||||
try db.put(k, out.items);
|
||||
}
|
||||
|
||||
pub fn setExpireSeconds(db: *Db, allocator: Allocator, key: []const u8, nowNanos: i128, seconds: u64) !void {
|
||||
try setExpireAt(db, allocator, key, nowNanos + @as(i128, @intCast(seconds)) * secondNanos);
|
||||
}
|
||||
|
||||
pub fn getExpiresAt(db: *Db, allocator: Allocator, key: []const u8) !?i128 {
|
||||
const k = try markerKey(allocator, key);
|
||||
defer allocator.free(k);
|
||||
const raw = try db.get(allocator, k) orelse return null;
|
||||
defer allocator.free(raw);
|
||||
return try readI128(raw);
|
||||
}
|
||||
|
||||
pub fn clearExpire(db: *Db, allocator: Allocator, key: []const u8) !void {
|
||||
const k = try markerKey(allocator, key);
|
||||
defer allocator.free(k);
|
||||
try db.delete(k);
|
||||
}
|
||||
|
||||
pub fn isDue(expiresAtNanos: i128, nowNanos: i128) bool {
|
||||
return nowNanos >= expiresAtNanos;
|
||||
}
|
||||
|
||||
/// Redis TTL(2) semantics for the CLI: `.noKey` mirrors return -2 (the key
|
||||
/// itself is absent), `.noExpiry` mirrors -1 (it exists but carries no
|
||||
/// expiry), `.seconds` is the whole seconds left, floored at 0 so a marker
|
||||
/// that is due but not yet reaped never reports a negative remainder.
|
||||
pub const TtlResult = union(enum) {
|
||||
noKey,
|
||||
noExpiry,
|
||||
seconds: u64,
|
||||
};
|
||||
|
||||
pub fn remaining(db: *Db, allocator: Allocator, key: []const u8, nowNanos: i128) !TtlResult {
|
||||
const existing = try db.get(allocator, key) orelse return .noKey;
|
||||
allocator.free(existing);
|
||||
const expiresAt = try getExpiresAt(db, allocator, key) orelse return .noExpiry;
|
||||
if (nowNanos >= expiresAt) return .{ .seconds = 0 };
|
||||
return .{ .seconds = @intCast(@divFloor(expiresAt - nowNanos, secondNanos)) };
|
||||
}
|
||||
|
||||
/// Lazily enforces expiry on read: if `key` carries a due marker, deletes
|
||||
/// both the marker and the underlying key (a real delete through the normal
|
||||
/// WAL/tombstone path, not just "pretend it's gone") and reports true.
|
||||
pub fn expireIfDue(db: *Db, allocator: Allocator, key: []const u8, nowNanos: i128) !bool {
|
||||
const expiresAt = try getExpiresAt(db, allocator, key) orelse return false;
|
||||
if (!isDue(expiresAt, nowNanos)) return false;
|
||||
try db.delete(key);
|
||||
try clearExpire(db, allocator, key);
|
||||
return true;
|
||||
}
|
||||
|
||||
/// db.get with lazy expiry applied first - the read-path counterpart of
|
||||
/// putWithTtl, so a caller never observes a logically-expired value just
|
||||
/// because no sweep has run yet.
|
||||
pub fn getLive(db: *Db, allocator: Allocator, key: []const u8, nowNanos: i128) !?[]u8 {
|
||||
_ = try expireIfDue(db, allocator, key, nowNanos);
|
||||
return db.get(allocator, key);
|
||||
}
|
||||
|
||||
pub fn putWithTtl(db: *Db, allocator: Allocator, key: []const u8, value: []const u8, nowNanos: i128, seconds: u64) !void {
|
||||
try db.put(key, value);
|
||||
try setExpireSeconds(db, allocator, key, nowNanos, seconds);
|
||||
}
|
||||
|
||||
/// Active sweep: scans every marker under the 'T' namespace and reaps every
|
||||
/// one that is due, deleting both the underlying key and its marker.
|
||||
/// Returns the number reaped. Two passes on purpose - a RangeIterator's
|
||||
/// merge heap holds live cursor positions directly into the memtable being
|
||||
/// walked (only the Entry it hands back through next() is arena-copied),
|
||||
/// so issuing db.delete mid-scan would mutate the very structure the
|
||||
/// iterator's cursors are threaded through. Collect the due keys first,
|
||||
/// mutate only after the iterator (and its cursors) are closed.
|
||||
pub fn sweepExpired(db: *Db, allocator: Allocator, nowNanos: i128) !usize {
|
||||
var due = std.ArrayList([]u8).init(allocator);
|
||||
defer {
|
||||
for (due.items) |k| allocator.free(k);
|
||||
due.deinit();
|
||||
}
|
||||
{
|
||||
const prefix = &[_]u8{nsTag};
|
||||
var iter = try db.scanPrefix(allocator, prefix);
|
||||
defer iter.deinit();
|
||||
while (try iter.next()) |entry| {
|
||||
const expiresAt = try readI128(entry.value);
|
||||
if (isDue(expiresAt, nowNanos)) {
|
||||
try due.append(try allocator.dupe(u8, entry.key[1..]));
|
||||
}
|
||||
}
|
||||
}
|
||||
for (due.items) |userKey| {
|
||||
try db.delete(userKey);
|
||||
try clearExpire(db, allocator, userKey);
|
||||
}
|
||||
return due.items.len;
|
||||
}
|
||||
|
||||
fn openTemp(tmp: *std.testing.TmpDir, allocator: Allocator) !Db {
|
||||
var pathBuf: [std.fs.max_path_bytes]u8 = undefined;
|
||||
const base = try tmp.dir.realpath(".", &pathBuf);
|
||||
var dbPathBuf: [std.fs.max_path_bytes]u8 = undefined;
|
||||
const dbPath = try std.fmt.bufPrint(&dbPathBuf, "{s}/store", .{base});
|
||||
return Db.open(allocator, dbPath, .{});
|
||||
}
|
||||
|
||||
test "remaining reports noKey, noExpiry, and counted seconds" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
try std.testing.expectEqual(TtlResult.noKey, try remaining(&db, allocator, "missing", 0));
|
||||
|
||||
try db.put("plain", "v");
|
||||
try std.testing.expectEqual(TtlResult.noExpiry, try remaining(&db, allocator, "plain", 0));
|
||||
|
||||
try setExpireSeconds(&db, allocator, "plain", 1000, 30);
|
||||
const r = try remaining(&db, allocator, "plain", 1000);
|
||||
try std.testing.expectEqual(@as(u64, 30), r.seconds);
|
||||
}
|
||||
|
||||
test "expireIfDue removes value and marker only once due" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
try db.put("session:1", "payload");
|
||||
try setExpireSeconds(&db, allocator, "session:1", 0, 5);
|
||||
|
||||
try std.testing.expect(!try expireIfDue(&db, allocator, "session:1", 4 * secondNanos));
|
||||
const stillThere = (try db.get(allocator, "session:1")).?;
|
||||
allocator.free(stillThere);
|
||||
|
||||
try std.testing.expect(try expireIfDue(&db, allocator, "session:1", 5 * secondNanos));
|
||||
try std.testing.expectEqual(@as(?[]u8, null), try db.get(allocator, "session:1"));
|
||||
try std.testing.expectEqual(@as(?i128, null), try getExpiresAt(&db, allocator, "session:1"));
|
||||
}
|
||||
|
||||
test "getLive lazily expires exactly at the deadline" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
try putWithTtl(&db, allocator, "cache:x", "hot", 0, 2);
|
||||
|
||||
const before = (try getLive(&db, allocator, "cache:x", 1 * secondNanos)).?;
|
||||
defer allocator.free(before);
|
||||
try std.testing.expectEqualStrings("hot", before);
|
||||
|
||||
try std.testing.expectEqual(@as(?[]u8, null), try getLive(&db, allocator, "cache:x", 2 * secondNanos));
|
||||
// A plain db.get after the lazy expiry proves the delete was real, not
|
||||
// getLive-only bookkeeping - the tombstone is now the source of truth.
|
||||
try std.testing.expectEqual(@as(?[]u8, null), try db.get(allocator, "cache:x"));
|
||||
}
|
||||
|
||||
test "clearExpire makes a key immortal again" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
try putWithTtl(&db, allocator, "note", "keep me", 0, 1);
|
||||
try clearExpire(&db, allocator, "note");
|
||||
|
||||
try std.testing.expectEqual(TtlResult.noExpiry, try remaining(&db, allocator, "note", 100 * secondNanos));
|
||||
const v = (try getLive(&db, allocator, "note", 100 * secondNanos)).?;
|
||||
defer allocator.free(v);
|
||||
try std.testing.expectEqualStrings("keep me", v);
|
||||
}
|
||||
|
||||
test "sweepExpired reaps every due key without touching live or unrelated keys" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
try putWithTtl(&db, allocator, "a", "1", 0, 1);
|
||||
try putWithTtl(&db, allocator, "b", "2", 0, 100);
|
||||
try putWithTtl(&db, allocator, "c", "3", 0, 1);
|
||||
try db.put("d", "forever");
|
||||
|
||||
const reaped = try sweepExpired(&db, allocator, 2 * secondNanos);
|
||||
try std.testing.expectEqual(@as(usize, 2), reaped);
|
||||
|
||||
try std.testing.expectEqual(@as(?[]u8, null), try db.get(allocator, "a"));
|
||||
try std.testing.expectEqual(@as(?[]u8, null), try db.get(allocator, "c"));
|
||||
|
||||
const b = (try db.get(allocator, "b")).?;
|
||||
defer allocator.free(b);
|
||||
try std.testing.expectEqualStrings("2", b);
|
||||
|
||||
const d = (try db.get(allocator, "d")).?;
|
||||
defer allocator.free(d);
|
||||
try std.testing.expectEqualStrings("forever", d);
|
||||
|
||||
// Idempotent: a second sweep at the same instant finds nothing left.
|
||||
try std.testing.expectEqual(@as(usize, 0), try sweepExpired(&db, allocator, 2 * secondNanos));
|
||||
}
|
||||
|
||||
test "sweep survives a wal reopen: markers are durable, not in-memory bookkeeping" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
{
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
try putWithTtl(&db, allocator, "durable-ttl", "v", 0, 5);
|
||||
}
|
||||
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
const r = try remaining(&db, allocator, "durable-ttl", 1 * secondNanos);
|
||||
try std.testing.expectEqual(@as(u64, 4), r.seconds);
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 1), try sweepExpired(&db, allocator, 10 * secondNanos));
|
||||
try std.testing.expectEqual(@as(?[]u8, null), try db.get(allocator, "durable-ttl"));
|
||||
}
|
||||
|
||||
test "setExpireAt on a nonexistent key is a real marker, not silently rejected" {
|
||||
// Documents the PEXPIREAT-like looseness of setExpireAt directly (the
|
||||
// CLI layer is what adds the Redis EXPIRE existence check on top).
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
try setExpireSeconds(&db, allocator, "ghost", 0, 10);
|
||||
try std.testing.expectEqual(TtlResult.noKey, try remaining(&db, allocator, "ghost", 0));
|
||||
const expiresAt = (try getExpiresAt(&db, allocator, "ghost")).?;
|
||||
try std.testing.expectEqual(@as(i128, 10 * secondNanos), expiresAt);
|
||||
}
|
||||
+127
@@ -81,6 +81,62 @@ pub const Transaction = struct {
|
||||
try self.reads.put(keyCopy, {});
|
||||
}
|
||||
|
||||
/// One key/value pair from `scanPrefix`, owned by the caller (free with
|
||||
/// `freeEntries`).
|
||||
pub const Entry = struct {
|
||||
key: []u8,
|
||||
value: []u8,
|
||||
|
||||
fn free(self: Entry, allocator: Allocator) void {
|
||||
allocator.free(self.key);
|
||||
allocator.free(self.value);
|
||||
}
|
||||
};
|
||||
|
||||
pub fn freeEntries(allocator: Allocator, entries: []Entry) void {
|
||||
for (entries) |e| e.free(allocator);
|
||||
allocator.free(entries);
|
||||
}
|
||||
|
||||
/// Prefix scan that reflects this transaction's own uncommitted writes,
|
||||
/// the same read-your-writes promise `get` makes for a single key
|
||||
/// extended to a whole key range: a row staged here shadows whatever the
|
||||
/// committed snapshot has under that exact key (a staged tombstone hides
|
||||
/// it entirely, a staged value replaces it), and a brand-new key staged
|
||||
/// here that falls under `prefix` shows up even though the underlying
|
||||
/// `Db` has never durably seen it. Every entry the underlying snapshot
|
||||
/// contributes is also recorded as a read for OCC purposes, same as
|
||||
/// `get` - a concurrent external write to one of those existing keys
|
||||
/// between now and `commit()` is still caught as `error.Conflict`. A
|
||||
/// concurrent external INSERT of a brand-new key under `prefix` is not
|
||||
/// caught (there is no key to register a read against before it
|
||||
/// exists) - the same phantom-read gap every OCC scheme without
|
||||
/// range-locks has, not something specific to this method.
|
||||
pub fn scanPrefix(self: *Transaction, allocator: Allocator, prefix: []const u8) ![]Entry {
|
||||
var results = std.ArrayList(Entry).init(allocator);
|
||||
errdefer {
|
||||
for (results.items) |e| e.free(allocator);
|
||||
results.deinit();
|
||||
}
|
||||
|
||||
var it = try self.db.scanPrefixAsOf(allocator, prefix, self.snapshotSeq);
|
||||
defer it.deinit();
|
||||
while (try it.next()) |entry| {
|
||||
if (self.writes.contains(entry.key)) continue;
|
||||
try self.recordRead(entry.key);
|
||||
try results.append(.{ .key = try allocator.dupe(u8, entry.key), .value = try allocator.dupe(u8, entry.value) });
|
||||
}
|
||||
|
||||
var writeEntries = self.writes.iterator();
|
||||
while (writeEntries.next()) |entry| {
|
||||
if (!std.mem.startsWith(u8, entry.key_ptr.*, prefix)) continue;
|
||||
if (entry.value_ptr.kind == .tombstone) continue;
|
||||
try results.append(.{ .key = try allocator.dupe(u8, entry.key_ptr.*), .value = try allocator.dupe(u8, entry.value_ptr.value) });
|
||||
}
|
||||
|
||||
return results.toOwnedSlice();
|
||||
}
|
||||
|
||||
pub fn commit(self: *Transaction) !void {
|
||||
var readKeys = self.reads.keyIterator();
|
||||
while (readKeys.next()) |key| {
|
||||
@@ -186,3 +242,74 @@ test "conflicting commit is rejected but clean commit succeeds" {
|
||||
defer allocator.free(note);
|
||||
try std.testing.expectEqualStrings("ok", note);
|
||||
}
|
||||
|
||||
test "scanPrefix sees committed rows and a staged brand-new key under the same prefix" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
try db.put("idx:a", "1");
|
||||
try db.put("idx:b", "2");
|
||||
try db.put("other:z", "9");
|
||||
|
||||
var txn = Transaction.begin(&db);
|
||||
defer txn.deinit();
|
||||
try txn.put("idx:c", "3");
|
||||
|
||||
const entries = try txn.scanPrefix(allocator, "idx:");
|
||||
defer Transaction.freeEntries(allocator, entries);
|
||||
try std.testing.expectEqual(@as(usize, 3), entries.len);
|
||||
|
||||
var seenC = false;
|
||||
for (entries) |e| {
|
||||
if (std.mem.eql(u8, e.key, "idx:c")) {
|
||||
try std.testing.expectEqualStrings("3", e.value);
|
||||
seenC = true;
|
||||
}
|
||||
}
|
||||
try std.testing.expect(seenC);
|
||||
}
|
||||
|
||||
test "scanPrefix shadows a committed key with its staged replacement and hides a staged tombstone" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
try db.put("idx:a", "old-a");
|
||||
try db.put("idx:b", "old-b");
|
||||
|
||||
var txn = Transaction.begin(&db);
|
||||
defer txn.deinit();
|
||||
try txn.put("idx:a", "new-a");
|
||||
try txn.delete("idx:b");
|
||||
|
||||
const entries = try txn.scanPrefix(allocator, "idx:");
|
||||
defer Transaction.freeEntries(allocator, entries);
|
||||
try std.testing.expectEqual(@as(usize, 1), entries.len);
|
||||
try std.testing.expectEqualStrings("idx:a", entries[0].key);
|
||||
try std.testing.expectEqualStrings("new-a", entries[0].value);
|
||||
}
|
||||
|
||||
test "scanPrefix registers committed keys as reads, so a concurrent external write under it still conflicts at commit" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
var db = try openTemp(&tmp, allocator);
|
||||
defer db.close();
|
||||
|
||||
try db.put("idx:a", "1");
|
||||
|
||||
var txn = Transaction.begin(&db);
|
||||
defer txn.deinit();
|
||||
const entries = try txn.scanPrefix(allocator, "idx:");
|
||||
Transaction.freeEntries(allocator, entries);
|
||||
|
||||
try db.put("idx:a", "external-change");
|
||||
|
||||
try txn.put("unrelated", "1");
|
||||
try std.testing.expectError(error.Conflict, txn.commit());
|
||||
}
|
||||
+1029
File diff suppressed because it is too large.
Load diff
+658
@@ -0,0 +1,658 @@
|
||||
const std = @import("std");
|
||||
const encode = @import("encode.zig");
|
||||
|
||||
const Allocator = std.mem.Allocator;
|
||||
|
||||
// VersionSet/VersionEdit: append-only log of incremental changes to the
|
||||
// live sstable set (fileId -> level), replacing a full-state rewrite on
|
||||
// every flush/compact. Modeled on LevelDB's VersionSet: a CURRENT file
|
||||
// names the active manifest log, the log holds a sequence of edit
|
||||
// records, and recovery replays them into a map. The log is periodically
|
||||
// rotated (a single snapshot record replaces the whole history) so it
|
||||
// never grows without bound.
|
||||
|
||||
pub const VersionError = error{
|
||||
Corrupt,
|
||||
Truncated,
|
||||
};
|
||||
|
||||
pub const FileEdit = struct {
|
||||
fileId: u64,
|
||||
level: u32,
|
||||
};
|
||||
|
||||
pub const VersionEdit = struct {
|
||||
allocator: Allocator,
|
||||
added: std.ArrayList(FileEdit),
|
||||
deleted: std.ArrayList(u64),
|
||||
/// Optional counter watermarks, stamped by the caller alongside a file
|
||||
/// add/delete: the exact `Db.nextFileId`/`Db.nextSeq` value in effect
|
||||
/// right after this edit was built. Recorded so recovery can recover
|
||||
/// these counters from the manifest instead of relying solely on
|
||||
/// scanning which sstable files currently happen to be on disk - a scan
|
||||
/// alone regresses whenever a compaction fully garbage-collects a run of
|
||||
/// tombstones into zero output files (see runMerge: no output file means
|
||||
/// no fresh fileId is minted, and if the GC'd inputs held the highest
|
||||
/// seq/fileId ever used, a post-restart scan would compute a lower
|
||||
/// floor than what was really already consumed, risking seq reuse
|
||||
/// across a restart).
|
||||
nextFileId: ?u64 = null,
|
||||
nextSeq: ?u64 = null,
|
||||
|
||||
pub fn init(allocator: Allocator) VersionEdit {
|
||||
return .{
|
||||
.allocator = allocator,
|
||||
.added = std.ArrayList(FileEdit).init(allocator),
|
||||
.deleted = std.ArrayList(u64).init(allocator),
|
||||
};
|
||||
}
|
||||
|
||||
pub fn deinit(self: *VersionEdit) void {
|
||||
self.added.deinit();
|
||||
self.deleted.deinit();
|
||||
}
|
||||
|
||||
pub fn addFile(self: *VersionEdit, fileId: u64, level: usize) !void {
|
||||
try self.added.append(.{ .fileId = fileId, .level = @intCast(level) });
|
||||
}
|
||||
|
||||
pub fn deleteFile(self: *VersionEdit, fileId: u64) !void {
|
||||
try self.deleted.append(fileId);
|
||||
}
|
||||
|
||||
pub fn setNextFileId(self: *VersionEdit, nextFileId: u64) void {
|
||||
self.nextFileId = nextFileId;
|
||||
}
|
||||
|
||||
pub fn setNextSeq(self: *VersionEdit, nextSeq: u64) void {
|
||||
self.nextSeq = nextSeq;
|
||||
}
|
||||
|
||||
pub fn isEmpty(self: *const VersionEdit) bool {
|
||||
return self.added.items.len == 0 and self.deleted.items.len == 0;
|
||||
}
|
||||
|
||||
fn encode_(self: *const VersionEdit, out: *std.ArrayList(u8)) !void {
|
||||
const writer = encode.ByteWriter.init(out);
|
||||
try writer.putU32(@intCast(self.added.items.len));
|
||||
for (self.added.items) |f| {
|
||||
try writer.putU64(f.fileId);
|
||||
try writer.putU32(f.level);
|
||||
}
|
||||
try writer.putU32(@intCast(self.deleted.items.len));
|
||||
for (self.deleted.items) |id| try writer.putU64(id);
|
||||
try writer.byte(if (self.nextFileId != null) 1 else 0);
|
||||
if (self.nextFileId) |v| try writer.putU64(v);
|
||||
try writer.byte(if (self.nextSeq != null) 1 else 0);
|
||||
if (self.nextSeq) |v| try writer.putU64(v);
|
||||
}
|
||||
};
|
||||
|
||||
/// Applies an edit's added/deleted entries onto a live fileId->level map.
|
||||
pub fn applyEdit(state: *std.AutoHashMap(u64, usize), added: []const FileEdit, deleted: []const u64) !void {
|
||||
for (deleted) |id| _ = state.remove(id);
|
||||
for (added) |f| try state.put(f.fileId, f.level);
|
||||
}
|
||||
|
||||
/// Recovered counter watermarks, threaded out of `VersionSet.open`'s replay
|
||||
/// so `Db.open` can fold them into its own directory-scan-based recovery
|
||||
/// (taking the max of both sources - see `VersionEdit.nextFileId`/`nextSeq`
|
||||
/// doc comment for why the scan alone is not always sufficient).
|
||||
pub const Counters = struct {
|
||||
nextFileId: ?u64 = null,
|
||||
nextSeq: ?u64 = null,
|
||||
};
|
||||
|
||||
const currentName = "CURRENT";
|
||||
const manifestPrefix = "MANIFEST-";
|
||||
|
||||
fn manifestFileName(buf: []u8, id: u64) ![]const u8 {
|
||||
return std.fmt.bufPrint(buf, "{s}{d:0>10}", .{ manifestPrefix, id });
|
||||
}
|
||||
|
||||
pub const VersionSet = struct {
|
||||
allocator: Allocator,
|
||||
dir: std.fs.Dir,
|
||||
file: std.fs.File,
|
||||
manifestId: u64,
|
||||
/// Number of edit records appended to the current log since the last
|
||||
/// rotation (used as the rotation heuristic).
|
||||
pendingRecords: usize,
|
||||
/// Latest counter watermarks seen across every replayed (or since
|
||||
/// applied) edit - see `VersionEdit.nextFileId`/`nextSeq`. Carried
|
||||
/// forward into `rotate`'s snapshot record so they survive log rotation
|
||||
/// even though the individual edits that set them get discarded.
|
||||
counters: Counters,
|
||||
|
||||
/// Opens (or bootstraps) the version set, replaying the active
|
||||
/// manifest log into a fresh fileId->level map that the caller owns.
|
||||
pub fn open(allocator: Allocator, dir: std.fs.Dir) !struct { set: VersionSet, state: std.AutoHashMap(u64, usize) } {
|
||||
var state = std.AutoHashMap(u64, usize).init(allocator);
|
||||
errdefer state.deinit();
|
||||
|
||||
const currentContents = dir.readFileAlloc(allocator, currentName, 256) catch |err| switch (err) {
|
||||
error.FileNotFound => {
|
||||
const set = try bootstrap(allocator, dir);
|
||||
return .{ .set = set, .state = state };
|
||||
},
|
||||
else => return err,
|
||||
};
|
||||
defer allocator.free(currentContents);
|
||||
|
||||
const trimmed = std.mem.trim(u8, currentContents, " \t\r\n");
|
||||
const id = parseManifestId(trimmed) orelse return VersionError.Corrupt;
|
||||
|
||||
const file = dir.openFile(trimmed, .{ .mode = .read_write }) catch |err| switch (err) {
|
||||
error.FileNotFound => {
|
||||
const set = try bootstrap(allocator, dir);
|
||||
return .{ .set = set, .state = state };
|
||||
},
|
||||
else => return err,
|
||||
};
|
||||
errdefer file.close();
|
||||
|
||||
const data = try file.readToEndAlloc(allocator, std.math.maxInt(usize));
|
||||
defer allocator.free(data);
|
||||
|
||||
var pos: usize = 0;
|
||||
var goodEnd: usize = 0;
|
||||
var records: usize = 0;
|
||||
var counters = Counters{};
|
||||
while (pos + 8 <= data.len) {
|
||||
const storedCrc = encode.readU32(data, pos);
|
||||
const bodyStart = pos + 4;
|
||||
const bodyLen = recordBodyLen(data, bodyStart) orelse break;
|
||||
if (bodyStart + bodyLen > data.len) break;
|
||||
const body = data[bodyStart .. bodyStart + bodyLen];
|
||||
if (encode.crc(body) != storedCrc) break;
|
||||
|
||||
try replayRecord(&state, &counters, body);
|
||||
records += 1;
|
||||
pos = bodyStart + bodyLen;
|
||||
goodEnd = pos;
|
||||
}
|
||||
|
||||
if (goodEnd != data.len) {
|
||||
try file.setEndPos(goodEnd);
|
||||
try file.sync();
|
||||
}
|
||||
try file.seekFromEnd(0);
|
||||
|
||||
return .{
|
||||
.set = .{
|
||||
.allocator = allocator,
|
||||
.dir = dir,
|
||||
.file = file,
|
||||
.manifestId = id,
|
||||
.pendingRecords = records,
|
||||
.counters = counters,
|
||||
},
|
||||
.state = state,
|
||||
};
|
||||
}
|
||||
|
||||
fn bootstrap(allocator: Allocator, dir: std.fs.Dir) !VersionSet {
|
||||
var nameBuf: [64]u8 = undefined;
|
||||
const name = try manifestFileName(&nameBuf, 1);
|
||||
const file = try dir.createFile(name, .{ .read = true, .truncate = true });
|
||||
errdefer file.close();
|
||||
try writeCurrent(dir, name);
|
||||
return .{
|
||||
.allocator = allocator,
|
||||
.dir = dir,
|
||||
.file = file,
|
||||
.manifestId = 1,
|
||||
.pendingRecords = 0,
|
||||
.counters = .{},
|
||||
};
|
||||
}
|
||||
|
||||
pub fn close(self: *VersionSet) void {
|
||||
self.file.close();
|
||||
}
|
||||
|
||||
/// Applies `edit` to `state` and durably appends it to the manifest
|
||||
/// log. Rotates (compacts the log to a single snapshot record) once
|
||||
/// the number of pending records grows past the live file count.
|
||||
pub fn logAndApply(self: *VersionSet, edit: *const VersionEdit, state: *std.AutoHashMap(u64, usize)) !void {
|
||||
if (edit.isEmpty()) return;
|
||||
try applyEdit(state, edit.added.items, edit.deleted.items);
|
||||
if (edit.nextFileId) |v| self.counters.nextFileId = v;
|
||||
if (edit.nextSeq) |v| self.counters.nextSeq = v;
|
||||
try self.appendRecord(edit);
|
||||
self.pendingRecords += 1;
|
||||
|
||||
const liveCount = state.count();
|
||||
if (self.pendingRecords > @max(liveCount, 4) * 2) {
|
||||
try self.rotate(state);
|
||||
}
|
||||
}
|
||||
|
||||
fn appendRecord(self: *VersionSet, edit: *const VersionEdit) !void {
|
||||
var body = std.ArrayList(u8).init(self.allocator);
|
||||
defer body.deinit();
|
||||
try edit.encode_(&body);
|
||||
|
||||
var frame = std.ArrayList(u8).init(self.allocator);
|
||||
defer frame.deinit();
|
||||
const frameWriter = encode.ByteWriter.init(&frame);
|
||||
try frameWriter.putU32(encode.crc(body.items));
|
||||
try frameWriter.slice(body.items);
|
||||
|
||||
try self.file.seekFromEnd(0);
|
||||
try self.file.writeAll(frame.items);
|
||||
try self.file.sync();
|
||||
}
|
||||
|
||||
/// Replaces the whole log history with a single snapshot record that
|
||||
/// reproduces the current live state, then switches CURRENT to it and
|
||||
/// deletes the old log. Bounds manifest size to O(live files) instead
|
||||
/// of O(edits ever applied).
|
||||
fn rotate(self: *VersionSet, state: *const std.AutoHashMap(u64, usize)) !void {
|
||||
var snapshot = std.ArrayList(FileEdit).init(self.allocator);
|
||||
defer snapshot.deinit();
|
||||
var it = state.iterator();
|
||||
while (it.next()) |entry| {
|
||||
try snapshot.append(.{ .fileId = entry.key_ptr.*, .level = @intCast(entry.value_ptr.*) });
|
||||
}
|
||||
|
||||
const newId = self.manifestId + 1;
|
||||
var nameBuf: [64]u8 = undefined;
|
||||
const newName = try manifestFileName(&nameBuf, newId);
|
||||
|
||||
const newFile = try self.dir.createFile(newName, .{ .read = true, .truncate = true });
|
||||
errdefer newFile.close();
|
||||
|
||||
{
|
||||
var body = std.ArrayList(u8).init(self.allocator);
|
||||
defer body.deinit();
|
||||
const writer = encode.ByteWriter.init(&body);
|
||||
try writer.putU32(@intCast(snapshot.items.len));
|
||||
for (snapshot.items) |f| {
|
||||
try writer.putU64(f.fileId);
|
||||
try writer.putU32(f.level);
|
||||
}
|
||||
try writer.putU32(0);
|
||||
// Carry the latest known counters forward into the snapshot so
|
||||
// they survive rotation even though the edits that originally
|
||||
// set them are being discarded along with the old log.
|
||||
try writer.byte(if (self.counters.nextFileId != null) 1 else 0);
|
||||
if (self.counters.nextFileId) |v| try writer.putU64(v);
|
||||
try writer.byte(if (self.counters.nextSeq != null) 1 else 0);
|
||||
if (self.counters.nextSeq) |v| try writer.putU64(v);
|
||||
|
||||
var frame = std.ArrayList(u8).init(self.allocator);
|
||||
defer frame.deinit();
|
||||
const frameWriter = encode.ByteWriter.init(&frame);
|
||||
try frameWriter.putU32(encode.crc(body.items));
|
||||
try frameWriter.slice(body.items);
|
||||
try newFile.writeAll(frame.items);
|
||||
try newFile.sync();
|
||||
}
|
||||
|
||||
try writeCurrent(self.dir, newName);
|
||||
|
||||
var oldNameBuf: [64]u8 = undefined;
|
||||
const oldName = try manifestFileName(&oldNameBuf, self.manifestId);
|
||||
self.file.close();
|
||||
self.dir.deleteFile(oldName) catch {};
|
||||
|
||||
self.file = newFile;
|
||||
self.manifestId = newId;
|
||||
self.pendingRecords = snapshot.items.len;
|
||||
}
|
||||
};
|
||||
|
||||
fn writeCurrent(dir: std.fs.Dir, manifestName: []const u8) !void {
|
||||
const tmpName = "CURRENT.tmp";
|
||||
{
|
||||
const file = try dir.createFile(tmpName, .{ .truncate = true });
|
||||
defer file.close();
|
||||
try file.writeAll(manifestName);
|
||||
try file.sync();
|
||||
}
|
||||
try dir.rename(tmpName, currentName);
|
||||
}
|
||||
|
||||
fn parseManifestId(name: []const u8) ?u64 {
|
||||
if (!std.mem.startsWith(u8, name, manifestPrefix)) return null;
|
||||
return std.fmt.parseInt(u64, name[manifestPrefix.len..], 10) catch null;
|
||||
}
|
||||
|
||||
fn recordBodyLen(data: []const u8, bodyStart: usize) ?usize {
|
||||
if (bodyStart + 4 > data.len) return null;
|
||||
const addedCount = encode.readU32(data, bodyStart);
|
||||
var p = bodyStart + 4;
|
||||
var i: u32 = 0;
|
||||
while (i < addedCount) : (i += 1) {
|
||||
if (p + 12 > data.len) return null;
|
||||
p += 12;
|
||||
}
|
||||
if (p + 4 > data.len) return null;
|
||||
const deletedCount = encode.readU32(data, p);
|
||||
p += 4;
|
||||
var j: u32 = 0;
|
||||
while (j < deletedCount) : (j += 1) {
|
||||
if (p + 8 > data.len) return null;
|
||||
p += 8;
|
||||
}
|
||||
// Trailing counter watermarks: one presence flag byte plus an optional
|
||||
// u64 each, for nextFileId then nextSeq.
|
||||
var k: u32 = 0;
|
||||
while (k < 2) : (k += 1) {
|
||||
if (p + 1 > data.len) return null;
|
||||
const present = data[p];
|
||||
p += 1;
|
||||
if (present != 0) {
|
||||
if (p + 8 > data.len) return null;
|
||||
p += 8;
|
||||
}
|
||||
}
|
||||
return p - bodyStart;
|
||||
}
|
||||
|
||||
fn replayRecord(state: *std.AutoHashMap(u64, usize), counters: *Counters, body: []const u8) !void {
|
||||
const addedCount = encode.readU32(body, 0);
|
||||
var p: usize = 4;
|
||||
var i: u32 = 0;
|
||||
while (i < addedCount) : (i += 1) {
|
||||
const fileId = encode.readU64(body, p);
|
||||
const level = encode.readU32(body, p + 8);
|
||||
try state.put(fileId, level);
|
||||
p += 12;
|
||||
}
|
||||
const deletedCount = encode.readU32(body, p);
|
||||
p += 4;
|
||||
var j: u32 = 0;
|
||||
while (j < deletedCount) : (j += 1) {
|
||||
const fileId = encode.readU64(body, p);
|
||||
_ = state.remove(fileId);
|
||||
p += 8;
|
||||
}
|
||||
if (p < body.len) {
|
||||
const hasNextFileId = body[p];
|
||||
p += 1;
|
||||
if (hasNextFileId != 0) {
|
||||
counters.nextFileId = encode.readU64(body, p);
|
||||
p += 8;
|
||||
}
|
||||
const hasNextSeq = body[p];
|
||||
p += 1;
|
||||
if (hasNextSeq != 0) {
|
||||
counters.nextSeq = encode.readU64(body, p);
|
||||
p += 8;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
test "fresh directory bootstraps an empty version set" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
var opened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer opened.state.deinit();
|
||||
defer opened.set.close();
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 0), opened.state.count());
|
||||
try std.testing.expectEqual(@as(u64, 1), opened.set.manifestId);
|
||||
try std.testing.expect(tmp.dir.access(currentName, .{}) != error.FileNotFound);
|
||||
}
|
||||
|
||||
test "logAndApply records adds and survives reopen" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
{
|
||||
var opened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer opened.state.deinit();
|
||||
defer opened.set.close();
|
||||
|
||||
var edit = VersionEdit.init(allocator);
|
||||
defer edit.deinit();
|
||||
try edit.addFile(1, 0);
|
||||
try edit.addFile(2, 0);
|
||||
try opened.set.logAndApply(&edit, &opened.state);
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 2), opened.state.count());
|
||||
try std.testing.expectEqual(@as(usize, 0), opened.state.get(1).?);
|
||||
}
|
||||
|
||||
var reopened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer reopened.state.deinit();
|
||||
defer reopened.set.close();
|
||||
try std.testing.expectEqual(@as(usize, 2), reopened.state.count());
|
||||
try std.testing.expectEqual(@as(usize, 0), reopened.state.get(2).?);
|
||||
}
|
||||
|
||||
test "logAndApply replaces a deleted file with new outputs (compaction pattern)" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
var opened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer opened.state.deinit();
|
||||
defer opened.set.close();
|
||||
|
||||
var e1 = VersionEdit.init(allocator);
|
||||
defer e1.deinit();
|
||||
try e1.addFile(1, 0);
|
||||
try e1.addFile(2, 0);
|
||||
try opened.set.logAndApply(&e1, &opened.state);
|
||||
|
||||
var e2 = VersionEdit.init(allocator);
|
||||
defer e2.deinit();
|
||||
try e2.deleteFile(1);
|
||||
try e2.deleteFile(2);
|
||||
try e2.addFile(3, 1);
|
||||
try opened.set.logAndApply(&e2, &opened.state);
|
||||
|
||||
try std.testing.expectEqual(@as(usize, 1), opened.state.count());
|
||||
try std.testing.expectEqual(@as(usize, 1), opened.state.get(3).?);
|
||||
try std.testing.expectEqual(@as(?usize, null), opened.state.get(1));
|
||||
}
|
||||
|
||||
test "manifest rotates once pending records outgrow live file count" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
var opened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer opened.state.deinit();
|
||||
defer opened.set.close();
|
||||
|
||||
// Simulate steady-state leveled compaction: the live file count stays
|
||||
// fixed at 1 (each round replaces the previous output), but the edit
|
||||
// log keeps growing until it is worth rotating to a fresh snapshot.
|
||||
var fileId: u64 = 1;
|
||||
{
|
||||
var edit = VersionEdit.init(allocator);
|
||||
defer edit.deinit();
|
||||
try edit.addFile(fileId, 0);
|
||||
try opened.set.logAndApply(&edit, &opened.state);
|
||||
}
|
||||
|
||||
var n: usize = 0;
|
||||
while (n < 20) : (n += 1) {
|
||||
var edit = VersionEdit.init(allocator);
|
||||
defer edit.deinit();
|
||||
try edit.deleteFile(fileId);
|
||||
fileId += 1;
|
||||
try edit.addFile(fileId, 0);
|
||||
try opened.set.logAndApply(&edit, &opened.state);
|
||||
}
|
||||
|
||||
try std.testing.expect(opened.set.manifestId > 1);
|
||||
try std.testing.expectEqual(@as(usize, 1), opened.state.count());
|
||||
try std.testing.expectEqual(@as(usize, 0), opened.state.get(fileId).?);
|
||||
|
||||
var nameBuf: [64]u8 = undefined;
|
||||
const oldName = try manifestFileName(&nameBuf, 1);
|
||||
try std.testing.expectError(error.FileNotFound, tmp.dir.access(oldName, .{}));
|
||||
}
|
||||
|
||||
test "reopen after rotation still reflects live state" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
var lastId: u64 = 0;
|
||||
var finalFileId: u64 = 1;
|
||||
{
|
||||
var opened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer opened.state.deinit();
|
||||
defer opened.set.close();
|
||||
|
||||
{
|
||||
var edit = VersionEdit.init(allocator);
|
||||
defer edit.deinit();
|
||||
try edit.addFile(finalFileId, 0);
|
||||
try opened.set.logAndApply(&edit, &opened.state);
|
||||
}
|
||||
var n: usize = 0;
|
||||
while (n < 12) : (n += 1) {
|
||||
var edit = VersionEdit.init(allocator);
|
||||
defer edit.deinit();
|
||||
try edit.deleteFile(finalFileId);
|
||||
finalFileId += 1;
|
||||
try edit.addFile(finalFileId, 0);
|
||||
try opened.set.logAndApply(&edit, &opened.state);
|
||||
}
|
||||
lastId = opened.set.manifestId;
|
||||
}
|
||||
|
||||
try std.testing.expect(lastId > 1);
|
||||
|
||||
var reopened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer reopened.state.deinit();
|
||||
defer reopened.set.close();
|
||||
try std.testing.expectEqual(@as(usize, 1), reopened.state.count());
|
||||
try std.testing.expectEqual(@as(usize, 0), reopened.state.get(finalFileId).?);
|
||||
try std.testing.expectEqual(reopened.set.manifestId, lastId);
|
||||
}
|
||||
|
||||
test "torn trailing record is ignored on recovery" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
{
|
||||
var opened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer opened.state.deinit();
|
||||
defer opened.set.close();
|
||||
|
||||
var edit = VersionEdit.init(allocator);
|
||||
defer edit.deinit();
|
||||
try edit.addFile(7, 0);
|
||||
try opened.set.logAndApply(&edit, &opened.state);
|
||||
|
||||
try opened.set.file.seekFromEnd(0);
|
||||
try opened.set.file.writeAll(&[_]u8{ 0xde, 0xad, 0xbe });
|
||||
try opened.set.file.sync();
|
||||
}
|
||||
|
||||
var reopened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer reopened.state.deinit();
|
||||
defer reopened.set.close();
|
||||
try std.testing.expectEqual(@as(usize, 1), reopened.state.count());
|
||||
try std.testing.expectEqual(@as(usize, 0), reopened.state.get(7).?);
|
||||
}
|
||||
|
||||
test "logAndApply persists nextFileId and nextSeq counters across reopen" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
{
|
||||
var opened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer opened.state.deinit();
|
||||
defer opened.set.close();
|
||||
try std.testing.expectEqual(@as(?u64, null), opened.set.counters.nextFileId);
|
||||
|
||||
var edit = VersionEdit.init(allocator);
|
||||
defer edit.deinit();
|
||||
try edit.addFile(1, 0);
|
||||
edit.setNextFileId(2);
|
||||
edit.setNextSeq(500);
|
||||
try opened.set.logAndApply(&edit, &opened.state);
|
||||
|
||||
try std.testing.expectEqual(@as(?u64, 2), opened.set.counters.nextFileId);
|
||||
try std.testing.expectEqual(@as(?u64, 500), opened.set.counters.nextSeq);
|
||||
}
|
||||
|
||||
var reopened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer reopened.state.deinit();
|
||||
defer reopened.set.close();
|
||||
try std.testing.expectEqual(@as(?u64, 2), reopened.set.counters.nextFileId);
|
||||
try std.testing.expectEqual(@as(?u64, 500), reopened.set.counters.nextSeq);
|
||||
}
|
||||
|
||||
test "counters survive manifest rotation" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
var lastId: u64 = 0;
|
||||
{
|
||||
var opened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer opened.state.deinit();
|
||||
defer opened.set.close();
|
||||
|
||||
var fileId: u64 = 1;
|
||||
{
|
||||
var edit = VersionEdit.init(allocator);
|
||||
defer edit.deinit();
|
||||
try edit.addFile(fileId, 0);
|
||||
edit.setNextFileId(fileId + 1);
|
||||
edit.setNextSeq(10);
|
||||
try opened.set.logAndApply(&edit, &opened.state);
|
||||
}
|
||||
// Force rotation the same way the existing rotation test does, each
|
||||
// round bumping the seq counter so the final persisted value is
|
||||
// distinguishable from an earlier one if rotation dropped it.
|
||||
var n: usize = 0;
|
||||
while (n < 20) : (n += 1) {
|
||||
var edit = VersionEdit.init(allocator);
|
||||
defer edit.deinit();
|
||||
try edit.deleteFile(fileId);
|
||||
fileId += 1;
|
||||
try edit.addFile(fileId, 0);
|
||||
edit.setNextFileId(fileId + 1);
|
||||
edit.setNextSeq(10 + n + 1);
|
||||
try opened.set.logAndApply(&edit, &opened.state);
|
||||
}
|
||||
lastId = opened.set.manifestId;
|
||||
try std.testing.expectEqual(@as(?u64, fileId + 1), opened.set.counters.nextFileId);
|
||||
try std.testing.expectEqual(@as(?u64, 30), opened.set.counters.nextSeq);
|
||||
}
|
||||
try std.testing.expect(lastId > 1);
|
||||
|
||||
var reopened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer reopened.state.deinit();
|
||||
defer reopened.set.close();
|
||||
try std.testing.expectEqual(@as(?u64, 30), reopened.set.counters.nextSeq);
|
||||
}
|
||||
|
||||
test "missing CURRENT after previous run bootstraps fresh, empty state" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
{
|
||||
var opened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer opened.state.deinit();
|
||||
defer opened.set.close();
|
||||
var edit = VersionEdit.init(allocator);
|
||||
defer edit.deinit();
|
||||
try edit.addFile(1, 0);
|
||||
try opened.set.logAndApply(&edit, &opened.state);
|
||||
}
|
||||
|
||||
try tmp.dir.deleteFile(currentName);
|
||||
|
||||
var reopened = try VersionSet.open(allocator, tmp.dir);
|
||||
defer reopened.state.deinit();
|
||||
defer reopened.set.close();
|
||||
try std.testing.expectEqual(@as(usize, 0), reopened.state.count());
|
||||
}
|
||||
+96
@@ -165,6 +165,102 @@ fn applyFrame(memtable: *Memtable, body: []const u8) !u64 {
|
||||
return seq;
|
||||
}
|
||||
|
||||
const genFileName = "wal.gen";
|
||||
const genTmpName = "wal.gen.tmp";
|
||||
const genMagic = "CGEN";
|
||||
const genRecordLen = 4 + 8 + 4;
|
||||
|
||||
pub const GenerationError = error{ BadMagic, Truncated, Corrupt };
|
||||
|
||||
/// Reads the WAL's current generation counter from a db directory, next to
|
||||
/// wal.log itself. Returns 0 (not an error) if the generation file has never
|
||||
/// been written - a db that has never truncated its WAL is, by definition,
|
||||
/// still on generation 0, the same value any tailer starts at.
|
||||
pub fn currentGeneration(dir: std.fs.Dir) !u64 {
|
||||
var buf: [genRecordLen]u8 = undefined;
|
||||
const file = dir.openFile(genFileName, .{}) catch |err| switch (err) {
|
||||
error.FileNotFound => return 0,
|
||||
else => return err,
|
||||
};
|
||||
defer file.close();
|
||||
const n = try file.readAll(&buf);
|
||||
if (n != genRecordLen) return GenerationError.Truncated;
|
||||
if (!std.mem.eql(u8, buf[0..4], genMagic)) return GenerationError.BadMagic;
|
||||
const gen = encode.readU64(&buf, 4);
|
||||
const storedCrc = encode.readU32(&buf, 12);
|
||||
if (encode.crc(buf[0..12]) != storedCrc) return GenerationError.Corrupt;
|
||||
return gen;
|
||||
}
|
||||
|
||||
/// Bumps and atomically persists the WAL generation counter. Called once per
|
||||
/// truncation (see reset(), invoked from Db.flush()) so a tailer that reads
|
||||
/// this stamp can tell "same WAL contents since I last looked" from "primary
|
||||
/// truncated, and may have already regrown past whatever byte offset I had"
|
||||
/// from an explicit marker instead of a byte-size heuristic - a pure
|
||||
/// endPos-vs-offset comparison is fooled exactly when the post-truncation
|
||||
/// file reaches or exceeds the old offset before the next poll.
|
||||
pub fn bumpGeneration(dir: std.fs.Dir) !u64 {
|
||||
const gen = (try currentGeneration(dir)) + 1;
|
||||
var buf: [genRecordLen]u8 = undefined;
|
||||
@memcpy(buf[0..4], genMagic);
|
||||
std.mem.writeInt(u64, buf[4..12], gen, encode.endian);
|
||||
std.mem.writeInt(u32, buf[12..16], encode.crc(buf[0..12]), encode.endian);
|
||||
{
|
||||
const file = try dir.createFile(genTmpName, .{ .truncate = true });
|
||||
defer file.close();
|
||||
try file.writeAll(&buf);
|
||||
try file.sync();
|
||||
}
|
||||
try dir.rename(genTmpName, genFileName);
|
||||
return gen;
|
||||
}
|
||||
|
||||
test "currentGeneration is 0 before any bump" {
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
try std.testing.expectEqual(@as(u64, 0), try currentGeneration(tmp.dir));
|
||||
}
|
||||
|
||||
test "bumpGeneration increments and persists across reopen" {
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
try std.testing.expectEqual(@as(u64, 1), try bumpGeneration(tmp.dir));
|
||||
try std.testing.expectEqual(@as(u64, 2), try bumpGeneration(tmp.dir));
|
||||
try std.testing.expectEqual(@as(u64, 3), try bumpGeneration(tmp.dir));
|
||||
try std.testing.expectEqual(@as(u64, 3), try currentGeneration(tmp.dir));
|
||||
}
|
||||
|
||||
test "currentGeneration rejects a corrupted generation file" {
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
_ = try bumpGeneration(tmp.dir);
|
||||
|
||||
const file = try tmp.dir.openFile(genFileName, .{ .mode = .read_write });
|
||||
defer file.close();
|
||||
try file.seekTo(4);
|
||||
try file.writeAll(&[_]u8{ 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF });
|
||||
|
||||
try std.testing.expectError(GenerationError.Corrupt, currentGeneration(tmp.dir));
|
||||
}
|
||||
|
||||
test "Wal.reset does not itself touch the generation counter" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
var wal = try Wal.open(allocator, tmp.dir, "wal.log");
|
||||
defer wal.close();
|
||||
try wal.append(.value, 1, "k", "v");
|
||||
try wal.reset();
|
||||
|
||||
// reset() only truncates the WAL file itself - generation bumping is a
|
||||
// separate, explicit call made by Db.flush() at the same call site, kept
|
||||
// out of Wal so tests exercising reset() in isolation (elsewhere in this
|
||||
// file) do not need a directory-wide side effect they never asked for.
|
||||
try std.testing.expectEqual(@as(u64, 0), try currentGeneration(tmp.dir));
|
||||
}
|
||||
|
||||
test "wal append and replay" {
|
||||
const allocator = std.testing.allocator;
|
||||
var tmp = std.testing.tmpDir(.{});
|
||||
|
||||
+2222
File diff suppressed because it is too large.
Load diff
Reference in new issue
Block a user