tiny.sql.store

1pub const store = @import("store.zig");

1const std = @import("std");2const simd = @import("simd");3const trace = @import("trace.zig");45const Bytes = simd.ScalableTag(u8);67const Allocator = std.mem.Allocator;89const StoreError = error{10    InvalidRange,11    TransactionClosed,12    TransactionConflict,13};1415pub const Error = Allocator.Error || StoreError;1617pub const Entry = struct {18    key: []const u8,19    value: []const u8,20    generation: u64,21};2223const Version = struct {24    key: []u8,25    value: []u8,26    generation: u64,27    deleted: bool = false,28};2930const Mutation = struct {31    key: []u8,32    value: []u8 = &.{},33    deleted: bool = false,3435    fn deinit(self: *Mutation, allocator: Allocator) void {36        freeOwned(allocator, self.key);37        if (!self.deleted) freeOwned(allocator, self.value);38        self.* = .{ .key = &.{}, .value = &.{}, .deleted = true };39    }40};4142pub const Store = struct {43    allocator: Allocator,44    versions: std.ArrayList(Version) = .empty,45    generation: u64 = 0,4647    pub fn init(allocator: Allocator) Store {48        return .{ .allocator = allocator };49    }5051    pub fn deinit(self: *Store) void {52        for (self.versions.items) |*version| {53            freeOwned(self.allocator, version.key);54            if (!version.deleted) freeOwned(self.allocator, version.value);55        }56        self.versions.deinit(self.allocator);57        self.* = .{ .allocator = self.allocator };58    }5960    pub fn currentGeneration(self: *const Store) u64 {61        return self.generation;62    }6364    pub fn beginRead(self: *const Store) Snapshot {65        return .{ .store = self, .generation = self.generation };66    }6768    pub fn beginWrite(self: *Store) Write {69        return .{ .store = self, .snapshot_generation = self.generation };70    }7172    pub fn get(self: *const Store, key: []const u8) ?[]const u8 {73        const phase = trace.scope("store.get");74        defer phase.end();75        return valueFromIndex(self, latestVisibleIndex(self, key, self.generation));76    }7778    pub fn getAt(self: *const Store, key: []const u8, generation: u64) ?[]const u8 {79        return valueFromIndex(self, latestVisibleIndex(self, key, generation));80    }8182    pub fn range(self: *const Store, start: ?[]const u8, end: ?[]const u8) Error!Range {83        return rangeAt(self, start, end, self.generation);84    }8586    pub fn rangeAt(self: *const Store, start: ?[]const u8, end: ?[]const u8, generation: u64) Error!Range {87        if (start) |lower| {88            if (end) |upper| {89                if (simd.order(Bytes, lower, upper) == .gt) return error.InvalidRange;90            }91        }92        return .{93            .store = self,94            .start = start,95            .end = end,96            .generation = generation,97            .index = if (start) |lower| lowerBoundKey(self.versions.items, lower) else 0,98        };99    }100101    pub fn compact(self: *Store, oldest_visible_generation: u64) void {102        const phase = trace.scope("store.compact");103        defer phase.end();104105        var read_index: usize = 0;106        var write_index: usize = 0;107        while (read_index < self.versions.items.len) {108            const key = self.versions.items[read_index].key;109            var group_end = read_index;110            var has_oldest_version = false;111            while (group_end < self.versions.items.len and eqlKey(self.versions.items[group_end].key, key)) : (group_end += 1) {112                const version = self.versions.items[group_end];113                if (version.generation == oldest_visible_generation) has_oldest_version = true;114            }115116            var keep_floor: ?usize = null;117            var scan = read_index;118            while (scan < group_end) : (scan += 1) {119                const version = self.versions.items[scan];120                if (version.generation < oldest_visible_generation) {121                    if (keep_floor == null or self.versions.items[keep_floor.?].generation < version.generation) {122                        keep_floor = scan;123                    }124                }125            }126127            scan = read_index;128            while (scan < group_end) : (scan += 1) {129                const keep =130                    (!has_oldest_version and keep_floor != null and scan == keep_floor.?) or131                    self.versions.items[scan].generation >= oldest_visible_generation;132                if (keep) {133                    self.versions.items[write_index] = self.versions.items[scan];134                    write_index += 1;135                } else {136                    freeOwned(self.allocator, self.versions.items[scan].key);137                    if (!self.versions.items[scan].deleted) freeOwned(self.allocator, self.versions.items[scan].value);138                }139            }140141            read_index = group_end;142        }143        self.versions.shrinkRetainingCapacity(write_index);144    }145};146147pub const Snapshot = struct {148    store: *const Store,149    generation: u64,150151    pub fn get(self: Snapshot, key: []const u8) ?[]const u8 {152        const phase = trace.scope("snapshot.get");153        defer phase.end();154        return self.store.getAt(key, self.generation);155    }156157    pub fn range(self: Snapshot, start: ?[]const u8, end: ?[]const u8) Error!Range {158        return self.store.rangeAt(start, end, self.generation);159    }160};161162pub const Write = struct {163    store: *Store,164    snapshot_generation: u64,165    mutations: std.ArrayList(Mutation) = .empty,166    closed: bool = false,167168    pub fn deinit(self: *Write) void {169        for (self.mutations.items) |*mutation| mutation.deinit(self.store.allocator);170        self.mutations.deinit(self.store.allocator);171        self.closed = true;172    }173174    pub fn get(self: *const Write, key: []const u8) Error!?[]const u8 {175        if (self.closed) return error.TransactionClosed;176        if (self.localMutation(key)) |mutation| {177            if (mutation.deleted) return null;178            return mutation.value;179        }180        return self.store.getAt(key, self.snapshot_generation);181    }182183    pub fn put(self: *Write, key: []const u8, value: []const u8) Error!void {184        if (self.closed) return error.TransactionClosed;185        const phase = trace.scope("write.put");186        defer phase.end();187        try self.setMutation(key, value, false);188    }189190    pub fn delete(self: *Write, key: []const u8) Error!void {191        if (self.closed) return error.TransactionClosed;192        const phase = trace.scope("write.delete");193        defer phase.end();194        try self.setMutation(key, &.{}, true);195    }196197    pub fn commit(self: *Write) Error!u64 {198        if (self.closed) return error.TransactionClosed;199        const phase = trace.scope("write.commit");200        defer phase.end();201        if (self.mutations.items.len == 0) {202            self.closed = true;203            return self.store.generation;204        }205206        for (self.mutations.items) |mutation| {207            if (latestGenerationForKey(self.store, mutation.key)) |generation| {208                if (generation > self.snapshot_generation) return error.TransactionConflict;209            }210        }211212        const next_generation = self.store.generation + 1;213        try self.store.versions.ensureTotalCapacity(self.store.allocator, self.store.versions.items.len + self.mutations.items.len);214        for (self.mutations.items) |mutation| {215            const insert_index = lowerBoundVersion(self.store.versions.items, mutation.key, next_generation);216            self.store.versions.insertAssumeCapacity(insert_index, .{217                .key = mutation.key,218                .value = mutation.value,219                .generation = next_generation,220                .deleted = mutation.deleted,221            });222        }223        std.debug.assert(versionsAreOrdered(self.store.versions.items));224        self.store.generation = next_generation;225        self.mutations.clearRetainingCapacity();226        self.closed = true;227        trace.progress("write.commit.complete");228        return next_generation;229    }230231    fn setMutation(self: *Write, key: []const u8, value: []const u8, deleted: bool) Error!void {232        if (self.localMutationIndex(key)) |index| {233            var mutation = &self.mutations.items[index];234            if (!mutation.deleted) freeOwned(self.store.allocator, mutation.value);235            mutation.value = &.{};236            if (!deleted) mutation.value = try self.store.allocator.dupe(u8, value);237            mutation.deleted = deleted;238            return;239        }240241        const owned_key = try self.store.allocator.dupe(u8, key);242        errdefer freeOwned(self.store.allocator, owned_key);243        var owned_value: []u8 = &.{};244        if (!deleted) owned_value = try self.store.allocator.dupe(u8, value);245        errdefer if (!deleted) freeOwned(self.store.allocator, owned_value);246        try self.mutations.append(self.store.allocator, .{247            .key = owned_key,248            .value = owned_value,249            .deleted = deleted,250        });251        std.mem.sort(Mutation, self.mutations.items, {}, mutationLessThan);252    }253254    fn localMutation(self: *const Write, key: []const u8) ?Mutation {255        if (self.localMutationIndex(key)) |index| return self.mutations.items[index];256        return null;257    }258259    fn localMutationIndex(self: *const Write, key: []const u8) ?usize {260        var low: usize = 0;261        var high = self.mutations.items.len;262        while (low < high) {263            const mid = low + (high - low) / 2;264            switch (simd.order(Bytes, self.mutations.items[mid].key, key)) {265                .lt => low = mid + 1,266                .eq => return mid,267                .gt => high = mid,268            }269        }270        return null;271    }272};273274pub const Range = struct {275    store: *const Store,276    start: ?[]const u8,277    end: ?[]const u8,278    generation: u64,279    index: usize,280281    pub fn next(self: *Range) ?Entry {282        const phase = trace.scope("range.next");283        defer phase.end();284        while (self.index < self.store.versions.items.len) {285            const key = self.store.versions.items[self.index].key;286            if (self.start) |lower| {287                if (simd.order(Bytes, key, lower) == .lt) {288                    self.skipKey(key);289                    continue;290                }291            }292            if (self.end) |upper| {293                if (simd.order(Bytes, key, upper) != .lt) return null;294            }295296            var best_index: ?usize = null;297            while (self.index < self.store.versions.items.len and eqlKey(self.store.versions.items[self.index].key, key)) : (self.index += 1) {298                const candidate = self.store.versions.items[self.index];299                if (candidate.generation <= self.generation) {300                    if (best_index == null or self.store.versions.items[best_index.?].generation < candidate.generation) {301                        best_index = self.index;302                    }303                }304            }305            if (best_index) |index| {306                const version = self.store.versions.items[index];307                if (!version.deleted) {308                    return .{309                        .key = version.key,310                        .value = version.value,311                        .generation = version.generation,312                    };313                }314            }315        }316        return null;317    }318319    fn skipKey(self: *Range, key: []const u8) void {320        while (self.index < self.store.versions.items.len and eqlKey(self.store.versions.items[self.index].key, key)) {321            self.index += 1;322        }323    }324};325326fn latestVisibleIndex(store: *const Store, key: []const u8, generation: u64) ?usize {327    const start = lowerBoundKey(store.versions.items, key);328    var index = start;329    var best: ?usize = null;330    while (index < store.versions.items.len and eqlKey(store.versions.items[index].key, key)) : (index += 1) {331        const version = store.versions.items[index];332        if (version.generation <= generation) {333            if (best == null or store.versions.items[best.?].generation < version.generation) {334                best = index;335            }336        }337    }338    return best;339}340341fn valueFromIndex(store: *const Store, index: ?usize) ?[]const u8 {342    const found = index orelse return null;343    const version = store.versions.items[found];344    if (version.deleted) return null;345    return version.value;346}347348fn latestGenerationForKey(self: *const Store, key: []const u8) ?u64 {349    const start = lowerBoundKey(self.versions.items, key);350    var index = start;351    var latest: ?u64 = null;352    while (index < self.versions.items.len and eqlKey(self.versions.items[index].key, key)) : (index += 1) {353        latest = self.versions.items[index].generation;354    }355    return latest;356}357358fn lowerBoundKey(versions: []const Version, key: []const u8) usize {359    var low: usize = 0;360    var high = versions.len;361    while (low < high) {362        const mid = low + (high - low) / 2;363        switch (simd.order(Bytes, versions[mid].key, key)) {364            .lt => low = mid + 1,365            .eq, .gt => high = mid,366        }367    }368    return low;369}370371fn lowerBoundVersion(versions: []const Version, key: []const u8, generation: u64) usize {372    var low: usize = 0;373    var high = versions.len;374    while (low < high) {375        const mid = low + (high - low) / 2;376        if (versionPrecedesKeyGeneration(versions[mid], key, generation)) {377            low = mid + 1;378        } else {379            high = mid;380        }381    }382    return low;383}384385fn versionLessThan(_: void, left: Version, right: Version) bool {386    return versionPrecedesKeyGeneration(left, right.key, right.generation);387}388389fn versionPrecedesKeyGeneration(version: Version, key: []const u8, generation: u64) bool {390    return switch (simd.order(Bytes, version.key, key)) {391        .lt => true,392        .gt => false,393        .eq => version.generation < generation,394    };395}396397fn versionsAreOrdered(versions: []const Version) bool {398    if (versions.len < 2) return true;399    var index: usize = 1;400    while (index < versions.len) : (index += 1) {401        if (!versionLessThan({}, versions[index - 1], versions[index])) return false;402    }403    return true;404}405406fn versionCount(store: *const Store, key: []const u8) usize {407    const start = lowerBoundKey(store.versions.items, key);408    var index = start;409    var count: usize = 0;410    while (index < store.versions.items.len and eqlKey(store.versions.items[index].key, key)) : (index += 1) {411        count += 1;412    }413    return count;414}415416fn hasVersion(store: *const Store, key: []const u8, generation: u64) bool {417    const start = lowerBoundKey(store.versions.items, key);418    var index = start;419    while (index < store.versions.items.len and eqlKey(store.versions.items[index].key, key)) : (index += 1) {420        if (store.versions.items[index].generation == generation) return true;421    }422    return false;423}424425fn mutationLessThan(_: void, left: Mutation, right: Mutation) bool {426    return switch (simd.order(Bytes, left.key, right.key)) {427        .lt => true,428        .gt => false,429        .eq => false,430    };431}432433fn eqlKey(left: []const u8, right: []const u8) bool {434    return std.mem.eql(u8, left, right);435}436437fn freeOwned(allocator: Allocator, bytes: []u8) void {438    if (bytes.len != 0) allocator.free(bytes);439}440441fn putCommitted(store: *Store, key: []const u8, value: []const u8) !u64 {442    var tx = store.beginWrite();443    defer tx.deinit();444    try tx.put(key, value);445    return tx.commit();446}447448fn deleteCommitted(store: *Store, key: []const u8) !u64 {449    var tx = store.beginWrite();450    defer tx.deinit();451    try tx.delete(key);452    return tx.commit();453}454455test "committed writes are visible in key order" {456    var store = Store.init(std.testing.allocator);457    defer store.deinit();458459    _ = try putCommitted(&store, "b", "two");460    _ = try putCommitted(&store, "a", "one");461462    try std.testing.expectEqualStrings("one", store.get("a").?);463    try std.testing.expectEqualStrings("two", store.get("b").?);464465    var range = try store.range(null, null);466    const first = range.next().?;467    const second = range.next().?;468    try std.testing.expectEqualStrings("a", first.key);469    try std.testing.expectEqualStrings("b", second.key);470    try std.testing.expect(range.next() == null);471}472473test "snapshots preserve old values across later commits" {474    var store = Store.init(std.testing.allocator);475    defer store.deinit();476477    _ = try putCommitted(&store, "name", "first");478    const snapshot = store.beginRead();479    _ = try putCommitted(&store, "name", "second");480481    try std.testing.expectEqualStrings("first", snapshot.get("name").?);482    try std.testing.expectEqualStrings("second", store.get("name").?);483}484485test "delete creates a tombstone for new readers" {486    var store = Store.init(std.testing.allocator);487    defer store.deinit();488489    _ = try putCommitted(&store, "k", "v");490    const snapshot = store.beginRead();491    _ = try deleteCommitted(&store, "k");492493    try std.testing.expectEqualStrings("v", snapshot.get("k").?);494    try std.testing.expect(store.get("k") == null);495}496497test "write transactions detect same key conflicts" {498    var store = Store.init(std.testing.allocator);499    defer store.deinit();500501    var left = store.beginWrite();502    defer left.deinit();503    var right = store.beginWrite();504    defer right.deinit();505506    try left.put("k", "left");507    _ = try left.commit();508    try right.put("k", "right");509    try std.testing.expectError(error.TransactionConflict, right.commit());510}511512test "write transactions allow disjoint commits from the same snapshot" {513    var store = Store.init(std.testing.allocator);514    defer store.deinit();515516    var left = store.beginWrite();517    defer left.deinit();518    var right = store.beginWrite();519    defer right.deinit();520521    try left.put("a", "left");522    try right.put("b", "right");523    _ = try left.commit();524    _ = try right.commit();525526    try std.testing.expectEqualStrings("left", store.get("a").?);527    try std.testing.expectEqualStrings("right", store.get("b").?);528}529530test "range scan honors half open bounds and tombstones" {531    var store = Store.init(std.testing.allocator);532    defer store.deinit();533534    _ = try putCommitted(&store, "a", "1");535    _ = try putCommitted(&store, "b", "2");536    _ = try putCommitted(&store, "c", "3");537    _ = try deleteCommitted(&store, "b");538539    var range = try store.range("a", "c");540    const first = range.next().?;541    try std.testing.expectEqualStrings("a", first.key);542    try std.testing.expect(range.next() == null);543}544545test "multi key commits preserve canonical version order" {546    var store = Store.init(std.testing.allocator);547    defer store.deinit();548549    var seed = store.beginWrite();550    defer seed.deinit();551    try seed.put("b", "old-b");552    try seed.put("d", "old-d");553    _ = try seed.commit();554555    var tx = store.beginWrite();556    defer tx.deinit();557    try tx.put("c", "new-c");558    try tx.put("a", "new-a");559    try tx.put("b", "new-b");560    _ = try tx.commit();561562    try std.testing.expect(versionsAreOrdered(store.versions.items));563    try std.testing.expectEqual(@as(usize, 5), store.versions.items.len);564    try std.testing.expectEqualStrings("a", store.versions.items[0].key);565    try std.testing.expectEqual(@as(u64, 2), store.versions.items[0].generation);566    try std.testing.expectEqualStrings("b", store.versions.items[1].key);567    try std.testing.expectEqual(@as(u64, 1), store.versions.items[1].generation);568    try std.testing.expectEqualStrings("b", store.versions.items[2].key);569    try std.testing.expectEqual(@as(u64, 2), store.versions.items[2].generation);570    try std.testing.expectEqualStrings("c", store.versions.items[3].key);571    try std.testing.expectEqual(@as(u64, 2), store.versions.items[3].generation);572    try std.testing.expectEqualStrings("d", store.versions.items[4].key);573    try std.testing.expectEqual(@as(u64, 1), store.versions.items[4].generation);574}575576test "compaction drops obsolete floor when oldest generation is present" {577    var store = Store.init(std.testing.allocator);578    defer store.deinit();579580    _ = try putCommitted(&store, "k", "v1");581    const snapshot_generation = try putCommitted(&store, "k", "v2");582    _ = try putCommitted(&store, "k", "v3");583584    store.compact(snapshot_generation);585586    try std.testing.expect(versionsAreOrdered(store.versions.items));587    try std.testing.expectEqual(@as(usize, 2), versionCount(&store, "k"));588    try std.testing.expect(!hasVersion(&store, "k", 1));589    try std.testing.expectEqualStrings("v2", store.getAt("k", snapshot_generation).?);590    try std.testing.expectEqualStrings("v3", store.get("k").?);591}592593test "compaction keeps floor when oldest generation falls between versions" {594    var store = Store.init(std.testing.allocator);595    defer store.deinit();596597    _ = try putCommitted(&store, "k", "v1");598    _ = try putCommitted(&store, "other", "advance");599    _ = try putCommitted(&store, "k", "v3");600601    store.compact(2);602603    try std.testing.expect(versionsAreOrdered(store.versions.items));604    try std.testing.expectEqual(@as(usize, 2), versionCount(&store, "k"));605    try std.testing.expect(hasVersion(&store, "k", 1));606    try std.testing.expect(hasVersion(&store, "k", 3));607    try std.testing.expectEqualStrings("v1", store.getAt("k", 2).?);608    try std.testing.expectEqualStrings("v3", store.get("k").?);609}