File index writing

Author: lalinsky

src/InMemoryIndex.zig

@@ -10,14 +10,14 @@ const Item = common.Item;
 const SearchResults = common.SearchResults;
 const Change = common.Change;
 
-const Segment = @import("InMemorySegment.zig");
-const Segments = std.DoublyLinkedList(Segment);
+const InMemorySegment = @import("InMemorySegment.zig");
+const InMemorySegments = std.DoublyLinkedList(InMemorySegment);
 
 allocator: std.mem.Allocator,
 write_lock: std.Thread.RwLock,
 merge_lock: std.Thread.Mutex,
-segments: Segments,
-max_items_per_segment: usize = 100000,
+segments: InMemorySegments,
+max_items_per_segment: usize = 1_000_000,
 max_segments: usize = 16,
 auto_cleanup: bool = true,
 
@@ -37,28 +37,27 @@ pub fn deinit(self: *Self) void {
     defer self.write_lock.unlock();
 
     while (self.segments.popFirst()) |node| {
-        self.destroyNode(node);
+        self.destroySegment(node);
     }
 }
 
-fn destroyNode(self: *Self, node: *Segments.Node) void {
+fn createSegment(self: *Self) !*InMemorySegments.Node {
+    const node = try self.allocator.create(InMemorySegments.Node);
+    node.data = InMemorySegment.init(self.allocator);
+    return node;
+}
+
+fn destroySegment(self: *Self, node: *InMemorySegments.Node) void {
     node.data.deinit();
     self.allocator.destroy(node);
 }
 
 pub fn update(self: *Self, changes: []const Change) !void {
     var committed = false;
 
-    const node = try self.allocator.create(Segments.Node);
+    const node = try self.createSegment();
     defer {
-        if (!committed) self.allocator.destroy(node);
-    }
-
-    node.data = Segment.init(self.allocator);
-    defer {
-        if (!committed) {
-            node.data.deinit();
-        }
+        if (!committed) self.destroySegment(node);
     }
 
     var num_items: usize = 0;
@@ -134,9 +133,9 @@ fn hasNewerVersion(self: *Self, id: u32, version: u32) bool {
 }
 
 const Merge = struct {
-    first: *Segments.Node,
-    last: *Segments.Node,
-    replacement: *Segments.Node,
+    first: *InMemorySegments.Node,
+    last: *InMemorySegments.Node,
+    replacement: *InMemorySegments.Node,
 };
 
 fn prepareMerge(self: *Self) !?Merge {
@@ -158,7 +157,7 @@ fn prepareMerge(self: *Self) !?Merge {
         return null;
     }
 
-    var best_node: ?*Segments.Node = null;
+    var best_node: ?*InMemorySegments.Node = null;
     var best_score: f64 = std.math.inf(f64);
     segments_iter = self.segments.first;
     var level_size = @as(f64, @floatFromInt(total_size)) / 2;
@@ -187,23 +186,17 @@ fn prepareMerge(self: *Self) !?Merge {
 
     const segment1 = node1.data;
     const segment2 = node2.data;
-    const segments = [2]Segment{ segment1, segment2 };
+    const segments = [2]InMemorySegment{ segment1, segment2 };
 
     var committed = false;
 
-    const node = try self.allocator.create(Segments.Node);
+    const node = try self.createSegment();
     defer {
-        if (!committed) self.allocator.destroy(node);
+        if (!committed) self.destroySegment(node);
     }
 
     const merge = Merge{ .first = node1, .last = node2, .replacement = node };
 
-    node.data = Segment.init(self.allocator);
-    defer {
-        if (!committed) {
-            node.data.deinit();
-        }
-    }
     node.data.version = segment1.version;
 
     var total_docs: usize = 0;
@@ -267,11 +260,11 @@ fn commitMerge(self: *Self, merge: Merge) void {
 
     self.segments.insertAfter(merge.last, merge.replacement);
 
-    var iter: ?*Segments.Node = merge.first;
+    var iter: ?*InMemorySegments.Node = merge.first;
     while (iter) |node| {
         iter = node.next;
         self.segments.remove(node);
-        self.destroyNode(node);
+        self.destroySegment(node);
         if (node == merge.last) break;
     }
 
@@ -288,7 +281,7 @@ fn mergeSegments(self: *Self) !void {
     }
 }
 
-pub fn freezeFirstSegment(self: *Self) ?*Segment {
+pub fn freezeFirstSegment(self: *Self) ?*InMemorySegment {
     self.merge_lock.lock();
     defer self.merge_lock.unlock();
 
@@ -311,7 +304,7 @@ pub fn freezeFirstSegment(self: *Self) ?*Segment {
     return null;
 }
 
-pub fn removeFrozenSegment(self: *Self, segment: *Segment) void {
+pub fn removeSegment(self: *Self, segment: *InMemorySegment) void {
     self.merge_lock.lock();
     defer self.merge_lock.unlock();
 
@@ -323,7 +316,7 @@ pub fn removeFrozenSegment(self: *Self, segment: *Segment) void {
         if (&node.data == segment) {
             if (node.data.frozen) {
                 self.segments.remove(node);
-                self.destroyNode(node);
+                self.destroySegment(node);
                 return;
             }
         }

src/Index.zig

@@ -62,32 +62,61 @@ pub fn start(self: *Self) !void {
 pub fn deinit(self: *Self) void {
     self.stage.deinit();
     while (self.segments.popFirst()) |node| {
-        node.data.deinit();
-        self.allocator.destroy(node);
+        self.destroySegment(node);
     }
     self.scheduler.deinit();
 }
 
+fn createSegment(self: *Self) !*Segments.Node {
+    const node = try self.allocator.create(Segments.Node);
+    node.data = Segment.init(self.allocator);
+    return node;
+}
+
+fn destroySegment(self: *Self, node: *Segments.Node) void {
+    node.data.deinit();
+    self.allocator.destroy(node);
+}
+
+fn write(self: *Self) !void {
+    var file = try self.dir.atomicFile("index.dat", .{});
+    defer file.deinit();
+
+    var writer = file.file.writer();
+    try writer.writeInt(u32, @truncate(self.segments.len), .little);
+    var it = self.segments.first;
+    while (it) |node| : (it = node.next) {
+        try writer.writeInt(u32, node.data.version[0], .little);
+        try writer.writeInt(u32, node.data.version[1], .little);
+    }
+
+    try file.finish();
+}
+
 fn cleanup(self: *Self) !void {
     if (!self.run_cleanup) return;
 
     log.info("running cleanup", .{});
 
     // try self.stage.cleanup();
 
-    const staged_segment_or_null = self.stage.freezeFirstSegment();
-    if (staged_segment_or_null) |staged_segment| {
-        const segment = Segment.init(self.allocator);
-        try segment.convert(self.dir, staged_segment);
+    if (self.stage.freezeFirstSegment()) |segment| {
+        var commited = false;
+        const node = try self.createSegment();
+        defer {
+            if (!commited) self.destroySegment(node);
+        }
+
+        try node.data.convert(self.dir, segment);
 
         self.write_lock.lock();
         defer self.write_lock.unlock();
 
-        var node = try self.allocator.create(Segments.Node);
-        node.data = segment;
+        try self.write();
 
-        self.stage.removeFrozenSegment(staged_segment);
+        self.stage.removeSegment(segment);
         self.segments.append(node);
+        commited = true;
     }
 }
 

src/Segment.zig

@@ -9,8 +9,10 @@ const InMemorySegment = @import("InMemorySegment.zig");
 
 const filefmt = @import("filefmt.zig");
 
+pub const Version = std.meta.Tuple(&.{ u32, u32 });
+
 allocator: std.mem.Allocator,
-version: u32 = 0,
+version: Version = .{ 0, 0 },
 docs: std.AutoHashMap(u32, bool),
 index: std.ArrayList(u32),
 block_size: usize = 0,
@@ -68,14 +70,20 @@ pub fn search(self: *Self, hashes: []const u32, results: *SearchResults) !void {
             }
             const matches = std.sort.equalRange(Item, Item{ .hash = hash, .id = 0 }, block_items.items, {}, Item.cmpByHash);
             for (matches[0]..matches[1]) |i| {
-                try results.incr(block_items.items[i].id, self.version);
+                try results.incr(block_items.items[i].id, self.version[1]);
             }
         }
     }
 }
 
 const max_file_name_size = 255;
-const file_name_fmt = "segment-{d}.data";
+const file_name_fmt = "segment-{d}-{d}.dat";
+
+fn read(self: *Self, dir: std.fs.Dir, file_name: []const u8) !void {
+    const file = try dir.openFile(file_name, .{});
+    defer file.close();
+    try filefmt.readFile(file, self);
+}
 
 pub fn convert(self: *Self, dir: std.fs.Dir, source: *InMemorySegment) !void {
     if (!source.frozen) {
@@ -86,13 +94,31 @@ pub fn convert(self: *Self, dir: std.fs.Dir, source: *InMemorySegment) !void {
     }
 
     var file_name_buf: [max_file_name_size]u8 = undefined;
-    const file_name = try std.fmt.bufPrint(&file_name_buf, file_name_fmt, .{source.version});
+    const file_name = try std.fmt.bufPrint(&file_name_buf, file_name_fmt, .{ source.version, source.version });
 
-    const file = try dir.createFile(file_name, .{ .exclusive = true, .read = true });
-    defer file.close();
+    var file = try dir.atomicFile(file_name, .{});
+    defer file.deinit();
+    try filefmt.writeFile(file.file.writer(), source);
+    try file.finish();
 
-    try filefmt.writeFile(file.writer(), source);
-    try filefmt.readFile(file, self);
+    try self.read(dir, file_name);
+}
+
+pub fn merge(self: *Self, dir: std.fs.Dir, sources: [2]*Self) !void {
+    if (sources[0].version[1] + 1 != sources[1].version[0]) {
+        return error.SourceSegmentVersionMismatch;
+    }
+
+    self.version[0] = sources[0].version[0];
+    self.version[1] = sources[1].version[1];
+
+    var file_name_buf: [max_file_name_size]u8 = undefined;
+    const file_name = try std.fmt.bufPrint(&file_name_buf, file_name_fmt, .{ self.version[0], self.version[1] });
+
+    var file = try dir.atomicFile(file_name, .{});
+    defer file.deinit();
+
+    try self.read(dir, file_name);
 }
 
 test "convert" {
@@ -113,7 +139,8 @@ test "convert" {
 
     try segment.convert(tmpDir.dir, &source);
 
-    try std.testing.expectEqual(1, segment.version);
+    try std.testing.expectEqual(1, segment.version[0]);
+    try std.testing.expectEqual(1, segment.version[1]);
     try std.testing.expectEqual(1, segment.docs.count());
     try std.testing.expectEqual(1, segment.index.items.len);
 }

src/SegmentList.zig

@@ -342,7 +342,8 @@ pub fn readFile(file: fs.File, segment: *Segment) !void {
     }
     const num_blocks = blocks_data_size / header.block_size;
 
-    segment.version = header.version;
+    segment.version[0] = header.version;
+    segment.version[1] = header.version;
     segment.block_size = header.block_size;
 
     try segment.docs.ensureTotalCapacity(header.num_docs);
@@ -403,7 +404,8 @@ test "writeFile/readFile" {
 
         try readFile(file, &segment);
 
-        try testing.expectEqual(1, segment.version);
+        try testing.expectEqual(1, segment.version[0]);
+        try testing.expectEqual(1, segment.version[1]);
         try testing.expectEqual(1, segment.docs.count());
         try testing.expectEqual(1, segment.index.items.len);
         try testing.expectEqual(1, segment.index.items[0]);