Fix delayed index load

Author: lalinsky

build.zig

@@ -72,8 +72,11 @@ pub fn build(b: *std.Build) void {
     main_tests.root_module.addImport("msgpack", msgpack.module("msgpack"));
 
     const run_unit_tests = b.addRunArtifact(main_tests);
+    run_unit_tests.has_side_effects = true;
+
     const run_integration_tests = b.addSystemCommand(&[_][]const u8{ "pytest", "-vv", "tests/" });
     run_integration_tests.step.dependOn(b.getInstallStep());
+    run_integration_tests.has_side_effects = true;
 
     var unit_tests_step = b.step("unit-tests", "Run unit tests");
     unit_tests_step.dependOn(&run_unit_tests.step);

src/Index.zig

@@ -52,7 +52,7 @@ dir: std.fs.Dir,
 oplog: Oplog,
 
 open_lock: std.Thread.Mutex = .{},
-is_ready: std.atomic.Value(bool),
+is_ready: std.Thread.ResetEvent = .{},
 load_task: ?Scheduler.Task = null,
 
 segments_lock: std.Thread.RwLock = .{},
@@ -113,7 +113,6 @@ pub fn init(allocator: std.mem.Allocator, scheduler: *Scheduler, parent_dir: std
         .segments_lock = .{},
         .memory_segments = memory_segments,
         .file_segments = file_segments,
-        .is_ready = std.atomic.Value(bool).init(false),
     };
 }
 
@@ -287,13 +286,13 @@ fn maybeMergeMemorySegments(self: *Self) !bool {
 }
 
 pub fn open(self: *Self, create: bool) !void {
-    self.open_lock.lock();
-    defer self.open_lock.unlock();
-
-    if (self.is_ready.load(.monotonic)) {
+    if (self.is_ready.isSet()) {
         return;
     }
 
+    self.open_lock.lock();
+    defer self.open_lock.unlock();
+
     if (self.load_task != null) {
         return error.AlreadyOpening;
     }
@@ -333,11 +332,11 @@ fn load(self: *Self, manifest: []SegmentInfo) !void {
     self.checkpoint_task = try self.scheduler.createTask(.medium, checkpointTask, .{self});
     self.file_segment_merge_task = try self.scheduler.createTask(.low, fileSegmentMergeTask, .{self});
 
-    try self.oplog.open(1, updateInternal, self);
+    try self.oplog.open(last_commit_id + 1, updateInternal, self);
 
     log.info("index loaded", .{});
 
-    self.is_ready.store(true, .monotonic);
+    self.is_ready.set();
 }
 
 fn loadTask(self: *Self, manifest: []SegmentInfo) void {
@@ -390,14 +389,18 @@ fn readyForCheckpoint(self: *Self) ?MemorySegmentNode {
     return null;
 }
 
-fn checkIfReady(self: Self) !void {
-    if (!self.is_ready.load(.monotonic)) {
+pub fn waitForReady(self: *Self, timeout_ms: u32) !void {
+    try self.is_ready.timedWait(timeout_ms * std.time.us_per_ms);
+}
+
+pub fn checkReady(self: *Self) !void {
+    if (!self.is_ready.isSet()) {
         return error.IndexNotReady;
     }
 }
 
 pub fn update(self: *Self, changes: []const Change) !void {
-    try self.checkIfReady();
+    try self.checkReady();
     try self.updateInternal(changes, null);
 }
 
@@ -426,7 +429,7 @@ fn updateInternal(self: *Self, changes: []const Change, commit_id: ?u64) !void {
 }
 
 pub fn acquireReader(self: *Self) !IndexReader {
-    try self.checkIfReady();
+    try self.checkReady();
 
     self.segments_lock.lockShared();
     defer self.segments_lock.unlockShared();

src/index_tests.zig

@@ -4,6 +4,7 @@ const common = @import("common.zig");
 const Change = @import("change.zig").Change;
 const SearchResults = common.SearchResults;
 const Scheduler = @import("utils/Scheduler.zig");
+const Deadline = @import("utils/Deadline.zig");
 
 const Index = @import("Index.zig");
 
@@ -76,6 +77,8 @@ test "index create, update, reopen and search" {
     var scheduler = Scheduler.init(std.testing.allocator);
     defer scheduler.deinit();
 
+    try scheduler.start(2);
+
     var hashes: [100]u32 = undefined;
 
     {
@@ -95,6 +98,7 @@ test "index create, update, reopen and search" {
         defer index.deinit();
 
         try index.open(false);
+        try index.waitForReady(10000);
 
         var results = try index.search(generateRandomHashes(&hashes, 1), std.testing.allocator, .{});
         defer results.deinit();
@@ -115,6 +119,8 @@ test "index many updates" {
     var scheduler = Scheduler.init(std.testing.allocator);
     defer scheduler.deinit();
 
+    try scheduler.start(2);
+
     var hashes: [100]u32 = undefined;
 
     {
@@ -135,6 +141,7 @@ test "index many updates" {
     defer index.deinit();
 
     try index.open(false);
+    try index.waitForReady(100000);
 
     {
         var results = try index.search(generateRandomHashes(&hashes, 0), std.testing.allocator, .{});

src/segment_list.zig

@@ -230,6 +230,7 @@ pub fn SegmentListManager(Segment: type) type {
         }
 
         pub fn count(self: Self) usize {
+            // std.debug.print("count {}\n", .{self.segments.value.nodes});
             return self.segments.value.nodes.items.len;
         }
 

src/server.zig

@@ -233,11 +233,21 @@ fn handleNotFound(_: *Context, req: *httpz.Request, res: *httpz.Response) !void
 }
 
 fn handleError(_: *Context, req: *httpz.Request, res: *httpz.Response, err: anyerror) void {
-    log.err("unhandled error in {s}: {any}", .{ req.url.raw, err });
-    writeErrorResponse(500, err, req, res) catch {
-        res.status = 500;
-        res.body = "internal error";
-    };
+    switch (err) {
+        error.IndexNotReady => {
+            writeErrorResponse(503, err, req, res) catch {
+                res.status = 503;
+                res.body = "not ready yet";
+            };
+        },
+        else => {
+            log.err("unhandled error in {s}: {any}", .{ req.url.raw, err });
+            writeErrorResponse(500, err, req, res) catch {
+                res.status = 500;
+                res.body = "internal error";
+            };
+        },
+    }
 }
 
 fn writeErrorResponse(status: u16, err: anyerror, req: *httpz.Request, res: *httpz.Response) !void {
@@ -484,6 +494,8 @@ fn handleIndexHealth(ctx: *Context, req: *httpz.Request, res: *httpz.Response) !
     const index = try getIndex(ctx, req, res, false) orelse return;
     defer releaseIndex(ctx, index);
 
+    try index.checkReady();
+
     try res.writer().writeAll("OK\n");
 }
 

src/utils/Scheduler.zig

@@ -11,6 +11,7 @@ const Priority = enum(u8) {
 const TaskStatus = struct {
     reschedule: usize = 0,
     scheduled: bool = false,
+    running: bool = false,
     done: std.Thread.ResetEvent = .{},
     priority: Priority,
     ctx: *anyopaque,
@@ -88,12 +89,22 @@ pub fn createTask(self: *Self, priority: Priority, comptime func: anytype, args:
     return task;
 }
 
-pub fn destroyTask(self: *Self, task: Task) void {
+fn dequeue(self: *Self, task: Task) void {
     self.queue_mutex.lock();
+    defer self.queue_mutex.unlock();
+
     if (task.data.scheduled) {
         self.queue.remove(task);
+        task.next = null;
+        task.prev = null;
+        task.data.scheduled = false;
     }
-    self.queue_mutex.unlock();
+
+    task.data.reschedule = 0;
+}
+
+pub fn destroyTask(self: *Self, task: Task) void {
+    self.dequeue(task);
 
     task.data.done.wait();
 
@@ -108,15 +119,14 @@ pub fn scheduleTask(self: *Self, task: Task) void {
     self.queue_mutex.lock();
     defer self.queue_mutex.unlock();
 
-    if (task.data.scheduled) {
+    if (task.data.scheduled or task.data.running) {
         task.data.reschedule += 1;
-        return;
+    } else {
+        self.enqueue(task);
     }
-
-    self.addToQueue(task);
 }
 
-fn addToQueue(self: *Self, task: *Queue.Node) void {
+fn enqueue(self: *Self, task: *Queue.Node) void {
     task.data.scheduled = true;
     self.queue.prepend(task);
     self.queue_not_empty.signal();
@@ -133,6 +143,8 @@ fn getTaskToRun(self: *Self) ?*Queue.Node {
         };
         task.prev = null;
         task.next = null;
+        task.data.scheduled = false;
+        task.data.running = true;
         task.data.done.reset();
         return task;
     }
@@ -145,11 +157,10 @@ fn markAsDone(self: *Self, task: *Queue.Node) void {
 
     if (task.data.reschedule > 0) {
         task.data.reschedule -= 1;
-        self.addToQueue(task);
-    } else {
-        task.data.scheduled = false;
+        self.enqueue(task);
     }
 
+    task.data.running = false;
     task.data.done.set();
 }
 
@@ -207,7 +218,7 @@ test "Scheduler: smoke test" {
     };
     var counter: Counter = .{};
 
-    const task = try scheduler.createTask(.high, Counter.incr, &counter);
+    const task = try scheduler.createTask(.high, Counter.incr, .{&counter});
     defer scheduler.destroyTask(task);
 
     for (0..3) |_| {