lib/sql/src/store.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const simd = @import("simd");
  3 const trace = @import("trace.zig");
  4 
  5 const Bytes = simd.ScalableTag(u8);
  6 
  7 const Allocator = std.mem.Allocator;
  8 
  9 const StoreError = error{
 10     InvalidRange,
 11     TransactionClosed,
 12     TransactionConflict,
 13 };
 14 
 15 pub const Error = Allocator.Error || StoreError;
 16 
 17 pub const Entry = struct {
 18     key: []const u8,
 19     value: []const u8,
 20     generation: u64,
 21 };
 22 
 23 const Version = struct {
 24     key: []u8,
 25     value: []u8,
 26     generation: u64,
 27     deleted: bool = false,
 28 };
 29 
 30 const Mutation = struct {
 31     key: []u8,
 32     value: []u8 = &.{},
 33     deleted: bool = false,
 34 
 35     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 };
 41 
 42 pub const Store = struct {
 43     allocator: Allocator,
 44     versions: std.ArrayList(Version) = .empty,
 45     generation: u64 = 0,
 46 
 47     pub fn init(allocator: Allocator) Store {
 48         return .{ .allocator = allocator };
 49     }
 50 
 51     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     }
 59 
 60     pub fn currentGeneration(self: *const Store) u64 {
 61         return self.generation;
 62     }
 63 
 64     pub fn beginRead(self: *const Store) Snapshot {
 65         return .{ .store = self, .generation = self.generation };
 66     }
 67 
 68     pub fn beginWrite(self: *Store) Write {
 69         return .{ .store = self, .snapshot_generation = self.generation };
 70     }
 71 
 72     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     }
 77 
 78     pub fn getAt(self: *const Store, key: []const u8, generation: u64) ?[]const u8 {
 79         return valueFromIndex(self, latestVisibleIndex(self, key, generation));
 80     }
 81 
 82     pub fn range(self: *const Store, start: ?[]const u8, end: ?[]const u8) Error!Range {
 83         return rangeAt(self, start, end, self.generation);
 84     }
 85 
 86     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     }
100 
101     pub fn compact(self: *Store, oldest_visible_generation: u64) void {
102         const phase = trace.scope("store.compact");
103         defer phase.end();
104 
105         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             }
115 
116             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             }
126 
127             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.?) or
131                     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             }
140 
141             read_index = group_end;
142         }
143         self.versions.shrinkRetainingCapacity(write_index);
144     }
145 };
146 
147 pub const Snapshot = struct {
148     store: *const Store,
149     generation: u64,
150 
151     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     }
156 
157     pub fn range(self: Snapshot, start: ?[]const u8, end: ?[]const u8) Error!Range {
158         return self.store.rangeAt(start, end, self.generation);
159     }
160 };
161 
162 pub const Write = struct {
163     store: *Store,
164     snapshot_generation: u64,
165     mutations: std.ArrayList(Mutation) = .empty,
166     closed: bool = false,
167 
168     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     }
173 
174     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     }
182 
183     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     }
189 
190     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     }
196 
197     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         }
205 
206         for (self.mutations.items) |mutation| {
207             if (latestGenerationForKey(self.store, mutation.key)) |generation| {
208                 if (generation > self.snapshot_generation) return error.TransactionConflict;
209             }
210         }
211 
212         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     }
230 
231     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         }
240 
241         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     }
253 
254     fn localMutation(self: *const Write, key: []const u8) ?Mutation {
255         if (self.localMutationIndex(key)) |index| return self.mutations.items[index];
256         return null;
257     }
258 
259     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 };
273 
274 pub const Range = struct {
275     store: *const Store,
276     start: ?[]const u8,
277     end: ?[]const u8,
278     generation: u64,
279     index: usize,
280 
281     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             }
295 
296             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     }
318 
319     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 };
325 
326 fn 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 }
340 
341 fn 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 }
347 
348 fn 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 }
357 
358 fn 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 }
370 
371 fn 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 }
384 
385 fn versionLessThan(_: void, left: Version, right: Version) bool {
386     return versionPrecedesKeyGeneration(left, right.key, right.generation);
387 }
388 
389 fn 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 }
396 
397 fn 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 }
405 
406 fn 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 }
415 
416 fn 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 }
424 
425 fn 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 }
432 
433 fn eqlKey(left: []const u8, right: []const u8) bool {
434     return std.mem.eql(u8, left, right);
435 }
436 
437 fn freeOwned(allocator: Allocator, bytes: []u8) void {
438     if (bytes.len != 0) allocator.free(bytes);
439 }
440 
441 fn 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 }
447 
448 fn 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 }
454 
455 test "committed writes are visible in key order" {
456     var store = Store.init(std.testing.allocator);
457     defer store.deinit();
458 
459     _ = try putCommitted(&store, "b", "two");
460     _ = try putCommitted(&store, "a", "one");
461 
462     try std.testing.expectEqualStrings("one", store.get("a").?);
463     try std.testing.expectEqualStrings("two", store.get("b").?);
464 
465     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 }
472 
473 test "snapshots preserve old values across later commits" {
474     var store = Store.init(std.testing.allocator);
475     defer store.deinit();
476 
477     _ = try putCommitted(&store, "name", "first");
478     const snapshot = store.beginRead();
479     _ = try putCommitted(&store, "name", "second");
480 
481     try std.testing.expectEqualStrings("first", snapshot.get("name").?);
482     try std.testing.expectEqualStrings("second", store.get("name").?);
483 }
484 
485 test "delete creates a tombstone for new readers" {
486     var store = Store.init(std.testing.allocator);
487     defer store.deinit();
488 
489     _ = try putCommitted(&store, "k", "v");
490     const snapshot = store.beginRead();
491     _ = try deleteCommitted(&store, "k");
492 
493     try std.testing.expectEqualStrings("v", snapshot.get("k").?);
494     try std.testing.expect(store.get("k") == null);
495 }
496 
497 test "write transactions detect same key conflicts" {
498     var store = Store.init(std.testing.allocator);
499     defer store.deinit();
500 
501     var left = store.beginWrite();
502     defer left.deinit();
503     var right = store.beginWrite();
504     defer right.deinit();
505 
506     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 }
511 
512 test "write transactions allow disjoint commits from the same snapshot" {
513     var store = Store.init(std.testing.allocator);
514     defer store.deinit();
515 
516     var left = store.beginWrite();
517     defer left.deinit();
518     var right = store.beginWrite();
519     defer right.deinit();
520 
521     try left.put("a", "left");
522     try right.put("b", "right");
523     _ = try left.commit();
524     _ = try right.commit();
525 
526     try std.testing.expectEqualStrings("left", store.get("a").?);
527     try std.testing.expectEqualStrings("right", store.get("b").?);
528 }
529 
530 test "range scan honors half open bounds and tombstones" {
531     var store = Store.init(std.testing.allocator);
532     defer store.deinit();
533 
534     _ = try putCommitted(&store, "a", "1");
535     _ = try putCommitted(&store, "b", "2");
536     _ = try putCommitted(&store, "c", "3");
537     _ = try deleteCommitted(&store, "b");
538 
539     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 }
544 
545 test "multi key commits preserve canonical version order" {
546     var store = Store.init(std.testing.allocator);
547     defer store.deinit();
548 
549     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();
554 
555     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();
561 
562     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 }
575 
576 test "compaction drops obsolete floor when oldest generation is present" {
577     var store = Store.init(std.testing.allocator);
578     defer store.deinit();
579 
580     _ = try putCommitted(&store, "k", "v1");
581     const snapshot_generation = try putCommitted(&store, "k", "v2");
582     _ = try putCommitted(&store, "k", "v3");
583 
584     store.compact(snapshot_generation);
585 
586     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 }
592 
593 test "compaction keeps floor when oldest generation falls between versions" {
594     var store = Store.init(std.testing.allocator);
595     defer store.deinit();
596 
597     _ = try putCommitted(&store, "k", "v1");
598     _ = try putCommitted(&store, "other", "advance");
599     _ = try putCommitted(&store, "k", "v3");
600 
601     store.compact(2);
602 
603     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 }