dualfroz/cobalt
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
025f1708c6956e67677340428000f70ba3a95c5a
cobalt/src/collections.zig
T

1108 lines
40 KiB
Zig
Raw Blame History

const std = @import("std");
const dbmod = @import("db.zig");
const row = @import("row.zig");
const txnmod = @import("txn.zig");
const Allocator = std.mem.Allocator;
const Db = dbmod.Db;
const Transaction = txnmod.Transaction;
/// Redis-like typed collections layered over the flat KV keyspace, sharing
/// one Db with tables/indexes (row.zig/index.zig use tags 'R'/'I'/'M'; every
/// key here starts with 'X' so the two namespaces never collide):
/// "X" "C\0" name -> raw i64 (counter)
/// "X" "H\0" name "\0" field -> raw bytes (hash field)
/// "X" "S\0" name "\0" member -> empty marker (set member)
/// "X" "L\0" name "\0M" -> raw {headSeq i64, tailSeq i64} (list meta)
/// "X" "L\0" name "\0E\0" orderedInt(seq) -> raw bytes (list element)
/// "X" "Z\0" name "\0M\0" member -> raw f64 (zset member -> score)
/// "X" "Z\0" name "\0S\0" orderedFloat(score) orderedText(member) -> empty marker (zset score index)
pub const nsTag: u8 = 'X';
const sep: u8 = 0x00;
const maxConflictRetries: usize = 64;
fn appendSegment(out: *std.ArrayList(u8), segment: []const u8) !void {
try out.appendSlice(segment);
try out.append(sep);
}
fn putI64(out: *std.ArrayList(u8), v: i64) !void {
var b: [8]u8 = undefined;
std.mem.writeInt(i64, &b, v, .big);
try out.appendSlice(&b);
}
fn readI64(data: []const u8) !i64 {
if (data.len != 8) return error.Malformed;
return std.mem.readInt(i64, data[0..8], .big);
}
fn putF64(out: *std.ArrayList(u8), v: f64) !void {
var b: [8]u8 = undefined;
std.mem.writeInt(u64, &b, @bitCast(v), .big);
try out.appendSlice(&b);
}
fn readF64(data: []const u8) !f64 {
if (data.len != 8) return error.Malformed;
return @bitCast(std.mem.readInt(u64, data[0..8], .big));
}
/// Every atomic collection mutator below retries up to `maxConflictRetries`
/// times on error.Conflict, re-reading current state fresh through `txn.get`
/// on each attempt rather than replaying a stale value. `error.TooManyConflicts`
/// is returned only if the retry budget is exhausted.
/// Identifies which typed collection a name belongs to - used by collttl.zig
/// to expire a WHOLE collection (every physical key under its namespace),
/// since unlike a flat key a collection has no single Db key of its own
/// (except Counter, see `exists`/`drop` below).
pub const Kind = enum {
counter,
hash,
set,
list,
zset,
pub fn parse(s: []const u8) ?Kind {
if (std.mem.eql(u8, s, "counter")) return .counter;
if (std.mem.eql(u8, s, "hash")) return .hash;
if (std.mem.eql(u8, s, "set")) return .set;
if (std.mem.eql(u8, s, "list")) return .list;
if (std.mem.eql(u8, s, "zset")) return .zset;
return null;
}
pub fn label(self: Kind) []const u8 {
return switch (self) {
.counter => "counter",
.hash => "hash",
.set => "set",
.list => "list",
.zset => "zset",
};
}
pub fn tagByte(self: Kind) u8 {
return switch (self) {
.counter => 'C',
.hash => 'H',
.set => 'S',
.list => 'L',
.zset => 'Z',
};
}
/// Inverse of `tagByte`, null on any other byte. Used to decode a
/// `Kind` back out of a wire payload (`protocol.zig`'s `coll_ttl_op`)
/// or a marker key (`collttl.zig`'s `sweepExpired`), which previously
/// duplicated this exact switch inline.
pub fn fromTagByte(b: u8) ?Kind {
return switch (b) {
'C' => .counter,
'H' => .hash,
'S' => .set,
'L' => .list,
'Z' => .zset,
else => null,
};
}
};
/// Prefix common to every physical key a Hash/Set/List/ZSet named `name`
/// owns - each of those four kinds always segments as
/// nsTag+kindTag+sep+name+sep+<rest>, so this single prefix reaches every
/// field/member/element/score-entry at once. Counter is the one kind NOT
/// covered here: its single key has no trailing separator after `name`
/// (see Counter.key), so a same-named counter is handled separately in
/// `exists`/`drop` below rather than through a prefix scan.
fn genericPrefix(allocator: Allocator, k: Kind, name: []const u8) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, &[_]u8{k.tagByte()});
try appendSegment(&out, name);
return out.toOwnedSlice();
}
/// True if `name` currently has any physical presence under `k`'s namespace.
/// A counter counts as existing once `set`/`incrBy` has stored a value for
/// it, even 0 - that stored 0 is real data, not the "absent" default
/// Counter.get returns for a name nothing ever touched.
pub fn collectionExists(db: *Db, allocator: Allocator, k: Kind, name: []const u8) !bool {
if (k == .counter) {
const ck = try Counter.key(allocator, name);
defer allocator.free(ck);
const raw = try db.get(allocator, ck) orelse return false;
allocator.free(raw);
return true;
}
const p = try genericPrefix(allocator, k, name);
defer allocator.free(p);
var it = try db.scanPrefix(allocator, p);
defer it.deinit();
return (try it.next()) != null;
}
/// Deletes every physical key belonging to `name` under `k`'s namespace -
/// every field of a hash, the meta record and every element of a list, both
/// the member row and score-index row of every zset entry. Returns the
/// count of physical keys removed. This is what makes a collection-level
/// TTL (collttl.zig) actually mean something: expiring a collection has to
/// remove ALL of it, not just one key the way ttl.zig's flat-key expiry does.
pub fn dropCollection(db: *Db, allocator: Allocator, k: Kind, name: []const u8) !usize {
if (k == .counter) {
const ck = try Counter.key(allocator, name);
defer allocator.free(ck);
const raw = try db.get(allocator, ck) orelse return 0;
allocator.free(raw);
try db.delete(ck);
return 1;
}
const p = try genericPrefix(allocator, k, name);
defer allocator.free(p);
var keys = std.ArrayList([]u8).init(allocator);
defer {
for (keys.items) |kk| allocator.free(kk);
keys.deinit();
}
{
// Collect first, delete after: db.delete mid-scan would mutate the
// memtable structure a live RangeIterator's cursors are threaded
// through (same reasoning as ttl.zig's sweepExpired).
var it = try db.scanPrefix(allocator, p);
defer it.deinit();
while (try it.next()) |entry| {
try keys.append(try allocator.dupe(u8, entry.key));
}
}
for (keys.items) |kk| try db.delete(kk);
return keys.items.len;
}
fn prefixUpperBound(allocator: Allocator, prefix: []const u8) !?[]u8 {
var end = prefix.len;
while (end > 0 and prefix[end - 1] == 0xff) end -= 1;
if (end == 0) return null;
const upper = try allocator.alloc(u8, end);
@memcpy(upper, prefix[0..end]);
upper[end - 1] += 1;
return upper;
}
// ---------------------------------------------------------------------
// Counter
// ---------------------------------------------------------------------
pub const Counter = struct {
fn key(allocator: Allocator, name: []const u8) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, "C");
try out.appendSlice(name);
return out.toOwnedSlice();
}
pub fn get(db: *Db, allocator: Allocator, name: []const u8) !i64 {
const k = try key(allocator, name);
defer allocator.free(k);
const raw = try db.get(allocator, k) orelse return 0;
defer allocator.free(raw);
return readI64(raw);
}
pub fn set(db: *Db, allocator: Allocator, name: []const u8, value: i64) !void {
const k = try key(allocator, name);
defer allocator.free(k);
var out = std.ArrayList(u8).init(allocator);
defer out.deinit();
try putI64(&out, value);
try db.put(k, out.items);
}
/// Atomically adds `delta` to the counter (creating it at 0 if absent)
/// and returns the new value. Uses an OCC transaction so concurrent
/// increments never lose an update to a lost read-modify-write race.
pub fn incrBy(db: *Db, allocator: Allocator, name: []const u8, delta: i64) !i64 {
const k = try key(allocator, name);
defer allocator.free(k);
var attempt: usize = 0;
while (attempt < maxConflictRetries) : (attempt += 1) {
var txn = Transaction.begin(db);
defer txn.deinit();
const current: i64 = blk: {
const raw = try txn.get(allocator, k) orelse break :blk 0;
defer allocator.free(raw);
break :blk try readI64(raw);
};
const next = current + delta;
var out = std.ArrayList(u8).init(allocator);
defer out.deinit();
try putI64(&out, next);
try txn.put(k, out.items);
txn.commit() catch |err| {
if (err == error.Conflict) continue;
return err;
};
return next;
}
return error.TooManyConflicts;
}
};
// ---------------------------------------------------------------------
// Hash
// ---------------------------------------------------------------------
pub const HashField = struct {
field: []u8,
value: []u8,
pub fn free(self: HashField, allocator: Allocator) void {
allocator.free(self.field);
allocator.free(self.value);
}
};
pub const Hash = struct {
fn prefix(allocator: Allocator, name: []const u8) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, "H");
try appendSegment(&out, name);
return out.toOwnedSlice();
}
fn fieldKey(allocator: Allocator, name: []const u8, field: []const u8) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, "H");
try appendSegment(&out, name);
try out.appendSlice(field);
return out.toOwnedSlice();
}
/// Sets a field, returning true if the field was newly created.
pub fn set(db: *Db, allocator: Allocator, name: []const u8, field: []const u8, value: []const u8) !bool {
const k = try fieldKey(allocator, name, field);
defer allocator.free(k);
const existing = try db.get(allocator, k);
if (existing) |raw| allocator.free(raw);
try db.put(k, value);
return existing == null;
}
pub fn get(db: *Db, allocator: Allocator, name: []const u8, field: []const u8) !?[]u8 {
const k = try fieldKey(allocator, name, field);
defer allocator.free(k);
return db.get(allocator, k);
}
pub fn exists(db: *Db, allocator: Allocator, name: []const u8, field: []const u8) !bool {
const raw = try get(db, allocator, name, field) orelse return false;
allocator.free(raw);
return true;
}
/// Deletes a field, returning true if it existed.
pub fn del(db: *Db, allocator: Allocator, name: []const u8, field: []const u8) !bool {
const existed = try exists(db, allocator, name, field);
if (!existed) return false;
const k = try fieldKey(allocator, name, field);
defer allocator.free(k);
try db.delete(k);
return true;
}
pub fn len(db: *Db, allocator: Allocator, name: []const u8) !usize {
const p = try prefix(allocator, name);
defer allocator.free(p);
var it = try db.scanPrefix(allocator, p);
defer it.deinit();
var count: usize = 0;
while (try it.next()) |_| count += 1;
return count;
}
/// Returns every field/value pair currently stored under `name`. Caller
/// owns the returned slice and each entry's `field`/`value` buffers.
pub fn getAll(db: *Db, allocator: Allocator, name: []const u8) ![]HashField {
const p = try prefix(allocator, name);
defer allocator.free(p);
var results = std.ArrayList(HashField).init(allocator);
errdefer {
for (results.items) |f| f.free(allocator);
results.deinit();
}
var it = try db.scanPrefix(allocator, p);
defer it.deinit();
while (try it.next()) |entry| {
const field = try allocator.dupe(u8, entry.key[p.len..]);
errdefer allocator.free(field);
const value = try allocator.dupe(u8, entry.value);
try results.append(.{ .field = field, .value = value });
}
return results.toOwnedSlice();
}
};
// ---------------------------------------------------------------------
// Set
// ---------------------------------------------------------------------
pub const Set = struct {
fn prefix(allocator: Allocator, name: []const u8) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, "S");
try appendSegment(&out, name);
return out.toOwnedSlice();
}
fn memberKey(allocator: Allocator, name: []const u8, member: []const u8) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, "S");
try appendSegment(&out, name);
try out.appendSlice(member);
return out.toOwnedSlice();
}
/// Returns true if the member was newly added (false if already present).
pub fn add(db: *Db, allocator: Allocator, name: []const u8, member: []const u8) !bool {
if (try isMember(db, allocator, name, member)) return false;
const k = try memberKey(allocator, name, member);
defer allocator.free(k);
try db.put(k, &.{});
return true;
}
/// Returns true if the member was present and has been removed.
pub fn remove(db: *Db, allocator: Allocator, name: []const u8, member: []const u8) !bool {
if (!try isMember(db, allocator, name, member)) return false;
const k = try memberKey(allocator, name, member);
defer allocator.free(k);
try db.delete(k);
return true;
}
pub fn isMember(db: *Db, allocator: Allocator, name: []const u8, member: []const u8) !bool {
const k = try memberKey(allocator, name, member);
defer allocator.free(k);
const raw = try db.get(allocator, k) orelse return false;
allocator.free(raw);
return true;
}
pub fn card(db: *Db, allocator: Allocator, name: []const u8) !usize {
const members_ = try members(db, allocator, name);
defer {
for (members_) |m| allocator.free(m);
allocator.free(members_);
}
return members_.len;
}
/// Returns every member currently in the set. Caller owns the slice and
/// each element.
pub fn members(db: *Db, allocator: Allocator, name: []const u8) ![][]u8 {
const p = try prefix(allocator, name);
defer allocator.free(p);
var results = std.ArrayList([]u8).init(allocator);
errdefer {
for (results.items) |m| allocator.free(m);
results.deinit();
}
var it = try db.scanPrefix(allocator, p);
defer it.deinit();
while (try it.next()) |entry| {
try results.append(try allocator.dupe(u8, entry.key[p.len..]));
}
return results.toOwnedSlice();
}
fn containsSorted(haystack: [][]const u8, needle: []const u8) bool {
for (haystack) |h| {
if (std.mem.eql(u8, h, needle)) return true;
}
return false;
}
/// Members present in every one of `names`. Caller owns the result.
pub fn intersect(db: *Db, allocator: Allocator, names: []const []const u8) ![][]u8 {
var results = std.ArrayList([]u8).init(allocator);
errdefer {
for (results.items) |m| allocator.free(m);
results.deinit();
}
if (names.len == 0) return results.toOwnedSlice();
const first = try members(db, allocator, names[0]);
defer {
for (first) |m| allocator.free(m);
allocator.free(first);
}
outer: for (first) |candidate| {
for (names[1..]) |other| {
if (!try isMember(db, allocator, other, candidate)) continue :outer;
}
try results.append(try allocator.dupe(u8, candidate));
}
return results.toOwnedSlice();
}
/// Members present in any of `names`, deduplicated. Caller owns the result.
pub fn unionOf(db: *Db, allocator: Allocator, names: []const []const u8) ![][]u8 {
var results = std.ArrayList([]u8).init(allocator);
errdefer {
for (results.items) |m| allocator.free(m);
results.deinit();
}
for (names) |name| {
const ms = try members(db, allocator, name);
defer {
for (ms) |m| allocator.free(m);
allocator.free(ms);
}
for (ms) |m| {
var seen = false;
for (results.items) |existing| {
if (std.mem.eql(u8, existing, m)) {
seen = true;
break;
}
}
if (!seen) try results.append(try allocator.dupe(u8, m));
}
}
return results.toOwnedSlice();
}
/// Members in `name` that are not in any set of `others`. Caller owns the result.
pub fn difference(db: *Db, allocator: Allocator, name: []const u8, others: []const []const u8) ![][]u8 {
var results = std.ArrayList([]u8).init(allocator);
errdefer {
for (results.items) |m| allocator.free(m);
results.deinit();
}
const base = try members(db, allocator, name);
defer {
for (base) |m| allocator.free(m);
allocator.free(base);
}
outer: for (base) |candidate| {
for (others) |other| {
if (try isMember(db, allocator, other, candidate)) continue :outer;
}
try results.append(try allocator.dupe(u8, candidate));
}
return results.toOwnedSlice();
}
};
// ---------------------------------------------------------------------
// List (double-ended queue keyed by a monotonic sequence range)
// ---------------------------------------------------------------------
const ListMeta = struct {
// Elements occupy sequence numbers (head+1 .. tail-1) inclusive; an
// empty list has tail == head + 1. Pushing left decrements head and
// writes at the new head; pushing right increments tail and writes at
// the slot just below the new tail.
head: i64,
tail: i64,
};
pub const List = struct {
fn metaKey(allocator: Allocator, name: []const u8) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, "L");
try appendSegment(&out, name);
try out.append('M');
return out.toOwnedSlice();
}
fn elementPrefix(allocator: Allocator, name: []const u8) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, "L");
try appendSegment(&out, name);
try appendSegment(&out, "E");
return out.toOwnedSlice();
}
fn elementKey(allocator: Allocator, name: []const u8, seq: i64) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, "L");
try appendSegment(&out, name);
try appendSegment(&out, "E");
try row.appendOrdered(&out, .{ .int = seq });
return out.toOwnedSlice();
}
fn readMeta(txn: *Transaction, allocator: Allocator, k: []const u8) !ListMeta {
const raw = try txn.get(allocator, k) orelse return ListMeta{ .head = 0, .tail = 1 };
defer allocator.free(raw);
if (raw.len != 16) return error.Malformed;
return .{
.head = std.mem.readInt(i64, raw[0..8], .big),
.tail = std.mem.readInt(i64, raw[8..16], .big),
};
}
fn writeMeta(txn: *Transaction, k: []const u8, meta: ListMeta) !void {
var buf: [16]u8 = undefined;
std.mem.writeInt(i64, buf[0..8], meta.head, .big);
std.mem.writeInt(i64, buf[8..16], meta.tail, .big);
try txn.put(k, &buf);
}
pub fn pushLeft(db: *Db, allocator: Allocator, name: []const u8, value: []const u8) !void {
const mk = try metaKey(allocator, name);
defer allocator.free(mk);
var attempt: usize = 0;
while (attempt < maxConflictRetries) : (attempt += 1) {
var txn = Transaction.begin(db);
defer txn.deinit();
var meta = try readMeta(&txn, allocator, mk);
const seq = meta.head;
meta.head -= 1;
const ek = try elementKey(allocator, name, seq);
defer allocator.free(ek);
try txn.put(ek, value);
try writeMeta(&txn, mk, meta);
txn.commit() catch |err| {
if (err == error.Conflict) continue;
return err;
};
return;
}
return error.TooManyConflicts;
}
pub fn pushRight(db: *Db, allocator: Allocator, name: []const u8, value: []const u8) !void {
const mk = try metaKey(allocator, name);
defer allocator.free(mk);
var attempt: usize = 0;
while (attempt < maxConflictRetries) : (attempt += 1) {
var txn = Transaction.begin(db);
defer txn.deinit();
var meta = try readMeta(&txn, allocator, mk);
const seq = meta.tail;
meta.tail += 1;
const ek = try elementKey(allocator, name, seq);
defer allocator.free(ek);
try txn.put(ek, value);
try writeMeta(&txn, mk, meta);
txn.commit() catch |err| {
if (err == error.Conflict) continue;
return err;
};
return;
}
return error.TooManyConflicts;
}
pub fn popLeft(db: *Db, allocator: Allocator, name: []const u8) !?[]u8 {
const mk = try metaKey(allocator, name);
defer allocator.free(mk);
var attempt: usize = 0;
while (attempt < maxConflictRetries) : (attempt += 1) {
var txn = Transaction.begin(db);
defer txn.deinit();
var meta = try readMeta(&txn, allocator, mk);
if (meta.tail - meta.head <= 1) return null;
const seq = meta.head + 1;
meta.head += 1;
const ek = try elementKey(allocator, name, seq);
defer allocator.free(ek);
const value = try txn.get(allocator, ek) orelse return error.Malformed;
errdefer allocator.free(value);
try txn.delete(ek);
try writeMeta(&txn, mk, meta);
txn.commit() catch |err| {
allocator.free(value);
if (err == error.Conflict) continue;
return err;
};
return value;
}
return error.TooManyConflicts;
}
pub fn popRight(db: *Db, allocator: Allocator, name: []const u8) !?[]u8 {
const mk = try metaKey(allocator, name);
defer allocator.free(mk);
var attempt: usize = 0;
while (attempt < maxConflictRetries) : (attempt += 1) {
var txn = Transaction.begin(db);
defer txn.deinit();
var meta = try readMeta(&txn, allocator, mk);
if (meta.tail - meta.head <= 1) return null;
const seq = meta.tail - 1;
meta.tail -= 1;
const ek = try elementKey(allocator, name, seq);
defer allocator.free(ek);
const value = try txn.get(allocator, ek) orelse return error.Malformed;
errdefer allocator.free(value);
try txn.delete(ek);
try writeMeta(&txn, mk, meta);
txn.commit() catch |err| {
allocator.free(value);
if (err == error.Conflict) continue;
return err;
};
return value;
}
return error.TooManyConflicts;
}
pub fn len(db: *Db, allocator: Allocator, name: []const u8) !usize {
const mk = try metaKey(allocator, name);
defer allocator.free(mk);
const raw = try db.get(allocator, mk) orelse return 0;
defer allocator.free(raw);
if (raw.len != 16) return error.Malformed;
const head = std.mem.readInt(i64, raw[0..8], .big);
const tail = std.mem.readInt(i64, raw[8..16], .big);
return @intCast(tail - head - 1);
}
/// Redis-style inclusive range with negative indices counting from the
/// end (-1 is the last element). Out-of-range indices are clamped;
/// start > stop after clamping yields an empty result. Caller owns the
/// returned slice and each element.
pub fn range(db: *Db, allocator: Allocator, name: []const u8, start: i64, stop: i64) ![][]u8 {
const total: i64 = @intCast(try len(db, allocator, name));
var results = std.ArrayList([]u8).init(allocator);
errdefer {
for (results.items) |m| allocator.free(m);
results.deinit();
}
if (total == 0) return results.toOwnedSlice();
var lo = if (start < 0) total + start else start;
var hi = if (stop < 0) total + stop else stop;
if (lo < 0) lo = 0;
if (hi > total - 1) hi = total - 1;
if (lo > hi) return results.toOwnedSlice();
const p = try elementPrefix(allocator, name);
defer allocator.free(p);
var it = try db.scanPrefix(allocator, p);
defer it.deinit();
var idx: i64 = 0;
while (try it.next()) |entry| {
if (idx >= lo and idx <= hi) {
try results.append(try allocator.dupe(u8, entry.value));
}
idx += 1;
}
return results.toOwnedSlice();
}
};
// ---------------------------------------------------------------------
// ZSet (sorted set: member -> score, plus a score-ordered index for
// range/rank queries)
// ---------------------------------------------------------------------
pub const ZMember = struct {
member: []u8,
score: f64,
pub fn free(self: ZMember, allocator: Allocator) void {
allocator.free(self.member);
}
};
pub const ZSet = struct {
fn memberKey(allocator: Allocator, name: []const u8, member: []const u8) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, "Z");
try appendSegment(&out, name);
try appendSegment(&out, "M");
try out.appendSlice(member);
return out.toOwnedSlice();
}
fn scorePrefix(allocator: Allocator, name: []const u8) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, "Z");
try appendSegment(&out, name);
try appendSegment(&out, "S");
return out.toOwnedSlice();
}
fn scoreEntryKey(allocator: Allocator, name: []const u8, scoreVal: f64, member: []const u8) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, "Z");
try appendSegment(&out, name);
try appendSegment(&out, "S");
try row.appendOrdered(&out, .{ .float = scoreVal });
try row.appendOrdered(&out, .{ .text = member });
return out.toOwnedSlice();
}
/// Score-index key truncated to just the ordered score (no member
/// suffix), used as a scan bound. Every real entry is a strict byte
/// extension of this (the member part is never empty), so a key equal to
/// `scoreBound(s)` never itself denotes a stored member.
fn scoreBound(allocator: Allocator, name: []const u8, scoreVal: f64) ![]u8 {
var out = std.ArrayList(u8).init(allocator);
errdefer out.deinit();
try out.append(nsTag);
try appendSegment(&out, "Z");
try appendSegment(&out, name);
try appendSegment(&out, "S");
try row.appendOrdered(&out, .{ .float = scoreVal });
return out.toOwnedSlice();
}
/// Upserts `member`'s score, atomically re-pointing the score index if
/// the member already existed with a different score.
pub fn add(db: *Db, allocator: Allocator, name: []const u8, member: []const u8, scoreVal: f64) !void {
const mk = try memberKey(allocator, name, member);
defer allocator.free(mk);
var attempt: usize = 0;
while (attempt < maxConflictRetries) : (attempt += 1) {
var txn = Transaction.begin(db);
defer txn.deinit();
const old = try txn.get(allocator, mk);
if (old) |raw| {
defer allocator.free(raw);
const oldScore = try readF64(raw);
if (oldScore == scoreVal) return;
const oldEntryKey = try scoreEntryKey(allocator, name, oldScore, member);
defer allocator.free(oldEntryKey);
try txn.delete(oldEntryKey);
}
var scoreBuf = std.ArrayList(u8).init(allocator);
defer scoreBuf.deinit();
try putF64(&scoreBuf, scoreVal);
try txn.put(mk, scoreBuf.items);
const entryKey = try scoreEntryKey(allocator, name, scoreVal, member);
defer allocator.free(entryKey);
try txn.put(entryKey, &.{});
txn.commit() catch |err| {
if (err == error.Conflict) continue;
return err;
};
return;
}
return error.TooManyConflicts;
}
pub fn score(db: *Db, allocator: Allocator, name: []const u8, member: []const u8) !?f64 {
const mk = try memberKey(allocator, name, member);
defer allocator.free(mk);
const raw = try db.get(allocator, mk) orelse return null;
defer allocator.free(raw);
return try readF64(raw);
}
pub fn incrBy(db: *Db, allocator: Allocator, name: []const u8, member: []const u8, delta: f64) !f64 {
const mk = try memberKey(allocator, name, member);
defer allocator.free(mk);
var attempt: usize = 0;
while (attempt < maxConflictRetries) : (attempt += 1) {
var txn = Transaction.begin(db);
defer txn.deinit();
const oldRaw = try txn.get(allocator, mk);
const oldScore: f64 = if (oldRaw) |raw| blk: {
defer allocator.free(raw);
break :blk try readF64(raw);
} else 0.0;
const newScore = oldScore + delta;
if (oldRaw != null) {
const oldEntryKey = try scoreEntryKey(allocator, name, oldScore, member);
defer allocator.free(oldEntryKey);
try txn.delete(oldEntryKey);
}
var scoreBuf = std.ArrayList(u8).init(allocator);
defer scoreBuf.deinit();
try putF64(&scoreBuf, newScore);
try txn.put(mk, scoreBuf.items);
const entryKey = try scoreEntryKey(allocator, name, newScore, member);
defer allocator.free(entryKey);
try txn.put(entryKey, &.{});
txn.commit() catch |err| {
if (err == error.Conflict) continue;
return err;
};
return newScore;
}
return error.TooManyConflicts;
}
/// Returns true if the member was present and has been removed.
pub fn remove(db: *Db, allocator: Allocator, name: []const u8, member: []const u8) !bool {
const mk = try memberKey(allocator, name, member);
defer allocator.free(mk);
var attempt: usize = 0;
while (attempt < maxConflictRetries) : (attempt += 1) {
var txn = Transaction.begin(db);
defer txn.deinit();
const raw = try txn.get(allocator, mk) orelse return false;
defer allocator.free(raw);
const existingScore = try readF64(raw);
try txn.delete(mk);
const entryKey = try scoreEntryKey(allocator, name, existingScore, member);
defer allocator.free(entryKey);
try txn.delete(entryKey);
txn.commit() catch |err| {
if (err == error.Conflict) continue;
return err;
};
return true;
}
return error.TooManyConflicts;
}
pub fn card(db: *Db, allocator: Allocator, name: []const u8) !usize {
const p = try scorePrefix(allocator, name);
defer allocator.free(p);
var it = try db.scanPrefix(allocator, p);
defer it.deinit();
var count: usize = 0;
while (try it.next()) |_| count += 1;
return count;
}
fn decodeScoreEntry(allocator: Allocator, name: []const u8, key: []const u8) !ZMember {
const p = try scorePrefix(allocator, name);
defer allocator.free(p);
if (key.len <= p.len or !std.mem.startsWith(u8, key, p)) return error.Malformed;
const rest = key[p.len..];
// rest = [tag u8][8-byte ordered score][tag u8][escaped member][0x00 0x00]
if (rest.len < 9 or rest[0] != @intFromEnum(row.ValueTag.float)) return error.Malformed;
var bits = std.mem.readInt(u64, rest[1..9], .big);
if (bits & 0x8000000000000000 != 0) {
bits &= 0x7fffffffffffffff;
} else {
bits = ~bits;
}
const decodedScore: f64 = @bitCast(bits);
const memberPart = rest[9..];
if (memberPart.len < 3 or memberPart[0] != @intFromEnum(row.ValueTag.text)) return error.Malformed;
var unescaped = std.ArrayList(u8).init(allocator);
errdefer unescaped.deinit();
var i: usize = 1;
while (i < memberPart.len) {
if (memberPart[i] == 0x00) {
if (i + 1 >= memberPart.len) return error.Malformed;
if (memberPart[i + 1] == 0x00) break;
if (memberPart[i + 1] == 0xFF) {
try unescaped.append(0x00);
i += 2;
continue;
}
return error.Malformed;
}
try unescaped.append(memberPart[i]);
i += 1;
}
return .{ .member = try unescaped.toOwnedSlice(), .score = decodedScore };
}
/// Ascending-by-score range selected by rank (0-based, inclusive on both
/// ends), Redis-style: negative indices count from the end. Caller owns
/// the returned slice and each member's buffer.
pub fn rangeByRank(db: *Db, allocator: Allocator, name: []const u8, start: i64, stop: i64) ![]ZMember {
const total: i64 = @intCast(try card(db, allocator, name));
var results = std.ArrayList(ZMember).init(allocator);
errdefer {
for (results.items) |m| m.free(allocator);
results.deinit();
}
if (total == 0) return results.toOwnedSlice();
var lo = if (start < 0) total + start else start;
var hi = if (stop < 0) total + stop else stop;
if (lo < 0) lo = 0;
if (hi > total - 1) hi = total - 1;
if (lo > hi) return results.toOwnedSlice();
const p = try scorePrefix(allocator, name);
defer allocator.free(p);
var it = try db.scanPrefix(allocator, p);
defer it.deinit();
var idx: i64 = 0;
while (try it.next()) |entry| {
if (idx >= lo and idx <= hi) {
try results.append(try decodeScoreEntry(allocator, name, entry.key));
}
idx += 1;
}
return results.toOwnedSlice();
}
/// All members whose score falls within [min, max], found via a real
/// bounded scan over the score index (not a full-collection scan): the
/// start bound is `scoreBound(min)` and the end bound is the byte
/// immediately after `scoreBound(max)`, which is enough to include every
/// entry whose score equals `max` since real entries strictly extend
/// their score bound (see `scoreBound`'s doc comment). Caller owns the
/// returned slice and each member's buffer.
pub fn rangeByScore(db: *Db, allocator: Allocator, name: []const u8, min: f64, max: f64) ![]ZMember {
var results = std.ArrayList(ZMember).init(allocator);
errdefer {
for (results.items) |m| m.free(allocator);
results.deinit();
}
const startBound = try scoreBound(allocator, name, min);
defer allocator.free(startBound);
const maxBound = try scoreBound(allocator, name, max);
defer allocator.free(maxBound);
const endBound = try prefixUpperBound(allocator, maxBound) orelse try allocator.dupe(u8, maxBound);
defer allocator.free(endBound);
var it = try db.scanRange(allocator, .{
.start = startBound,
.end = endBound,
.startInclusive = true,
.endInclusive = false,
});
defer it.deinit();
while (try it.next()) |entry| {
try results.append(try decodeScoreEntry(allocator, name, entry.key));
}
return results.toOwnedSlice();
}
/// 0-based ascending-by-score rank of `member`, or null if absent.
pub fn rank(db: *Db, allocator: Allocator, name: []const u8, member: []const u8) !?usize {
const target = try score(db, allocator, name, member) orelse return null;
const p = try scorePrefix(allocator, name);
defer allocator.free(p);
var it = try db.scanPrefix(allocator, p);
defer it.deinit();
var idx: usize = 0;
while (try it.next()) |entry| {
const decoded = try decodeScoreEntry(allocator, name, entry.key);
defer decoded.free(allocator);
if (decoded.score == target and std.mem.eql(u8, decoded.member, member)) return idx;
idx += 1;
}
return null;
}
};
// ---------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------
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 "counter incrBy accumulates and get defaults to zero" {
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(@as(i64, 0), try Counter.get(&db, allocator, "hits"));
try std.testing.expectEqual(@as(i64, 5), try Counter.incrBy(&db, allocator, "hits", 5));
try std.testing.expectEqual(@as(i64, 3), try Counter.incrBy(&db, allocator, "hits", -2));
try std.testing.expectEqual(@as(i64, 3), try Counter.get(&db, allocator, "hits"));
try Counter.set(&db, allocator, "hits", 100);
try std.testing.expectEqual(@as(i64, 100), try Counter.get(&db, allocator, "hits"));
}
test "counter incrBy rejects a real external conflict" {
const allocator = std.testing.allocator;
var tmp = std.testing.tmpDir(.{});
defer tmp.cleanup();
var db = try openTemp(&tmp, allocator);
defer db.close();
// incrBy retries internally, so exercise the underlying primitive
// directly: begin a txn, let an external writer land in between, then
// confirm the stale commit is rejected rather than silently overwriting.
const k = try Counter.key(allocator, "race");
defer allocator.free(k);
var txn = Transaction.begin(&db);
defer txn.deinit();
_ = try txn.get(allocator, k);
try db.put(k, "external");
try txn.put(k, "mine");
try std.testing.expectError(error.Conflict, txn.commit());
}
Reference in new issue View Git Blame Copy Permalink