Better merge performance

Author: lalinsky

src/Index.zig

@@ -31,7 +31,7 @@ const Self = @This();
 
 const Options = struct {
     create: bool = false,
-    min_segment_size: usize = 1_000_000,
+    min_segment_size: usize = 250_000,
     max_segment_size: usize = 1_000_000_000,
 };
 
@@ -86,14 +86,16 @@ pub fn init(allocator: std.mem.Allocator, dir: std.fs.Dir, options: Options) !Se
         .min_segment_size = options.min_segment_size,
         .max_segment_size = options.max_segment_size,
         .segments_per_level = 10,
-        .segments_per_merge = 2, // TODO increase to 10
+        .segments_per_merge = 10,
+        .strategy = .balanced,
     };
 
     const memory_segment_merge_policy = TieredMergePolicy(MemorySegment){
         .min_segment_size = 100,
         .max_segment_size = options.min_segment_size,
         .segments_per_level = 10,
-        .segments_per_merge = 2, // TODO increase to 5
+        .segments_per_merge = 10,
+        .strategy = .aggressive,
     };
 
     return .{
@@ -270,7 +272,7 @@ fn loadSegments(self: *Self) !void {
 fn doCheckpoint(self: *Self) !bool {
     const start_time = std.time.milliTimestamp();
 
-    var src = try self.readyForCheckpoint() orelse return false;
+    var src = self.readyForCheckpoint() orelse return false;
 
     var src_reader = src.data.reader();
     defer src_reader.close();
@@ -322,7 +324,7 @@ fn checkpointThreadFn(self: *Self) void {
 
         if (self.doCheckpoint()) |successful| {
             if (successful) {
-                self.file_segment_merge_condition.signal();
+                self.signalFileSegmentMerge();
                 continue;
             }
         } else |err| {
@@ -502,7 +504,7 @@ const Checkpoint = struct {
     dest: ?*FileSegmentNode = null,
 };
 
-fn readyForCheckpoint(self: *Self) !?*MemorySegmentNode {
+fn readyForCheckpoint(self: *Self) ?*MemorySegmentNode {
     self.segments_lock.lockShared();
     defer self.segments_lock.unlockShared();
 
@@ -515,15 +517,27 @@ fn readyForCheckpoint(self: *Self) !?*MemorySegmentNode {
     return null;
 }
 
-pub fn update(self: *Self, changes: []const Change) !void {
-    //const t1 = std.time.milliTimestamp();
-    if (try self.maybeMergeMemorySegments()) {
-        self.checkpoint_condition.signal();
+fn signalMemorySegmentMerge(self: *Self) void {
+    self.segments_lock.lockShared();
+    defer self.segments_lock.unlockShared();
+
+    if (self.memory_segments.needsMerge()) {
+        self.memory_segment_merge_condition.signal();
     }
-    //const t2 = std.time.milliTimestamp();
+}
+
+fn signalFileSegmentMerge(self: *Self) void {
+    self.segments_lock.lockShared();
+    defer self.segments_lock.unlockShared();
+
+    if (self.file_segments.needsMerge()) {
+        self.file_segment_merge_condition.signal();
+    }
+}
+
+pub fn update(self: *Self, changes: []const Change) !void {
     try self.oplog.write(changes, Updater{ .index = self });
-    // const t3 = std.time.milliTimestamp();
-    //log.info("merge: {}ms, update: {}ms", .{ t2 - t1, t3 - t2 });
+    self.signalMemorySegmentMerge();
 }
 
 pub fn search(self: *Self, hashes: []const u32, allocator: std.mem.Allocator, deadline: Deadline) !SearchResults {

src/filefmt.zig

@@ -91,7 +91,7 @@ test "check writeVarint32" {
 }
 
 pub const max_file_name_size = 64;
-const segment_file_name_fmt = "segment-{d}-{d}.dat";
+const segment_file_name_fmt = "segment-{x:0>8}-{x:0>8}.dat";
 pub const index_file_name = "index.dat";
 
 pub fn buildSegmentFileName(buf: []u8, version: common.SegmentID) []u8 {

src/segment_list.zig

@@ -19,11 +19,14 @@ pub fn SegmentList(Segment: type) type {
         merge_policy: MergePolicy,
         segments: List,
 
+        num_allowed_segments: std.atomic.Value(usize),
+
         pub fn init(allocator: std.mem.Allocator, merge_policy: MergePolicy) Self {
             return .{
                 .allocator = allocator,
                 .merge_policy = merge_policy,
                 .segments = .{},
+                .num_allowed_segments = std.atomic.Value(usize).init(0),
             };
         }
 
@@ -134,6 +137,10 @@ pub fn SegmentList(Segment: type) type {
             }
         }
 
+        pub fn needsMerge(self: *Self) bool {
+            return self.segments.len > self.num_allowed_segments.load(.monotonic);
+        }
+
         pub const SegmentsToMerge = MergePolicy.Candidate;
 
         pub const PreparedMerge = struct {
@@ -143,7 +150,10 @@ pub fn SegmentList(Segment: type) type {
         };
 
         pub fn prepareMerge(self: *Self) !?PreparedMerge {
-            const sources = self.merge_policy.findSegmentsToMerge(self.segments) orelse return null;
+            const result = self.merge_policy.findSegmentsToMerge(self.segments);
+            self.num_allowed_segments.store(result.num_allowed_segments, .monotonic);
+
+            const sources = result.candidate orelse return null;
 
             var merger = SegmentMerger(Segment).init(self.allocator, self);
             errdefer merger.deinit();
@@ -157,7 +167,11 @@ pub fn SegmentList(Segment: type) type {
             try merger.prepare();
 
             const target = try self.createSegment();
-            return .{ .sources = sources, .merger = merger, .target = target };
+            return .{
+                .sources = sources,
+                .merger = merger,
+                .target = target,
+            };
         }
 
         pub fn cleanupAfterMerge(self: *Self, merge: PreparedMerge, cleanup_args: anytype) void {

src/segment_merge_policy.zig

@@ -1,7 +1,7 @@
 const std = @import("std");
 const log = std.log.scoped(.segment_merge_policy);
 
-const verbose = false;
+const verbose = true;
 
 pub fn MergeCandidate(comptime Segment: type) type {
     return struct {
@@ -24,6 +24,13 @@ pub fn TieredMergePolicy(comptime T: type) type {
         segments_per_merge: u32 = 10,
         segments_per_level: u32 = 10,
 
+        strategy: Strategy = .balanced,
+
+        const Strategy = enum {
+            balanced,
+            aggressive,
+        };
+
         const SegmentList = std.DoublyLinkedList(T);
         const SegmentNode = SegmentList.Node;
 
@@ -74,21 +81,34 @@ pub fn TieredMergePolicy(comptime T: type) type {
             return num_allowed_segments + num_oversized_segments;
         }
 
-        pub fn findSegmentsToMerge(self: Self, segments: SegmentList) ?Candidate {
+        pub const FindSegmentsToMergeResult = struct {
+            num_allowed_segments: usize,
+            candidate: ?Candidate,
+        };
+
+        pub fn findSegmentsToMerge(self: Self, segments: SegmentList) FindSegmentsToMergeResult {
             const num_segments = segments.len;
             const num_allowed_segments = self.calculateBudget(segments);
             log.debug("budget: {} segments", .{num_allowed_segments});
 
+            var result = FindSegmentsToMergeResult{
+                .num_allowed_segments = num_allowed_segments,
+                .candidate = null,
+            };
+
             if (num_allowed_segments >= segments.len) {
-                return null;
+                return result;
             }
 
             const merge_factor = @min(self.segments_per_merge, self.segments_per_level);
             const log_merge_factor = @log2(@as(f64, @floatFromInt(merge_factor)));
             const log_min_segment_size = @log2(@as(f64, @floatFromInt(self.min_segment_size)));
 
             const tier_scaling_factor = @as(f64, @floatFromInt(num_allowed_segments)) / @as(f64, @floatFromInt(num_segments)) / @as(f64, @floatFromInt(self.segments_per_level));
-            var tier = @as(f64, @floatFromInt(num_segments - 1)) * tier_scaling_factor;
+            const top_tier = @as(f64, @floatFromInt(num_segments)) * tier_scaling_factor;
+            var tier = top_tier;
+
+            var segment_no: usize = 0;
 
             var best_candidate: ?Candidate = null;
             var best_score: f64 = 0.0;
@@ -98,6 +118,7 @@ pub fn TieredMergePolicy(comptime T: type) type {
             var iter = segments.first;
             while (iter) |current_node| : (iter = current_node.next) {
                 tier -= tier_scaling_factor;
+                segment_no += 1;
 
                 if (current_node.data.getSize() > self.max_segment_size) {
                     // skip oversized segments
@@ -126,7 +147,16 @@ pub fn TieredMergePolicy(comptime T: type) type {
 
                     const log_size = @log2(@as(f64, @floatFromInt(candidate.size)));
                     const candidate_tier = (log_size - log_min_segment_size) / log_merge_factor;
-                    const score = candidate_tier - tier;
+                    var score = candidate_tier - tier;
+
+                    const adjustment_factor: f64 = switch (self.strategy) {
+                        .balanced => 1.2,
+                        .aggressive => 1.8,
+                    };
+
+                    const adjustment = @as(f64, @floatFromInt(candidate.num_segments)) / @as(f64, @floatFromInt(self.segments_per_merge));
+                    score = score - adjustment_factor * adjustment;
+
                     // std.debug.print("candidate {}-{}: len={} size={} candidate_tier={}, score={d}\n", .{ candidate.start.data.id, candidate.end.data.id, candidate.num_segments, candidate.size, candidate_tier, score });
                     if (score < best_score or best_candidate == null) {
                         best_candidate = candidate;
@@ -142,7 +172,8 @@ pub fn TieredMergePolicy(comptime T: type) type {
                 max_merge_size = current_node.data.getSize();
             }
 
-            return best_candidate;
+            result.candidate = best_candidate;
+            return result;
         }
     };
 }
@@ -186,6 +217,7 @@ test "TieredMergePolicy" {
         .max_segment_size = 100000,
         .segments_per_merge = 10,
         .segments_per_level = 5,
+        .strategy = .aggressive,
     };
 
     var last_id: u64 = 1;
@@ -200,6 +232,9 @@ test "TieredMergePolicy" {
         last_id += 1;
     }
 
+    var total_merge_size: u64 = 0;
+    var total_merge_count: u64 = 0;
+
     for (0..1000) |_| {
         if (verbose) {
             std.debug.print("---\n", .{});
@@ -221,14 +256,22 @@ test "TieredMergePolicy" {
             }
         }
 
-        const candidate = policy.findSegmentsToMerge(segments) orelse continue;
+        const result = policy.findSegmentsToMerge(segments);
+        const candidate = result.candidate orelse continue;
+
+        total_merge_size += candidate.num_segments;
+        total_merge_count += 1;
 
         if (verbose) {
             std.debug.print("merging {}-{}\n", .{ candidate.start.data.id, candidate.end.data.id });
         }
         try applyMerge(MockSegment, &segments, candidate, std.testing.allocator);
     }
 
+    if (verbose) {
+        std.debug.print("avg merge size: {}\n", .{total_merge_size / total_merge_count});
+    }
+
     const num_allowed_segmens = policy.calculateBudget(segments);
     try std.testing.expect(num_allowed_segmens >= segments.len);
 }