diff --git a/README.md b/README.md index 9fc8864..c8f089a 100644 --- a/README.md +++ b/README.md @@ -45,6 +45,8 @@ npm install turbodatabase # Node.js | io_uring / kqueue | ✅ Working | Async I/O, event-loop server | | Binary wire protocol | ✅ Working | TCP_NODELAY, pipelining, batch ops | | JSON REST API | ✅ Working | MongoDB-inspired routes on :27017 | +| **Bulk insert (NDJSON + SIMD)** | ✅ Working | **Batch NDJSON with SIMD newline scanning, ~20K docs/s** | +| **WebSocket streaming** | ✅ Working | **Fire-and-forget ingest for high-latency links** | | **Authentication** | ✅ Working | **API key + BLAKE3 hashing, per-key permissions** | | **Schema validation** | ✅ Working | **Required fields, type checking (string/number/bool/object/array)** | | **TTL / document expiry** | ✅ Working | **Per-doc TTL with background reaper** | @@ -154,6 +156,37 @@ curl "http://localhost:27018/db/users?limit=20&filter={\"age\":{\"\$gt\":25}}" curl http://localhost:27018/health ``` +### Bulk Ingestion + +```bash +# NDJSON bulk insert — one request, many documents +curl -X POST http://localhost:27018/db/users/bulk --data-binary @- <<'EOF' +{"key":"alice","value":{"name":"Alice","age":30}} +{"key":"bob","value":{"name":"Bob","age":25}} +{"key":"carol","value":{"name":"Carol","age":35}} +EOF +# → {"inserted":3,"errors":0,"collection":"users","tenant":"default"} +``` + +For high-latency connections (SSH tunnels, WAN), use **WebSocket streaming** to eliminate per-batch round-trips: + +```python +# pip install websocket-client +from turbodb.stream import StreamInserter + +with StreamInserter("ws://host:27018", "users") as s: + for doc in documents: + s.add(doc["key"], doc["value"]) # buffered, fire-and-forget +print(s.result) # {"inserted": 12000, "errors": 0, "frames": 6} +``` + +| Method | Best for | Throughput | +|--------|----------|-----------| +| HTTP single insert | Real-time writes | ~10K/s (localhost) | +| HTTP bulk (NDJSON) | Batch imports | ~20K/s (localhost) | +| HTTP bulk parallel | Bulk over WAN (high bandwidth) | ~1.2K/s (tunnel, 4 conns) | +| WS streaming | Continuous ingest, low-latency fire-and-forget | ~18K/s (localhost) | + ### Python (FFI — no network overhead) ```bash diff --git a/python/turbodb/stream.py b/python/turbodb/stream.py new file mode 100644 index 0000000..b25167a --- /dev/null +++ b/python/turbodb/stream.py @@ -0,0 +1,91 @@ +"""WebSocket streaming bulk insert for TurboDB. + +Fire-and-forget protocol: client streams NDJSON frames without waiting +for per-frame responses. ~10-20x faster than HTTP batches over high-latency +connections (SSH tunnels, WAN). + +Usage:: + + from turbodb.stream import StreamInserter + + with StreamInserter("ws://host:port", "mycollection") as s: + for doc in docs: + s.add(doc["key"], doc["value"]) + print(s.result) # {"inserted": 12000, "errors": 0, "frames": 6} + +Requires: pip install websocket-client +""" + +import json + +try: + import websocket +except ImportError: + websocket = None + + +class StreamInserter: + """Fire-and-forget WebSocket streaming bulk inserter.""" + + def __init__(self, url, collection, tenant=None, batch_size=2000): + if websocket is None: + raise ImportError("pip install websocket-client") + self.url = url + self.collection = collection + self.tenant = tenant + self.batch_size = batch_size + self._ws = None + self._buf = [] + self._sent_frames = 0 + self.result = None + + def open(self): + """Connect and enter STREAM mode.""" + self._ws = websocket.create_connection(self.url) + init = f"STREAM /db/{self.collection}/bulk" + if self.tenant: + init += f" tenant={self.tenant}" + self._ws.send(init) + ack = json.loads(self._ws.recv()) + if "error" in ack: + raise RuntimeError(f"Stream init failed: {ack['error']}") + return self + + def add(self, key, value): + """Buffer a document. Auto-flushes at batch_size.""" + if isinstance(value, (dict, list)): + value = json.dumps(value) + self._buf.append(json.dumps({"key": key, "value": value})) + if len(self._buf) >= self.batch_size: + self.flush() + + def send_ndjson(self, ndjson_str): + """Send a pre-built NDJSON string as one frame (fire-and-forget).""" + self._ws.send(ndjson_str) + self._sent_frames += 1 + + def flush(self): + """Send buffered docs as one NDJSON frame.""" + if not self._buf: + return + self._ws.send("\n".join(self._buf)) + self._sent_frames += 1 + self._buf.clear() + + def close(self): + """Flush remaining, close WS, return final summary.""" + self.flush() + self._ws.send_close() + try: + self.result = json.loads(self._ws.recv()) + except Exception: + self.result = {"inserted": 0, "errors": 0, "frames": self._sent_frames} + self._ws.close() + return self.result + + def __enter__(self): + self.open() + return self + + def __exit__(self, *args): + self.close() diff --git a/src/art.zig b/src/art.zig index 7197d21..eaa76d8 100644 --- a/src/art.zig +++ b/src/art.zig @@ -490,9 +490,13 @@ pub const ART = struct { } if (mismatch < prefix_len) { - // Prefix mismatch: split the node - hdr.writeUnlock(); - const new_node = try Node4.create(self.alloc); + // Prefix mismatch: split the node. + // Keep write lock held through the entire operation to prevent + // TOCTOU races (another thread modifying prefix between unlock/relock). + const new_node = Node4.create(self.alloc) catch |err| { + hdr.writeUnlock(); + return err; + }; // New node's prefix is the common part new_node.header.prefix_len = @intCast(mismatch); @@ -508,7 +512,6 @@ pub const ART = struct { } // Adjust existing node's prefix (remove consumed part) - hdr.writeLock(); const remaining = prefix_len - mismatch - 1; if (remaining > 0 and mismatch + 1 < MAX_PREFIX_LEN) { var dst: [MAX_PREFIX_LEN]u8 = [_]u8{0} ** MAX_PREFIX_LEN; diff --git a/src/auth.zig b/src/auth.zig index 10f0550..41694d1 100644 --- a/src/auth.zig +++ b/src/auth.zig @@ -44,34 +44,36 @@ pub const AuthContext = struct { pub const AuthStore = struct { keys: [MAX_KEYS]KeyEntry = undefined, count: u32 = 0, - enabled: bool = false, + enabled: std.atomic.Value(bool) = std.atomic.Value(bool).init(false), lock: std.Thread.RwLock = .{}, - /// Add an API key. Returns the BLAKE3 hash for storage. - pub fn addKey(self: *AuthStore, raw_key: []const u8, name: []const u8, perm: Permission) [32]u8 { + /// Add an API key. Returns the BLAKE3 hash if added, null if at MAX_KEYS. + pub fn addKey(self: *AuthStore, raw_key: []const u8, name: []const u8, perm: Permission) ?[32]u8 { return self.addKeyForTenant(raw_key, name, "default", perm); } - pub fn addKeyForTenant(self: *AuthStore, raw_key: []const u8, name: []const u8, tenant_id: []const u8, perm: Permission) [32]u8 { + /// Add an API key for a specific tenant. Returns the BLAKE3 hash if added, + /// null if the key store is full (MAX_KEYS reached). + pub fn addKeyForTenant(self: *AuthStore, raw_key: []const u8, name: []const u8, tenant_id: []const u8, perm: Permission) ?[32]u8 { self.lock.lock(); defer self.lock.unlock(); + if (self.count >= MAX_KEYS) return null; + const hash = crypto.blake3(raw_key); - if (self.count < MAX_KEYS) { - var entry = KeyEntry{ - .hash = hash, - .name = undefined, - .name_len = @intCast(@min(name.len, 64)), - .tenant_id = undefined, - .tenant_id_len = @intCast(@min(tenant_id.len, 64)), - .perm = perm, - }; - @memcpy(entry.name[0..entry.name_len], name[0..entry.name_len]); - @memcpy(entry.tenant_id[0..entry.tenant_id_len], tenant_id[0..entry.tenant_id_len]); - self.keys[self.count] = entry; - self.count += 1; - self.enabled = true; - } + var entry = KeyEntry{ + .hash = hash, + .name = undefined, + .name_len = @intCast(@min(name.len, 64)), + .tenant_id = undefined, + .tenant_id_len = @intCast(@min(tenant_id.len, 64)), + .perm = perm, + }; + @memcpy(entry.name[0..entry.name_len], name[0..entry.name_len]); + @memcpy(entry.tenant_id[0..entry.tenant_id_len], tenant_id[0..entry.tenant_id_len]); + self.keys[self.count] = entry; + self.count += 1; + self.enabled.store(true, .release); return hash; } @@ -82,7 +84,9 @@ pub const AuthStore = struct { } pub fn resolve(self: *AuthStore, raw_key: []const u8) ?AuthContext { - if (!self.enabled) { + // Fast-path: if no keys have been added, skip hashing and locking. + // Uses acquire to see the latest store from addKeyForTenant. + if (!self.enabled.load(.acquire)) { return AuthContext{ .perm = .admin, .tenant_id = [_]u8{0} ** 64, @@ -108,7 +112,7 @@ pub const AuthStore = struct { /// Check if auth is enabled. pub fn isEnabled(self: *AuthStore) bool { - return self.enabled; + return self.enabled.load(.acquire); } /// Extract API key from HTTP headers. diff --git a/src/branch.zig b/src/branch.zig index 9fa7f88..7ba6226 100644 --- a/src/branch.zig +++ b/src/branch.zig @@ -22,13 +22,12 @@ pub const Branch = struct { agent_id: [64]u8, // which agent owns this branch agent_id_len: u8, - // Branch-local storage: key_hash -> value (only modified keys) + // Branch-local storage: actual key string -> value (only modified keys) // This is the CoW layer — unmodified keys fall through to main - writes: std.AutoHashMap(u64, BranchWrite), + writes: std.StringHashMap(BranchWrite), allocator: Allocator, pub const BranchWrite = struct { - key: []const u8, // owned copy value: []const u8, // owned copy deleted: bool, // true = tombstone (deleted on branch) epoch: u64, // when this write happened @@ -45,25 +44,30 @@ pub const Branch = struct { pub fn deinit(self: *Branch) void { var it = self.writes.iterator(); while (it.next()) |entry| { - if (entry.value_ptr.key.len > 0) self.allocator.free(entry.value_ptr.key); if (entry.value_ptr.value.len > 0) self.allocator.free(entry.value_ptr.value); + self.allocator.free(@constCast(entry.key_ptr.*)); } self.writes.deinit(); } /// Write a key-value pair on this branch (CoW — only stores the delta) pub fn write(self: *Branch, key: []const u8, value: []const u8, epoch: u64) !void { - const key_hash = fnv1a(key); - // Free old write if exists - if (self.writes.getPtr(key_hash)) |old| { - if (old.key.len > 0) self.allocator.free(old.key); - if (old.value.len > 0) self.allocator.free(old.value); - } + // Allocate new copies BEFORE freeing old ones — if alloc fails, + // the existing entry stays valid. const owned_key = try self.allocator.dupe(u8, key); errdefer self.allocator.free(owned_key); const owned_val = try self.allocator.dupe(u8, value); - try self.writes.put(key_hash, .{ - .key = owned_key, + errdefer self.allocator.free(owned_val); + // Free old write if exists (safe — new copies already allocated) + if (self.writes.getPtr(key)) |old| { + if (old.value.len > 0) self.allocator.free(old.value); + // Free the old key that was used as the map key. + const old_map_key = self.writes.getKey(key).?; + self.allocator.free(@constCast(old_map_key)); + // Remove old entry so we can insert with new owned key. + _ = self.writes.remove(key); + } + try self.writes.put(owned_key, .{ .value = owned_val, .deleted = false, .epoch = epoch, @@ -72,14 +76,14 @@ pub const Branch = struct { /// Mark a key as deleted on this branch pub fn delete(self: *Branch, key: []const u8, epoch: u64) !void { - const key_hash = fnv1a(key); - if (self.writes.getPtr(key_hash)) |old| { - if (old.key.len > 0) self.allocator.free(old.key); + if (self.writes.getPtr(key)) |old| { if (old.value.len > 0) self.allocator.free(old.value); + const old_map_key = self.writes.getKey(key).?; + self.allocator.free(@constCast(old_map_key)); + _ = self.writes.remove(key); } const owned_key = try self.allocator.dupe(u8, key); - try self.writes.put(key_hash, .{ - .key = owned_key, + try self.writes.put(owned_key, .{ .value = &.{}, .deleted = true, .epoch = epoch, @@ -88,8 +92,7 @@ pub const Branch = struct { /// Read a key on this branch. Returns branch-local value or null (fall through to main). pub fn read(self: *const Branch, key: []const u8) ?BranchRead { - const key_hash = fnv1a(key); - if (self.writes.get(key_hash)) |w| { + if (self.writes.get(key)) |w| { if (w.deleted) return .{ .deleted = true, .value = null }; return .{ .deleted = false, .value = w.value }; } @@ -108,7 +111,7 @@ pub const Branch = struct { while (it.next()) |entry| { const w = entry.value_ptr.*; try entries.append(alloc, .{ - .key = w.key, + .key = entry.key_ptr.*, .value = w.value, .deleted = w.deleted, .epoch = w.epoch, @@ -165,7 +168,7 @@ pub fn compareBranches(branches: []*const Branch, alloc: Allocator) !CompareResu for (branches) |br| { var it = br.writes.iterator(); while (it.next()) |entry| { - try all_keys.put(entry.value_ptr.key, {}); + try all_keys.put(entry.key_ptr.*, {}); } } @@ -183,7 +186,7 @@ pub fn compareBranches(branches: []*const Branch, alloc: Allocator) !CompareResu .agent_id = br.getAgentId(), .value = r.value, .deleted = r.deleted, - .epoch = if (br.writes.get(fnv1a(key))) |w| w.epoch else 0, + .epoch = if (br.writes.get(key)) |w| w.epoch else 0, }); } } @@ -432,7 +435,7 @@ pub const BranchManager = struct { .status = .active, .agent_id = undefined, .agent_id_len = aid_len, - .writes = std.AutoHashMap(u64, Branch.BranchWrite).init(self.allocator), + .writes = std.StringHashMap(Branch.BranchWrite).init(self.allocator), .allocator = self.allocator, }; @memcpy(branch.name[0..name_len], name_arg[0..name_len]); diff --git a/src/bwtree.zig b/src/bwtree.zig index 2789e13..50817c2 100644 --- a/src/bwtree.zig +++ b/src/bwtree.zig @@ -49,9 +49,11 @@ pub const BwTree = struct { next_page_id: std.atomic.Value(usize), allocator: std.mem.Allocator, // Deferred reclamation: old chains retired after consolidation are parked here - // and freed on the next consolidation or deinit (simple two-phase approach). + // and freed once they are at least 2 epochs old (epoch-based safe reclamation). retired: std.ArrayList(*Page), + retired_epochs: std.ArrayList(u64), retired_mu: std.Thread.Mutex, + epoch: u64, pub fn init(allocator: std.mem.Allocator) !BwTree { var tree: BwTree = undefined; @@ -59,7 +61,9 @@ pub const BwTree = struct { tree.root_pid = 0; tree.next_page_id = std.atomic.Value(usize).init(1); tree.retired = std.ArrayList(*Page).init(allocator); + tree.retired_epochs = std.ArrayList(u64).init(allocator); tree.retired_mu = .{}; + tree.epoch = 0; // Zero out all mapping slots var i: usize = 0; @@ -78,9 +82,10 @@ pub const BwTree = struct { } pub fn deinit(self: *BwTree) void { - // Free all retired chains first - self.drainRetired(); + // Free all retired chains unconditionally on shutdown + self.drainAllRetired(); self.retired.deinit(); + self.retired_epochs.deinit(); var i: usize = 0; while (i < MAX_PAGES) : (i += 1) { const ptr_val = self.mapping[i].load(.acquire); @@ -107,12 +112,32 @@ pub const BwTree = struct { /// Drain the retired list — frees chains that were parked on previous consolidations. fn drainRetired(self: *BwTree) void { + self.retired_mu.lock(); + defer self.retired_mu.unlock(); + const current_epoch = self.epoch; + // Free entries that are at least 2 epochs old -- any concurrent reader + // that started before the retirement has completed by now. + var i: usize = 0; + while (i < self.retired.items.len) { + if (current_epoch -| self.retired_epochs.items[i] >= 2) { + self.freeChain(self.retired.items[i]); + _ = self.retired.swapRemove(i); + _ = self.retired_epochs.swapRemove(i); + } else { + i += 1; + } + } + } + + /// Drain ALL retired entries unconditionally (used by deinit). + fn drainAllRetired(self: *BwTree) void { self.retired_mu.lock(); defer self.retired_mu.unlock(); for (self.retired.items) |page| { self.freeChain(page); } self.retired.clearRetainingCapacity(); + self.retired_epochs.clearRetainingCapacity(); } // ─── allocPage ─────────────────────────────────────────────────────── @@ -206,6 +231,7 @@ pub const BwTree = struct { } else { // CAS failed — another thread won; free our delta and retry self.allocator.destroy(new_delta); + std.atomic.spinLoopHint(); } } } @@ -237,6 +263,7 @@ pub const BwTree = struct { return; } else { self.allocator.destroy(new_delta); + std.atomic.spinLoopHint(); } } } @@ -245,8 +272,10 @@ pub const BwTree = struct { /// When delta chain exceeds MAX_DELTA_CHAIN, merge into a new base page. pub fn consolidate(self: *BwTree, page_id: usize) void { - // Drain previously retired chains — they've survived at least one full - // consolidation cycle, so readers from the previous epoch are done. + // Advance global epoch so drainRetired can age out old entries. + self.epoch +%= 1; + + // Drain previously retired chains that are old enough. self.drainRetired(); const old = self.mapping[page_id].load(.acquire); @@ -290,11 +319,16 @@ pub const BwTree = struct { ) == null) { // Success — park old chain head for deferred reclamation. // Concurrent readers may still be traversing it; it will be freed - // on the next consolidation cycle (two-phase epoch approach). + // once the entry is at least 2 epochs old. self.retired_mu.lock(); defer self.retired_mu.unlock(); self.retired.append(page) catch { - // If we can't track it, leak it — better than use-after-free. + // If we can't track it, leak it -- better than use-after-free. + return; + }; + self.retired_epochs.append(self.epoch) catch { + // Keep arrays in sync: undo the page append on epoch failure. + _ = self.retired.pop(); }; } else { // Another thread consolidated first; discard our work diff --git a/src/cdc.zig b/src/cdc.zig index 65fe906..b017cfb 100644 --- a/src/cdc.zig +++ b/src/cdc.zig @@ -13,12 +13,13 @@ pub const Event = struct { tenant_id_len: u8, collection: [64]u8, collection_len: u8, - key: [128]u8, - key_len: u8, - value: [512]u8, + key: [256]u8, + key_len: u16, + value: [4096]u8, value_len: u16, doc_id: u64, op: Op, + truncated: bool, pub fn tenant(self: *const Event) []const u8 { return self.tenant_id[0..self.tenant_id_len]; @@ -75,7 +76,7 @@ pub const Delivery = struct { webhook_url: [256]u8, webhook_url_len: u16, signature_hex: [64]u8, - payload: [1024]u8, + payload: [8192]u8, payload_len: u16, pub fn tenant(self: *const Delivery) []const u8 { @@ -102,6 +103,7 @@ pub const CDCManager = struct { deliveries: std.ArrayList(Delivery), next_subscription_id: std.atomic.Value(u64), next_seq: std.atomic.Value(u64), + events_dropped: std.atomic.Value(u64), mu: std.Thread.Mutex, cond: std.Thread.Condition, running: std.atomic.Value(bool), @@ -115,6 +117,7 @@ pub const CDCManager = struct { .deliveries = .empty, .next_subscription_id = std.atomic.Value(u64).init(1), .next_seq = std.atomic.Value(u64).init(1), + .events_dropped = std.atomic.Value(u64).init(0), .mu = .{}, .cond = .{}, .running = std.atomic.Value(bool).init(false), @@ -143,6 +146,8 @@ pub const CDCManager = struct { self.worker = null; } + const MAX_SUBSCRIPTIONS: usize = 1024; + pub fn registerWebhook(self: *CDCManager, tenant_id: []const u8, collection: []const u8, webhook_url: []const u8, secret: []const u8) !u64 { var sub = std.mem.zeroes(Subscription); sub.id = self.next_subscription_id.fetchAdd(1, .monotonic); @@ -153,23 +158,41 @@ pub const CDCManager = struct { self.mu.lock(); defer self.mu.unlock(); + if (self.subscriptions.items.len >= MAX_SUBSCRIPTIONS) return error.TooManySubscriptions; try self.subscriptions.append(self.allocator, sub); return sub.id; } + pub fn removeWebhook(self: *CDCManager, sub_id: u64) bool { + self.mu.lock(); + defer self.mu.unlock(); + for (self.subscriptions.items, 0..) |sub, i| { + if (sub.id == sub_id) { + _ = self.subscriptions.swapRemove(i); + return true; + } + } + return false; + } + pub fn emit(self: *CDCManager, tenant_id: []const u8, collection: []const u8, key: []const u8, value: []const u8, doc_id: u64, op: Op) void { var ev = std.mem.zeroes(Event); ev.seq = self.next_seq.fetchAdd(1, .monotonic); ev.doc_id = doc_id; ev.op = op; + ev.truncated = false; + const key_trunc = fillTracked(&ev.key, &ev.key_len, key); + const val_trunc = fillU16Tracked(&ev.value, &ev.value_len, value); + if (key_trunc or val_trunc) ev.truncated = true; fill(&ev.tenant_id, &ev.tenant_id_len, tenant_id); fill(&ev.collection, &ev.collection_len, collection); - fill(&ev.key, &ev.key_len, key); - fillU16(&ev.value, &ev.value_len, value); self.mu.lock(); defer self.mu.unlock(); - self.pending.append(self.allocator, ev) catch return; + self.pending.append(self.allocator, ev) catch { + _ = self.events_dropped.fetchAdd(1, .monotonic); + return; + }; self.cond.signal(); } @@ -197,20 +220,34 @@ pub const CDCManager = struct { self.mu.unlock(); return; } - const ev = self.pending.orderedRemove(0); - const subs = self.subscriptions.items; + // Drain all pending events in one batch to avoid O(n) orderedRemove(0) per event. + var batch = self.pending; + self.pending = .empty; + const sub_count = self.subscriptions.items.len; self.mu.unlock(); - for (subs) |sub| { - if (!matches(sub, ev)) continue; - const delivery = makeDelivery(sub, ev); - self.mu.lock(); - self.deliveries.append(self.allocator, delivery) catch {}; - if (self.deliveries.items.len > 4096) { - _ = self.deliveries.orderedRemove(0); + for (batch.items) |ev| { + var si: usize = 0; + while (si < sub_count) : (si += 1) { + self.mu.lock(); + if (si >= self.subscriptions.items.len) { + self.mu.unlock(); + break; + } + const sub = self.subscriptions.items[si]; + self.mu.unlock(); + + if (!matches(sub, ev)) continue; + const delivery = makeDelivery(sub, ev); + self.mu.lock(); + self.deliveries.append(self.allocator, delivery) catch {}; + if (self.deliveries.items.len > 4096) { + _ = self.deliveries.orderedRemove(0); + } + self.mu.unlock(); } - self.mu.unlock(); } + batch.deinit(self.allocator); } } }; @@ -228,11 +265,24 @@ fn makeDelivery(sub: Subscription, ev: Event) Delivery { .update => "update", .delete => "delete", }; + + // Escape user-controlled fields to prevent JSON injection + var esc_tenant_buf: [256]u8 = undefined; + var esc_col_buf: [256]u8 = undefined; + var esc_key_buf: [1024]u8 = undefined; + const esc_tenant = cdcJsonEscape(ev.tenant(), &esc_tenant_buf); + const esc_col = cdcJsonEscape(ev.collectionName(), &esc_col_buf); + const esc_key = cdcJsonEscape(ev.keySlice(), &esc_key_buf); + const payload = std.fmt.bufPrint( &delivery.payload, "{{\"seq\":{d},\"tenant\":\"{s}\",\"collection\":\"{s}\",\"op\":\"{s}\",\"key\":\"{s}\",\"doc_id\":{d},\"value\":{s}}}", - .{ ev.seq, ev.tenant(), ev.collectionName(), op_str, ev.keySlice(), ev.doc_id, if (ev.value_len > 0) ev.valueSlice() else "null" }, - ) catch "{}"; + .{ ev.seq, esc_tenant, esc_col, op_str, esc_key, ev.doc_id, if (ev.value_len > 0) ev.valueSlice() else "null" }, + ) catch std.fmt.bufPrint( + &delivery.payload, + "{{\"seq\":{d},\"tenant\":\"{s}\",\"collection\":\"{s}\",\"op\":\"{s}\",\"key\":\"{s}\",\"doc_id\":{d},\"value\":null,\"truncated\":true}}", + .{ ev.seq, esc_tenant, esc_col, op_str, esc_key, ev.doc_id }, + ) catch return delivery; delivery.payload_len = @intCast(payload.len); delivery.signature_hex = crypto.hmacSha256Hex(sub.secretSlice(), payload); return delivery; @@ -256,6 +306,44 @@ fn fillU16(dest: anytype, len_ptr: *u16, src: []const u8) void { len_ptr.* = @intCast(n); } +/// Like fill, but returns true if src was truncated to fit dest. +fn fillTracked(dest: anytype, len_ptr: anytype, src: []const u8) bool { + const n = @min(dest.len, src.len); + @memcpy(dest[0..n], src[0..n]); + len_ptr.* = @intCast(n); + return src.len > dest.len; +} + +/// Like fillU16, but returns true if src was truncated to fit dest. +fn fillU16Tracked(dest: anytype, len_ptr: *u16, src: []const u8) bool { + const n = @min(dest.len, src.len); + @memcpy(dest[0..n], src[0..n]); + len_ptr.* = @intCast(n); + return src.len > dest.len; +} + +/// Escape a string for safe JSON interpolation (handles " and \). +fn cdcJsonEscape(raw: []const u8, buf: []u8) []const u8 { + var needs_escape = false; + for (raw) |ch| { + if (ch == '"' or ch == '\\') { needs_escape = true; break; } + } + if (!needs_escape) return raw; + var pos: usize = 0; + for (raw) |ch| { + if (ch == '"' or ch == '\\') { + if (pos + 2 > buf.len) return raw; + buf[pos] = '\\'; + pos += 1; + } else { + if (pos + 1 > buf.len) return raw; + } + buf[pos] = ch; + pos += 1; + } + return buf[0..pos]; +} + test "cdc filters by tenant and collection and signs payload" { const alloc = std.testing.allocator; var cdc = CDCManager.init(alloc); diff --git a/src/codeindex.zig b/src/codeindex.zig index 8d4e0fc..47795ec 100644 --- a/src/codeindex.zig +++ b/src/codeindex.zig @@ -14,6 +14,13 @@ pub const WordIndex = struct { /// path → set of words contributed (for efficient re-index cleanup) file_words: std.StringHashMap(std.StringHashMap(void)), allocator: std.mem.Allocator, + /// Mutex for concurrent access (background indexer writes, search reads). + mu: std.Thread.Mutex = .{}, + + /// Cap hits per word to bound memory. Common words ("the", "var", "if") + /// accumulate thousands of hits — beyond this they waste memory for + /// negligible search value. + const MAX_HITS_PER_WORD: usize = 5_000; pub fn init(allocator: std.mem.Allocator) WordIndex { return .{ @@ -77,9 +84,11 @@ pub const WordIndex = struct { /// Index a file's content — tokenizes into words and records hits. pub fn indexFile(self: *WordIndex, path: []const u8, content: []const u8) !void { + if (content.len > TrigramIndex.MAX_INDEX_FILE_SIZE) return; + self.mu.lock(); + defer self.mu.unlock(); // Clean up old entries first self.removeFile(path); - // Dupe path so hits survive after the caller's buffer is freed. const owned_path = try self.allocator.dupe(u8, path); @@ -102,6 +111,9 @@ pub const WordIndex = struct { gop.value_ptr.* = .{}; } + // Skip overly common words to bound memory. + if (gop.value_ptr.items.len >= MAX_HITS_PER_WORD) continue; + if (gop.value_ptr.items.len > 0) { const last = gop.value_ptr.items[gop.value_ptr.items.len - 1]; if (std.mem.eql(u8, last.path, owned_path) and last.line_num == line_num) { @@ -129,43 +141,54 @@ pub const WordIndex = struct { try self.file_words.put(owned_path, words_set); } - /// Look up all hits for a word. O(1) lookup + O(hits) iteration. - pub fn search(self: *WordIndex, word: []const u8) []const WordHit { + /// Look up all hits for a word. Returns an owned copy that the caller + /// must free with `allocator.free(result)`. The copy is made while the + /// mutex is held so the internal ArrayList buffer cannot be reallocated + /// by a concurrent indexFile() call. + pub fn search(self: *WordIndex, word: []const u8, allocator: std.mem.Allocator) ![]const WordHit { + self.mu.lock(); + defer self.mu.unlock(); if (self.index.get(word)) |hits| { - return hits.items; + const copy = try allocator.alloc(WordHit, hits.items.len); + @memcpy(copy, hits.items); + return copy; } return &.{}; } /// Look up hits, returning results allocated by the caller. - /// Deduplicates by (path, line_num). -pub fn searchDeduped(self: *WordIndex, word: []const u8, allocator: std.mem.Allocator) ![]const WordHit { - const hits = self.search(word); - if (hits.len == 0) return try allocator.alloc(WordHit, 0); - if (hits.len == 1) { - var out = try allocator.alloc(WordHit, 1); - out[0] = hits[0]; - return out; - } + /// Deduplicates by (path, line_num). The mutex is held for the + /// duration of the read so the internal buffer stays valid. + pub fn searchDeduped(self: *WordIndex, word: []const u8, allocator: std.mem.Allocator) ![]const WordHit { + self.mu.lock(); + defer self.mu.unlock(); - const DedupKey = struct { path_ptr: usize, line_num: u32 }; - var seen = std.AutoHashMap(DedupKey, void).init(allocator); - defer seen.deinit(); - try seen.ensureTotalCapacity(@intCast(hits.len)); + const items = if (self.index.get(word)) |hits| hits.items else return try allocator.alloc(WordHit, 0); + if (items.len == 0) return try allocator.alloc(WordHit, 0); + if (items.len == 1) { + var out = try allocator.alloc(WordHit, 1); + out[0] = items[0]; + return out; + } - var result: std.ArrayList(WordHit) = .{}; - errdefer result.deinit(allocator); - try result.ensureTotalCapacity(allocator, hits.len); + const DedupKey = struct { path_ptr: usize, line_num: u32 }; + var seen = std.AutoHashMap(DedupKey, void).init(allocator); + defer seen.deinit(); + try seen.ensureTotalCapacity(@intCast(items.len)); + + var result: std.ArrayList(WordHit) = .{}; + errdefer result.deinit(allocator); + try result.ensureTotalCapacity(allocator, items.len); - for (hits) |hit| { - const key = DedupKey{ .path_ptr = @intFromPtr(hit.path.ptr), .line_num = hit.line_num }; - const gop = try seen.getOrPut(key); - if (!gop.found_existing) { - result.appendAssumeCapacity(hit); + for (items) |hit| { + const key = DedupKey{ .path_ptr = @intFromPtr(hit.path.ptr), .line_num = hit.line_num }; + const gop = try seen.getOrPut(key); + if (!gop.found_existing) { + result.appendAssumeCapacity(hit); + } } + return result.toOwnedSlice(allocator); } - return result.toOwnedSlice(allocator); -} }; // ── Trigram index ─────────────────────────────────────────── @@ -265,6 +288,18 @@ pub const TrigramIndex = struct { /// Mutex for concurrent access (background indexer writes, search reads). mu: std.Thread.Mutex = .{}, + /// Cap posting list size per trigram. Once a trigram appears in this many files + /// it has no discrimination value for search — stop tracking it to bound memory. + const MAX_FILES_PER_TRIGRAM: usize = 10_000; + + /// Cap unique trigrams stored per file. Large generated/minified files can produce + /// 10K+ unique trigrams — beyond this limit the memory cost outweighs search value. + const MAX_TRIGRAMS_PER_FILE: usize = 8_000; + + /// Skip files larger than this for trigram indexing. Minified JS, generated code, + /// and data files waste memory for negligible search value. + pub const MAX_INDEX_FILE_SIZE: usize = 256 * 1024; + pub fn init(allocator: std.mem.Allocator) TrigramIndex { return .{ .index = std.AutoHashMap(Trigram, PostingList).init(allocator), @@ -290,9 +325,9 @@ pub const TrigramIndex = struct { } self.file_trigrams.deinit(); - // Free owned path strings in id_to_path + // Free owned path strings in id_to_path (skip empty slots from removed files) for (self.id_to_path.items) |p| { - self.allocator.free(p); + if (p.len > 0) self.allocator.free(p); } self.id_to_path.deinit(self.allocator); @@ -330,6 +365,16 @@ pub const TrigramIndex = struct { } trigrams.deinit(self.allocator); _ = self.file_trigrams.remove(file_id); + + // Clean up path_to_id and id_to_path mappings + if (file_id < self.id_to_path.items.len) { + const path = self.id_to_path.items[file_id]; + if (path.len > 0) { + _ = self.path_to_id.remove(path); + self.allocator.free(path); + self.id_to_path.items[file_id] = &.{}; + } + } } /// Legacy wrapper for callers that pass a path string. @@ -340,6 +385,7 @@ pub const TrigramIndex = struct { pub fn indexFile(self: *TrigramIndex, path: []const u8, content: []const u8) !void { if (content.len < 3) return; + if (content.len > MAX_INDEX_FILE_SIZE) return; // ── Pass 1 (no lock): collect unique trigrams into a temporary hashmap ── var unique_tris = std.AutoHashMap(Trigram, void).init(self.allocator); @@ -355,6 +401,8 @@ pub const TrigramIndex = struct { normalizeChar(content[i + 2]), ); unique_tris.put(tri, {}) catch {}; + // Stop collecting if we hit the per-file cap + if (unique_tris.count() >= MAX_TRIGRAMS_PER_FILE) break; } // ── Pass 2 (locked): update the global index once per unique trigram ── @@ -375,12 +423,15 @@ pub const TrigramIndex = struct { var it = unique_tris.keyIterator(); while (it.next()) |tri_ptr| { const tri = tri_ptr.*; - tri_list.appendAssumeCapacity(tri); const idx_gop = try self.index.getOrPut(tri); if (!idx_gop.found_existing) { idx_gop.value_ptr.* = .{}; } + // Skip trigrams that already appear in too many files + if (idx_gop.value_ptr.items.len >= MAX_FILES_PER_TRIGRAM) continue; + + tri_list.appendAssumeCapacity(tri); postingSortedInsert(idx_gop.value_ptr, self.allocator, .{ .file_id = file_id, .mask = PostingMask{ .loc_mask = 0xFF, .next_mask = 0xFF } }); } @@ -400,11 +451,18 @@ pub const TrigramIndex = struct { // Per-doc trigram sets collected outside the lock const BatchEntry = struct { tris: []Trigram }; var batch: [64]BatchEntry = undefined; + const batch_len = @min(paths.len, 64); - for (paths, contents, 0..) |path, content, di| { + var allocated_count: usize = 0; + errdefer for (0..allocated_count) |i| { + if (batch[i].tris.len > 0) self.allocator.free(batch[i].tris); + }; + + for (paths[0..batch_len], contents[0..batch_len], 0..) |path, content, di| { _ = path; - if (content.len < 3) { + if (content.len < 3 or content.len > MAX_INDEX_FILE_SIZE) { batch[di].tris = &.{}; + allocated_count += 1; continue; } reusable_tris.clearRetainingCapacity(); @@ -418,6 +476,7 @@ pub const TrigramIndex = struct { normalizeChar(content[i + 2]), ); reusable_tris.put(tri, {}) catch {}; + if (reusable_tris.count() >= MAX_TRIGRAMS_PER_FILE) break; } // Copy unique trigrams to owned slice for this doc @@ -429,13 +488,14 @@ pub const TrigramIndex = struct { idx += 1; } batch[di].tris = tris; + allocated_count += 1; } // Single lock acquisition for all docs self.mu.lock(); defer self.mu.unlock(); - for (paths, contents, 0..) |path, _, di| { + for (paths[0..batch_len], contents[0..batch_len], 0..) |path, _, di| { const tris = batch[di].tris; defer if (tris.len > 0) self.allocator.free(tris); if (tris.len == 0) continue; @@ -450,15 +510,16 @@ pub const TrigramIndex = struct { tri_list.ensureTotalCapacity(self.allocator, tris.len) catch continue; for (tris) |tri| { - tri_list.appendAssumeCapacity(tri); - const idx_gop = self.index.getOrPut(tri) catch continue; if (!idx_gop.found_existing) { idx_gop.value_ptr.* = .{}; } + // Skip trigrams that already appear in too many files + if (idx_gop.value_ptr.items.len >= MAX_FILES_PER_TRIGRAM) continue; + + tri_list.appendAssumeCapacity(tri); postingSortedInsert(idx_gop.value_ptr, self.allocator, .{ .file_id = file_id, .mask = PostingMask{ .loc_mask = 0xFF, .next_mask = 0xFF } }); } - self.file_trigrams.put(file_id, tri_list) catch {}; } } @@ -1582,11 +1643,13 @@ pub fn buildCoveringSet(query: []const u8, allocator: std.mem.Allocator) ![]Spar /// In-memory sparse n-gram index. Mirrors the TrigramIndex API so it can /// be used as a drop-in acceleration layer alongside the trigram index. pub const SparseNgramIndex = struct { - /// ngram hash → set of file paths that contain the n-gram + /// ngram hash -> set of file paths that contain the n-gram index: std.AutoHashMap(u64, std.StringHashMap(void)), - /// path → list of ngram hashes contributed (for cleanup on re-index) + /// path -> list of ngram hashes contributed (for cleanup on re-index) file_ngrams: std.StringHashMap(std.ArrayList(u64)), allocator: std.mem.Allocator, + /// Mutex for concurrent access. + mu: std.Thread.Mutex = .{}, pub fn init(allocator: std.mem.Allocator) SparseNgramIndex { return .{ @@ -1599,6 +1662,11 @@ pub const SparseNgramIndex = struct { pub fn deinit(self: *SparseNgramIndex) void { var iter = self.index.iterator(); while (iter.next()) |entry| { + // Free owned path keys inside each file_set + var key_iter = entry.value_ptr.keyIterator(); + while (key_iter.next()) |k| { + self.allocator.free(k.*); + } entry.value_ptr.deinit(); } self.index.deinit(); @@ -1611,10 +1679,19 @@ pub const SparseNgramIndex = struct { } pub fn removeFile(self: *SparseNgramIndex, path: []const u8) void { + self.mu.lock(); + defer self.mu.unlock(); + self.removeFileInner(path); + } + + fn removeFileInner(self: *SparseNgramIndex, path: []const u8) void { const ngrams = self.file_ngrams.getPtr(path) orelse return; for (ngrams.items) |hash| { if (self.index.getPtr(hash)) |file_set| { - _ = file_set.remove(path); + // Find and free the owned key before removing from the set + if (file_set.fetchRemove(path)) |kv| { + self.allocator.free(kv.key); + } if (file_set.count() == 0) { file_set.deinit(); _ = self.index.remove(hash); @@ -1626,7 +1703,10 @@ pub const SparseNgramIndex = struct { } pub fn indexFile(self: *SparseNgramIndex, path: []const u8, content: []const u8) !void { - self.removeFile(path); + self.mu.lock(); + defer self.mu.unlock(); + + self.removeFileInner(path); const ngrams = try extractSparseNgrams(content, self.allocator); defer self.allocator.free(ngrams); @@ -1640,7 +1720,12 @@ pub const SparseNgramIndex = struct { if (!gop.found_existing) { gop.value_ptr.* = std.StringHashMap(void).init(self.allocator); } - _ = try gop.value_ptr.getOrPut(path); + // Dupe path so the index owns it + if (!gop.value_ptr.contains(path)) { + const owned_path = try self.allocator.dupe(u8, path); + errdefer self.allocator.free(owned_path); + try gop.value_ptr.put(owned_path, {}); + } _ = try seen.getOrPut(ng.hash); } @@ -1655,10 +1740,13 @@ pub const SparseNgramIndex = struct { /// Find candidate files that may contain the query string. /// Uses the sliding-window covering set from buildCoveringSet and returns - /// the UNION of all matching posting lists — a superset of true matches, + /// the UNION of all matching posting lists -- a superset of true matches, /// to be verified by content search. Returns null when the query is too /// short. Caller must free the returned slice. pub fn candidates(self: *SparseNgramIndex, query: []const u8, allocator: std.mem.Allocator) ?[]const []const u8 { + self.mu.lock(); + defer self.mu.unlock(); + const ngrams = buildCoveringSet(query, allocator) catch return null; defer allocator.free(ngrams); @@ -1700,3 +1788,129 @@ pub const SparseNgramIndex = struct { return @intCast(self.file_ngrams.count()); } }; + +// --- Tests ---------------------------------------------------------------- + +test "trigram index skips files exceeding MAX_INDEX_FILE_SIZE" { + const alloc = std.testing.allocator; + var idx = TrigramIndex.init(alloc); + defer idx.deinit(); + + // Small file: should be indexed + try idx.indexFile("small.zig", "fn main() void {}"); + try std.testing.expect(idx.fileCount() == 1); + + // Large file (>256KB): should be skipped + const big = try alloc.alloc(u8, TrigramIndex.MAX_INDEX_FILE_SIZE + 1); + defer alloc.free(big); + @memset(big, 'a'); + try idx.indexFile("big.bin", big); + // File count should still be 1 — big file was skipped + try std.testing.expect(idx.fileCount() == 1); +} + +test "trigram index caps unique trigrams per file at MAX_TRIGRAMS_PER_FILE" { + const alloc = std.testing.allocator; + var idx = TrigramIndex.init(alloc); + defer idx.deinit(); + + // Generate content with many unique trigrams: cycling through all printable ASCII + // Each 3-byte window is unique if we use a non-repeating sequence long enough. + const len = 40_000; // would produce ~40K raw trigrams, many unique + const content = try alloc.alloc(u8, len); + defer alloc.free(content); + for (0..len) |i| { + content[i] = @intCast(32 + (i % 95)); // printable ASCII cycle + } + + try idx.indexFile("many_trigrams.txt", content); + // The file_trigrams entry should be capped + const file_id = idx.path_to_id.get("many_trigrams.txt").?; + const tri_list = idx.file_trigrams.get(file_id).?; + try std.testing.expect(tri_list.items.len <= TrigramIndex.MAX_TRIGRAMS_PER_FILE); +} + +test "trigram posting list caps at MAX_FILES_PER_TRIGRAM" { + const alloc = std.testing.allocator; + var idx = TrigramIndex.init(alloc); + defer idx.deinit(); + + // Index many files with the same content to force one trigram's posting list to grow + const content = "aaa"; // only one unique trigram: "aaa" + const N = TrigramIndex.MAX_FILES_PER_TRIGRAM + 100; + var name_buf: [32]u8 = undefined; + for (0..N) |i| { + const name_len = std.fmt.formatIntBuf(&name_buf, i, 10, .lower, .{}); + try idx.indexFile(name_buf[0..name_len], content); + } + + // The posting list for trigram "aaa" should be capped + const tri = packTrigram('a', 'a', 'a'); + const posting = idx.index.get(tri).?; + try std.testing.expect(posting.items.len <= TrigramIndex.MAX_FILES_PER_TRIGRAM); +} + +test "word index skips files exceeding MAX_INDEX_FILE_SIZE" { + const alloc = std.testing.allocator; + var wi = WordIndex.init(alloc); + defer wi.deinit(); + + // Small file: should index words + try wi.indexFile("small.zig", "hello world function"); + const hits = try wi.search("hello", alloc); + defer if (hits.len > 0) alloc.free(hits); + try std.testing.expect(hits.len > 0); + + // Large file: should be skipped entirely + const big = try alloc.alloc(u8, TrigramIndex.MAX_INDEX_FILE_SIZE + 1); + defer alloc.free(big); + @memset(big, 'x'); + // Put a unique word near the start to verify it's NOT indexed + @memcpy(big[0..11], "uniquetoken"); + try wi.indexFile("big.bin", big); + const big_hits = try wi.search("uniquetoken", alloc); + defer if (big_hits.len > 0) alloc.free(big_hits); + try std.testing.expect(big_hits.len == 0); +} + +test "trigram indexBatch respects file size and per-file caps" { + const alloc = std.testing.allocator; + var idx = TrigramIndex.init(alloc); + defer idx.deinit(); + + // Create one normal file and one oversized file + const small_content = "fn main() void { return; }"; + const big_content = try alloc.alloc(u8, TrigramIndex.MAX_INDEX_FILE_SIZE + 1); + defer alloc.free(big_content); + @memset(big_content, 'z'); + + const paths = [_][]const u8{ "ok.zig", "toobig.bin" }; + const contents = [_][]const u8{ small_content, big_content }; + var reusable = std.AutoHashMap(Trigram, void).init(alloc); + defer reusable.deinit(); + + try idx.indexBatch(&paths, &contents, &reusable); + + // Only the small file should be indexed + try std.testing.expect(idx.path_to_id.contains("ok.zig")); + try std.testing.expect(!idx.path_to_id.contains("toobig.bin")); +} + +test "word index MAX_HITS_PER_WORD prevents unbounded growth" { + const alloc = std.testing.allocator; + var wi = WordIndex.init(alloc); + defer wi.deinit(); + + // Index many files that all contain the same word + var name_buf: [32]u8 = undefined; + const N = WordIndex.MAX_HITS_PER_WORD + 500; + for (0..N) |i| { + const name_len = std.fmt.formatIntBuf(&name_buf, i, 10, .lower, .{}); + const content = "common_word appears here"; + try wi.indexFile(name_buf[0..name_len], content); + } + + const hits = try wi.search("common_word", alloc); + defer if (hits.len > 0) alloc.free(hits); + try std.testing.expect(hits.len <= WordIndex.MAX_HITS_PER_WORD); +} diff --git a/src/collection.zig b/src/collection.zig index ac78fbb..7945a70 100644 --- a/src/collection.zig +++ b/src/collection.zig @@ -30,7 +30,7 @@ pub const MAX_COLLECTION_NAME_LEN: usize = 64; // Lock-free MPSC queue for deferring trigram+word index builds to a background thread. // Insert path pushes owned (key, value) pairs; background thread pops and indexes. pub const IndexQueue = struct { - const CAPACITY = 32768; // 32K entries per queue + const CAPACITY = 16384; // 16K entries — balances burst capacity vs memory const Entry = struct { key: []const u8, @@ -63,6 +63,7 @@ pub const IndexQueue = struct { if (next == self.tail.load(.acquire)) return false; // full // Atomically claim this slot by advancing head. if (self.head.cmpxchgWeak(h, next, .acq_rel, .monotonic)) |_| { + std.atomic.spinLoopHint(); continue; // CAS failed — another producer won, retry } // We own slot h — write data then signal readiness. @@ -162,9 +163,9 @@ fn extractJsonFloatArray(json: []const u8, field_name: []const u8, out: []f32) ? } // ─── Collection ────────────────────────────────────────────────────────── pub const Collection = struct { - pub const STRIPE_COUNT = 1024; + pub const STRIPE_COUNT = 256; // MVCC GC: trigger version chain cleanup every GC_INTERVAL inserts. - const GC_INTERVAL: u64 = 10_000; + const GC_INTERVAL: u64 = 500; name_buf: [128]u8, name_len: u8, @@ -176,9 +177,12 @@ pub const Collection = struct { epochs: *EpochManager, cdc: *cdc_mod.CDCManager, next_doc_id: std.atomic.Value(u64), - // Record-level latching: 1024 stripe locks replace single write_mu. - // Two writers to different keys proceed in parallel; same-key writes serialize. + // Record-level latching: stripe locks serialize same-key writes. + // shared_mu protects collection-wide structures (hash_idx, key_doc_ids, + // key_epochs, versions) that are accessed by ALL stripes — without this, + // concurrent inserts to different keys race on HashMap internals. stripe_locks: [STRIPE_COUNT]std.Thread.Mutex, + shared_mu: std.Thread.Mutex = .{}, hash_idx: std.AutoHashMap(u64, BTreeEntry), key_doc_ids: std.AutoHashMap(u64, u64), /// Per-key modification epoch — tracks when each key was last written on main. @@ -224,6 +228,10 @@ pub const Collection = struct { var path_buf: [512]u8 = undefined; const col = try alloc.create(Collection); errdefer alloc.destroy(col); + // Zero all bytes so that Mutex fields (os_unfair_lock) start as 0 + // before any field assignment. GPA fills with 0xAA which macOS + // detects as corrupted unfair_lock state. + @memset(std.mem.asBytes(col), 0); try ensureDirPath(data_dir); const page_path = try std.fmt.bufPrint(&path_buf, "{s}/{s}.pages", .{ data_dir, storage_name }); @@ -241,8 +249,11 @@ pub const Collection = struct { lock.* = .{}; } col.hash_idx = std.AutoHashMap(u64, BTreeEntry).init(alloc); + col.hash_idx.ensureTotalCapacity(4096) catch {}; col.key_doc_ids = std.AutoHashMap(u64, u64).init(alloc); + col.key_doc_ids.ensureTotalCapacity(4096) catch {}; col.key_epochs = std.AutoHashMap(u64, u64).init(alloc); + col.key_epochs.ensureTotalCapacity(4096) catch {}; col.alloc = alloc; col.cache = hot_cache.HotCache.init(); col.versions = mvcc_mod.VersionChain.init(alloc); @@ -261,10 +272,10 @@ pub const Collection = struct { col.name_len = @intCast(n); try col.rebuildIndexes(); - // Start background indexer thread. If spawn fails, inserts fall back to sync + // Start background indexer threads. If spawn fails, inserts fall back to sync // indexing via the queue-full path (queue stays empty, retries exhaust immediately). col.index_thread = std.Thread.spawn(.{}, indexWorkerQ, .{ col, &col.index_queue }) catch null; - col.index_thread2 = null; // Second worker slot — available for future use + col.index_thread2 = std.Thread.spawn(.{}, indexWorkerQ, .{ col, &col.index_queue2 }) catch null; return col; } @@ -339,7 +350,7 @@ pub const Collection = struct { } pub fn storageBytes(self: *const Collection) usize { - return self.pf.mm.len; + return self.pf.mm.dataLen(); } // ─── MVCC read transaction ────────────────────────────────────────── @@ -355,6 +366,8 @@ pub const Collection = struct { /// Insert a new document. Returns assigned doc_id. /// Uses per-key stripe lock for fine-grained concurrency. pub fn insert(self: *Collection, key: []const u8, value: []const u8) !u64 { + if (key.len == 0 or key.len > std.math.maxInt(u16)) return error.InvalidKeyLength; + if (value.len > std.math.maxInt(u32)) return error.InvalidValueLength; const key_hash = doc_mod.fnv1a(key); const stripe = stripeIndex(key_hash); self.stripe_locks[stripe].lock(); @@ -364,27 +377,28 @@ pub const Collection = struct { const hdr = doc_mod.newHeader(doc_id, key, value); const d = Doc{ .header = hdr, .key = key, .value = value }; - // Use stack buffer for small docs, heap for large ones (code files can be 100KB+) - var enc_buf: [65536]u8 = undefined; const total_size = DocHeader.size + key.len + value.len; - var heap_buf: ?[]u8 = null; - defer if (heap_buf) |hb| self.alloc.free(hb); - const buf = if (total_size <= enc_buf.len) - &enc_buf - else blk: { - heap_buf = try self.alloc.alloc(u8, total_size + 64); - break :blk heap_buf.?; - }; - const enc = try d.encodeBuf(buf); // Write to WAL buffer (background flusher will commit periodically). - const txn = self.wal_log.next_lsn.load(.monotonic); - _ = try self.wal_log.write(txn, .doc_insert, 0, 0, enc); + // Use stack buffer for WAL entry encoding. + var wal_buf: [65536]u8 = undefined; + var wal_heap: ?[]u8 = null; + defer if (wal_heap) |hb| self.alloc.free(hb); + const wal_enc_buf = if (total_size <= wal_buf.len) + &wal_buf + else blk: { + wal_heap = try self.alloc.alloc(u8, total_size + 64); + break :blk wal_heap.?; + }; + const wal_enc = try d.encodeBuf(wal_enc_buf); + _ = try self.wal_log.write(0, .doc_insert, 0, 0, wal_enc); - // Find (or allocate) a leaf page with enough space. - const pno = try self.findOrAllocLeaf(enc.len); - const page_off = self.pf.leafAppend(pno, enc) orelse + // Write directly to page memory (single copy, no intermediate buffer). + const pno = try self.findOrAllocLeaf(total_size); + const reserved = self.pf.leafReserve(pno, total_size) orelse return error.PageFull; + const page_off = reserved.off; + _ = d.encodeBuf(reserved.buf) catch return error.PageFull; // Index the document. const entry = BTreeEntry{ @@ -393,14 +407,33 @@ pub const Collection = struct { .page_no = pno, .page_off = page_off, }; - try self.idx.insert(entry); + + // Lock shared state — hash_idx, key_doc_ids, key_epochs, and versions + // are collection-wide and not protected by per-key stripe locks. + // Without this, concurrent inserts to different keys race on HashMap + // internals during resize, causing GPA panics or corruption. + self.shared_mu.lock(); + self.idx.insert(entry) catch |e| { + self.shared_mu.unlock(); + return e; + }; self.hash_idx.put(hdr.key_hash, entry) catch {}; // MVCC: register version in the version chain. const epoch = self.epochs.advance(); self.versions.appendVersion(self.alloc, doc_id, pno, page_off, epoch) catch {}; + self.versions.current_epoch.store(epoch, .release); self.key_doc_ids.put(hdr.key_hash, doc_id) catch {}; self.key_epochs.put(hdr.key_hash, doc_id) catch {}; // epoch = doc_id (monotonic) + const should_gc = self.gc_counter.fetchAdd(1, .monotonic) % GC_INTERVAL == GC_INTERVAL - 1; + self.shared_mu.unlock(); + + // Run GC outside shared_mu to avoid blocking all readers during collection. + if (should_gc) { + self.shared_mu.lock(); + _ = self.versions.gc(self.alloc); + self.shared_mu.unlock(); + } // Async trigram + word indexing — push to background queue, sync fallback if full or no worker. if (value.len >= 3) { @@ -412,11 +445,11 @@ pub const Collection = struct { const owned_key = self.alloc.dupe(u8, key) catch null; const owned_val = self.alloc.dupe(u8, value) catch null; if (owned_key != null and owned_val != null) { - const q = &self.index_queue; + const q = if (self.queue_toggle.fetchAdd(1, .monotonic) % 2 == 0) &self.index_queue else &self.index_queue2; var retries: u32 = 0; while (!q.push(owned_key.?, owned_val.?)) { retries += 1; - if (retries > 1000) { + if (retries > 64) { // Queue full — fall back to synchronous indexing. self.tri.indexFile(owned_key.?, owned_val.?) catch {}; self.words.indexFile(owned_key.?, owned_val.?) catch {}; @@ -424,7 +457,7 @@ pub const Collection = struct { self.alloc.free(owned_val.?); break; } - std.Thread.yield() catch {}; + std.atomic.spinLoopHint(); } } } @@ -452,14 +485,168 @@ pub const Collection = struct { emitChange(self, .insert, key, value, doc_id); - // Periodic MVCC version chain GC to prevent unbounded memory growth. - if (self.gc_counter.fetchAdd(1, .monotonic) % GC_INTERVAL == GC_INTERVAL - 1) { - _ = self.gcVersions(); - } - return doc_id; } + /// Fast-path batch insert for bulk operations. Takes shared_mu ONCE for all + /// HashMap operations, skips per-doc WAL txn overhead, and defers trigram + /// indexing entirely to background workers. + pub fn insertBatch(self: *Collection, keys: []const []const u8, values: []const []const u8) InsertBatchResult { + var result = InsertBatchResult{ .inserted = 0, .errors = 0 }; + if (keys.len == 0) return result; + const count = @min(keys.len, values.len); + + // Phase 1: encode + write pages (no shared lock needed — page-level locking). + // Increased from 1024 to match server-side MaxLines (4096) — one flush per HTTP batch. + const MaxBatch = 4096; + var entries: [MaxBatch]BTreeEntry = undefined; + var doc_ids: [MaxBatch]u64 = undefined; + var hashes: [MaxBatch]u64 = undefined; + var wal_payloads: [MaxBatch][]const u8 = undefined; + var batch_n: usize = 0; + + // Cache current leaf page — avoid per-doc findOrAllocLeaf scan. + var current_leaf: u32 = 0; + var leaf_valid = false; + + var i: usize = 0; + while (i < count) : (i += 1) { + if (batch_n >= MaxBatch) { + // Flush WAL batch (single lock), then index batch. + self.wal_log.writeBatch(0, .doc_insert, 0, 0, wal_payloads[0..batch_n]) catch {}; + self.flushBatch(entries[0..batch_n], doc_ids[0..batch_n], hashes[0..batch_n], keys, values, i - batch_n); + result.inserted += @intCast(batch_n); + batch_n = 0; + } + const key = keys[i]; + const value = values[i]; + if (key.len == 0 or key.len > std.math.maxInt(u16)) { result.errors += 1; continue; } + if (value.len > std.math.maxInt(u32)) { result.errors += 1; continue; } + + const doc_id = self.next_doc_id.fetchAdd(1, .monotonic); + const hdr = doc_mod.newHeader(doc_id, key, value); + const d = Doc{ .header = hdr, .key = key, .value = value }; + const total_size = DocHeader.size + key.len + value.len; + + // Page write — try cached leaf first, allocate fresh only when full. + if (!leaf_valid) { + current_leaf = self.findOrAllocLeaf(total_size) catch { result.errors += 1; continue; }; + leaf_valid = true; + } + var reserved = self.pf.leafReserve(current_leaf, total_size); + if (reserved == null) { + // Current leaf full — allocate a fresh page directly (skip backward scan + // since we know recent pages are full from this batch). + current_leaf = self.pf.allocPage(.leaf) catch { result.errors += 1; leaf_valid = false; continue; }; + reserved = self.pf.leafReserve(current_leaf, total_size); + if (reserved == null) { result.errors += 1; leaf_valid = false; continue; } + } + const res = reserved.?; + + // Encode ONCE — directly to the page. WAL reads from page memory. + _ = d.encodeBuf(res.buf) catch { result.errors += 1; continue; }; + + entries[batch_n] = BTreeEntry{ + .key_hash = hdr.key_hash, + .doc_id = doc_id, + .page_no = current_leaf, + .page_off = res.off, + }; + doc_ids[batch_n] = doc_id; + hashes[batch_n] = hdr.key_hash; + wal_payloads[batch_n] = res.buf[0..total_size]; + batch_n += 1; + } + // Flush remaining. + if (batch_n > 0) { + self.wal_log.writeBatch(0, .doc_insert, 0, 0, wal_payloads[0..batch_n]) catch {}; + self.flushBatch(entries[0..batch_n], doc_ids[0..batch_n], hashes[0..batch_n], keys, values, count - batch_n); + result.inserted += @intCast(batch_n); + } + return result; + } + + const InsertBatchResult = struct { inserted: u32, errors: u32 }; + + /// Phase 2 of batch insert: acquire shared_mu ONCE, insert all entries into + /// indexes, then push all docs to async indexer. + fn flushBatch( + self: *Collection, + entries: []const BTreeEntry, + doc_ids: []const u64, + hashes: []const u64, + keys: []const []const u8, + values: []const []const u8, + key_offset: usize, + ) void { + // Pre-grow HashMaps under lock to avoid data race on internal arrays. + // (Previously this was outside the lock — concurrent resizes corrupted the map.) + { + self.shared_mu.lock(); + const needed = entries.len + self.hash_idx.count(); + self.hash_idx.ensureTotalCapacity(@intCast(needed)) catch {}; + self.key_doc_ids.ensureTotalCapacity(@intCast(needed)) catch {}; + self.key_epochs.ensureTotalCapacity(@intCast(needed)) catch {}; + self.shared_mu.unlock(); + } + + // PostgreSQL pattern: small critical sections with frequent yields. + // Instead of holding shared_mu for the entire batch (4096 docs × 7 ops), + // chunk into sub-batches of 64 so parallel writers can interleave. + const ChunkSize = 64; + var last_epoch: u64 = 0; + var chunk_start: usize = 0; + while (chunk_start < entries.len) { + const chunk_end = @min(chunk_start + ChunkSize, entries.len); + + self.shared_mu.lock(); + for (chunk_start..chunk_end) |j| { + const entry = entries[j]; + self.idx.insert(entry) catch continue; + self.hash_idx.put(entry.key_hash, entry) catch {}; + const epoch = self.epochs.advance(); + self.versions.appendVersion(self.alloc, doc_ids[j], entry.page_no, entry.page_off, epoch) catch {}; + last_epoch = epoch; + self.key_doc_ids.put(hashes[j], doc_ids[j]) catch {}; + self.key_epochs.put(hashes[j], doc_ids[j]) catch {}; + } + if (last_epoch > 0) self.versions.current_epoch.store(last_epoch, .release); + self.shared_mu.unlock(); + + chunk_start = chunk_end; + } + + // GC if needed (outside lock). + const gc_val = self.gc_counter.fetchAdd(@intCast(entries.len), .monotonic); + if (gc_val >= GC_INTERVAL and gc_val % GC_INTERVAL < entries.len) { + self.shared_mu.lock(); + _ = self.versions.gc(self.alloc); + self.shared_mu.unlock(); + } + + // Push to async trigram/word index queues. + for (entries, 0..) |_, j| { + const ki = key_offset + j; + if (ki >= keys.len) break; + const value = values[ki]; + const key = keys[ki]; + if (value.len < 3) continue; + const owned_key = self.alloc.dupe(u8, key) catch continue; + const owned_val = self.alloc.dupe(u8, value) catch { self.alloc.free(owned_key); continue; }; + const q = if (self.queue_toggle.fetchAdd(1, .monotonic) % 2 == 0) &self.index_queue else &self.index_queue2; + if (!q.push(owned_key, owned_val)) { + self.alloc.free(owned_key); + self.alloc.free(owned_val); + } + } + + // Emit CDC for each doc. + for (entries, 0..) |_, j| { + const ki = key_offset + j; + if (ki >= keys.len) break; + emitChange(self, .insert, keys[ki], values[ki], doc_ids[j]); + } + } /// Insert a document with a pre-computed embedding (no JSON parsing needed). /// This is the fast path — embedding is passed directly as f32 slice. pub fn insertWithEmbedding(self: *Collection, key: []const u8, value: []const u8, embedding: []const f32) !u64 { @@ -475,7 +662,10 @@ pub const Collection = struct { vc.append(self.alloc, embedding) catch {}; // Look up the BTreeEntry for this doc const key_hash = doc_mod.fnv1a(key); - if (self.hash_idx.get(key_hash)) |entry| { + self.shared_mu.lock(); + const hash_result = self.hash_idx.get(key_hash); + self.shared_mu.unlock(); + if (hash_result) |entry| { self.vec_entries.append(self.alloc, entry) catch {}; } } @@ -497,14 +687,21 @@ pub const Collection = struct { } // L2: Hash index (O(1), but hash table lookup) - if (self.hash_idx.get(key_hash)) |entry| { + // shared_mu protects against concurrent HashMap resize during put(). + self.shared_mu.lock(); + const hash_entry = self.hash_idx.get(key_hash); + self.shared_mu.unlock(); + if (hash_entry) |entry| { if (self.readEntry(entry)) |d| { self.cache.insert(key_hash, entry.page_no, entry.page_off); return d; } } // L3: B-tree (O(log n)) - const entry = self.idx.search(key_hash) orelse return null; + self.shared_mu.lock(); + const btree_entry = self.idx.search(key_hash); + self.shared_mu.unlock(); + const entry = btree_entry orelse return null; const d = self.readEntry(entry) orelse return null; self.cache.insert(key_hash, entry.page_no, entry.page_off); return d; @@ -512,9 +709,19 @@ pub const Collection = struct { pub fn getAsOfEpoch(self: *Collection, key: []const u8, epoch: u64) ?Doc { const key_hash = doc_mod.fnv1a(key); - const doc_id = self.key_doc_ids.get(key_hash) orelse return null; - const ver = self.versions.getAtEpoch(doc_id, epoch) orelse return null; - return self.readLoc(ver.page_no, ver.page_off); + self.shared_mu.lock(); + const doc_id = self.key_doc_ids.get(key_hash) orelse { + self.shared_mu.unlock(); + return null; + }; + const ver = self.versions.getAtEpoch(doc_id, epoch) orelse { + self.shared_mu.unlock(); + return null; + }; + const pno = ver.page_no; + const poff = ver.page_off; + self.shared_mu.unlock(); + return self.readLoc(pno, poff); } pub fn getAsOfTimestamp(self: *Collection, key: []const u8, ts_ms: i64) ?Doc { @@ -548,12 +755,17 @@ pub const Collection = struct { /// Update an existing document (append new version, update index). /// Uses per-key stripe lock — concurrent updates to different keys proceed in parallel. pub fn update(self: *Collection, key: []const u8, new_value: []const u8) !bool { + if (key.len == 0 or key.len > std.math.maxInt(u16)) return error.InvalidKeyLength; + if (new_value.len > std.math.maxInt(u32)) return error.InvalidValueLength; const key_hash = doc_mod.fnv1a(key); const stripe = stripeIndex(key_hash); self.stripe_locks[stripe].lock(); defer self.stripe_locks[stripe].unlock(); - const old_entry = self.idx.search(key_hash) orelse return false; + self.shared_mu.lock(); + const old_entry_result = self.idx.search(key_hash); + self.shared_mu.unlock(); + const old_entry = old_entry_result orelse return false; const old_doc = self.readEntry(old_entry) orelse return false; const doc_id = old_doc.header.doc_id; @@ -566,8 +778,7 @@ pub const Collection = struct { var enc_buf: [65536]u8 = undefined; const enc = try d.encodeBuf(&enc_buf); - const txn = self.wal_log.next_lsn.load(.monotonic); - _ = try self.wal_log.write(txn, .doc_update, 0, 0, enc); + _ = try self.wal_log.write(0, .doc_update, 0, 0, enc); const pno = try self.findOrAllocLeaf(enc.len); const page_off = self.pf.leafAppend(pno, enc) orelse return error.PageFull; @@ -578,12 +789,23 @@ pub const Collection = struct { .page_no = pno, .page_off = page_off, }; + self.shared_mu.lock(); self.hash_idx.put(key_hash, new_entry) catch {}; + self.idx.insert(new_entry) catch {}; self.cache.invalidate(key_hash); // MVCC: register new version in the version chain (links to old automatically). const epoch = self.epochs.advance(); self.versions.appendVersion(self.alloc, doc_id, pno, page_off, epoch) catch {}; + self.versions.current_epoch.store(epoch, .release); + self.key_epochs.put(key_hash, doc_id) catch {}; + const should_gc = self.gc_counter.fetchAdd(1, .monotonic) % GC_INTERVAL == GC_INTERVAL - 1; + self.shared_mu.unlock(); + if (should_gc) { + self.shared_mu.lock(); + _ = self.versions.gc(self.alloc); + self.shared_mu.unlock(); + } emitChange(self, .update, key, new_value, doc_id); return true; @@ -597,25 +819,40 @@ pub const Collection = struct { self.stripe_locks[stripe].lock(); defer self.stripe_locks[stripe].unlock(); - const entry = self.idx.search(key_hash) orelse return false; + self.shared_mu.lock(); + const del_entry_result = self.idx.search(key_hash); + self.shared_mu.unlock(); + const entry = del_entry_result orelse return false; const old_doc = self.readEntry(entry) orelse return false; var tomb_hdr = doc_mod.newHeader(old_doc.header.doc_id, key, ""); tomb_hdr.flags |= DocHeader.DELETED; tomb_hdr.version = old_doc.header.version +% 1; tomb_hdr.next_ver = (@as(u64, entry.page_no) << 16) | entry.page_off; - const txn = self.wal_log.next_lsn.load(.monotonic); - _ = try self.wal_log.write(txn, .doc_delete, 0, 0, std.mem.asBytes(&tomb_hdr)); + _ = try self.wal_log.write(0, .doc_delete, 0, 0, std.mem.asBytes(&tomb_hdr)); const tomb_doc = Doc{ .header = tomb_hdr, .key = key, .value = "" }; var enc_buf: [65536]u8 = undefined; const enc = try tomb_doc.encodeBuf(&enc_buf); const pno = try self.findOrAllocLeaf(enc.len); const page_off = self.pf.leafAppend(pno, enc) orelse return error.PageFull; + self.shared_mu.lock(); const epoch = self.epochs.advance(); self.versions.appendVersion(self.alloc, old_doc.header.doc_id, pno, page_off, epoch) catch {}; + self.versions.current_epoch.store(epoch, .release); self.idx.delete(key_hash); _ = self.hash_idx.remove(key_hash); + _ = self.key_doc_ids.remove(key_hash); + _ = self.key_epochs.remove(key_hash); + const should_gc = self.gc_counter.fetchAdd(1, .monotonic) % GC_INTERVAL == GC_INTERVAL - 1; + self.shared_mu.unlock(); + if (should_gc) { + self.shared_mu.lock(); + _ = self.versions.gc(self.alloc); + self.shared_mu.unlock(); + } self.cache.invalidate(key_hash); + self.tri.removeFile(key); + self.words.removeFile(key); emitChange(self, .delete, key, "", old_doc.header.doc_id); return true; } @@ -712,9 +949,10 @@ pub const Collection = struct { }; } - /// O(1) word lookup using the inverted word index. - pub fn searchWord(self: *Collection, word: []const u8) []const codeindex.WordHit { - return self.words.search(word); + /// O(1) word lookup using the inverted word index. Returns an owned copy + /// that the caller must free with `allocator.free(result)`. + pub fn searchWord(self: *Collection, word: []const u8, allocator: std.mem.Allocator) ![]const codeindex.WordHit { + return self.words.search(word, allocator); } fn bruteForceSearch(self: *Collection, query: []const u8, limit: u32, alloc: std.mem.Allocator) !TextSearchResult { @@ -739,19 +977,30 @@ pub const Collection = struct { }; } + /// Lookup table for fast ASCII lowercasing — avoids branch per byte. + const lower_lut: [256]u8 = blk: { + var t: [256]u8 = undefined; + for (0..256) |i| t[i] = if (i >= 'A' and i <= 'Z') @intCast(i + 32) else @intCast(i); + break :blk t; + }; + fn containsInsensitive(haystack: []const u8, needle: []const u8) bool { if (needle.len == 0) return true; if (haystack.len < needle.len) return false; + // Pre-lowercase the needle once (avoids repeated table lookups per position). + var lower_needle: [256]u8 = undefined; + const nlen = @min(needle.len, 256); + for (needle[0..nlen], 0..) |c, k| lower_needle[k] = lower_lut[c]; + if (needle.len > 256) return false; // needle > 256 chars unsupported + const first = lower_needle[0]; var i: usize = 0; - while (i + needle.len <= haystack.len) : (i += 1) { + const end = haystack.len - needle.len; + while (i <= end) : (i += 1) { + // Fast skip: check first byte before full comparison. + if (lower_lut[haystack[i]] != first) continue; var match = true; - var j: usize = 0; - while (j < needle.len) : (j += 1) { - const hc = haystack[i + j]; - const nc = needle[j]; - const hl = if (hc >= 'A' and hc <= 'Z') hc + 32 else hc; - const nl = if (nc >= 'A' and nc <= 'Z') nc + 32 else nc; - if (hl != nl) { match = false; break; } + for (1..nlen) |j| { + if (lower_lut[haystack[i + j]] != lower_needle[j]) { match = false; break; } } if (match) return true; } @@ -810,13 +1059,17 @@ pub const Collection = struct { const total_pages = self.pf.next_alloc.load(.acquire); var skipped: u32 = 0; var pno: u32 = 0; + // Snapshot used_bytes per page under shared_mu to avoid reading + // partially-written documents from concurrent inserts. + self.shared_mu.lock(); outer: while (pno < total_pages) : (pno += 1) { const ph = self.pf.pageHeader(pno); if (@as(page_mod.PageType, @enumFromInt(ph.page_type)) != .leaf) continue; + const used = ph.used_bytes; const data = self.pf.pageData(pno); var pos: usize = 0; - while (pos + DocHeader.size <= ph.used_bytes) { - const rem = data[pos..ph.used_bytes]; + while (pos + DocHeader.size <= used) { + const rem = data[pos..used]; const decoded = doc_mod.decode(rem) catch break; const d = decoded.doc; pos += decoded.consumed; @@ -826,6 +1079,7 @@ pub const Collection = struct { if (results.items.len >= limit) break :outer; } } + self.shared_mu.unlock(); return ScanResult{ .docs = try results.toOwnedSlice(alloc), .alloc = alloc }; } @@ -834,9 +1088,23 @@ pub const Collection = struct { errdefer results.deinit(alloc); var skipped: u32 = 0; + // Snapshot key_doc_ids under shared_mu to avoid racing with insert(). + self.shared_mu.lock(); var it = self.key_doc_ids.iterator(); + // Collect doc_id/epoch pairs while holding the lock (fast — just integers). + const Pair = struct { doc_id: u64 }; + var pairs: std.ArrayList(Pair) = .empty; + defer pairs.deinit(alloc); while (it.next()) |entry| { const ver = self.versions.getAtEpoch(entry.value_ptr.*, epoch) orelse continue; + _ = ver; + pairs.append(alloc, .{ .doc_id = entry.value_ptr.* }) catch continue; + } + self.shared_mu.unlock(); + + // Now read pages outside the lock (safe — pages are stable after mmap rewrite). + for (pairs.items) |pair| { + const ver = self.versions.getAtEpoch(pair.doc_id, epoch) orelse continue; const doc = self.readLoc(ver.page_no, ver.page_off) orelse continue; if (skipped < offset) { skipped += 1; @@ -872,8 +1140,10 @@ pub const Collection = struct { } fn readEntry(self: *Collection, entry: BTreeEntry) ?Doc { - const raw = self.pf.leafRead(entry.page_no, entry.page_off, - DocHeader.size + 1024 + 65536); + const ph = self.pf.pageHeader(entry.page_no); + if (entry.page_off >= ph.used_bytes) return null; + const data = self.pf.pageData(entry.page_no); + const raw = data[entry.page_off..ph.used_bytes]; const decoded = doc_mod.decode(raw) catch return null; if (decoded.doc.header.flags & DocHeader.DELETED != 0) return null; return decoded.doc; @@ -918,7 +1188,9 @@ pub const Collection = struct { }; if (d.header.flags & DocHeader.DELETED == 0) { self.hash_idx.put(d.header.key_hash, entry) catch {}; + self.idx.insert(entry) catch {}; self.key_doc_ids.put(d.header.key_hash, d.header.doc_id) catch {}; + self.key_epochs.put(d.header.key_hash, d.header.doc_id) catch {}; if (d.value.len >= 3) { self.tri.indexFile(d.key, d.value) catch {}; self.words.indexFile(d.key, d.value) catch {}; @@ -1030,7 +1302,7 @@ pub const Collection = struct { const vc = self.vectors orelse return self.searchText(text_query, limit, result_alloc); // Phase 1: Text pre-filter — get candidate doc keys - const text_results = try self.searchText(text_query, limit * 3, result_alloc); + var text_results = try self.searchText(text_query, limit * 3, result_alloc); if (text_results.docs.len == 0 or vector_query.len != vc.dims) { return text_results; @@ -1077,10 +1349,8 @@ pub const Collection = struct { text_results.docs[i] = sd.doc; } - // If we have more docs than limit, shrink the slice if (text_results.docs.len > out_len) { - // We can't easily shrink, but the caller will respect .docs.len - // Just return the full text_results (already re-ranked in-place) + text_results.docs = text_results.docs[0..out_len]; } return text_results; @@ -1144,21 +1414,26 @@ pub const Collection = struct { var it = br.writes.iterator(); while (it.next()) |entry| { const w = entry.value_ptr.*; + const key = entry.key_ptr.*; if (w.deleted) { applied += 1; continue; } - const key_hash = doc_mod.fnv1a(w.key); + const key_hash = doc_mod.fnv1a(key); - // CONFLICT DETECTION: check if main modified this key after branch was created - if (self.key_epochs.get(key_hash)) |main_epoch| { + // CONFLICT DETECTION: check if main modified this key after branch was created. + // Acquire shared_mu to safely read key_epochs (concurrent inserts mutate it). + self.shared_mu.lock(); + const main_epoch_opt = self.key_epochs.get(key_hash); + self.shared_mu.unlock(); + if (main_epoch_opt) |main_epoch| { if (main_epoch > br.base_epoch) { // Main was modified after branch fork — this is a real conflict - const main_doc = self.get(w.key); + const main_doc = self.get(key); try conflicts_list.append(merge_alloc, .{ - .key = w.key, + .key = key, .branch_value = w.value, .main_value = if (main_doc) |d| d.value else "", }); @@ -1167,7 +1442,7 @@ pub const Collection = struct { } // No conflict — safe to apply - _ = self.insert(w.key, w.value) catch continue; + _ = self.insert(key, w.value) catch continue; applied += 1; } @@ -1228,20 +1503,21 @@ pub const Collection = struct { var it = br.writes.iterator(); while (it.next()) |entry| { const w = entry.value_ptr.*; + const key = entry.key_ptr.*; if (w.deleted) continue; // Check if query appears in the branch-local value or key - if (containsInsensitive(w.value, query) or containsInsensitive(w.key, query)) { + if (containsInsensitive(w.value, query) or containsInsensitive(key, query)) { try branch_matches.append(search_alloc, Doc{ .header = doc_mod.DocHeader{ .doc_id = w.epoch, - .key_hash = doc_mod.fnv1a(w.key), + .key_hash = doc_mod.fnv1a(key), .val_len = @intCast(w.value.len), - .key_len = @intCast(w.key.len), + .key_len = @intCast(key.len), .flags = 0, .version = 0, .next_ver = 0, }, - .key = w.key, + .key = key, .value = w.value, }); } @@ -1440,7 +1716,7 @@ fn indexWorkerQ(col: *Collection, queue: *IndexQueue) void { } if (n == 0) { _ = col.indexing_count.fetchSub(1, .release); - std.Thread.yield() catch {}; + std.Thread.sleep(1_000_000); // 1ms — avoid busy-spin when queue is empty continue; } // Batch trigram indexing (single lock acquisition for all docs). @@ -1603,6 +1879,13 @@ pub const Database = struct { kv.value.close(); self.alloc.free(kv.key); } + + // Delete the backing page file so data doesn't reappear on re-access. + var storage_name_buf: [MAX_TENANT_ID_LEN + MAX_COLLECTION_NAME_LEN + 4]u8 = undefined; + const storage_name = makeStorageName(&storage_name_buf, tenant_id, name) catch return; + var path_buf: [512]u8 = undefined; + const page_path = std.fmt.bufPrint(&path_buf, "{s}/{s}.pages", .{ self.dataDir(), storage_name }) catch return; + std.fs.cwd().deleteFile(page_path) catch {}; } pub fn listCollectionsForTenant(self: *Database, tenant_id: []const u8, alloc: std.mem.Allocator) !std.ArrayList([]const u8) { @@ -2003,3 +2286,122 @@ test "extractJsonFloatArray handles whitespace in array" { try std.testing.expectApproxEqAbs(@as(f32, 2.0), out[1], 0.001); try std.testing.expectApproxEqAbs(@as(f32, 3.0), out[2], 0.001); } + +test "concurrent inserts to same collection do not crash" { + const alloc = std.testing.allocator; + const db = try Database.open(alloc, "./test_concurrent_data"); + defer db.close(); + defer std.fs.cwd().deleteTree("./test_concurrent_data") catch {}; + + const col = try db.collection("stress"); + + // Spawn multiple threads doing concurrent inserts to different keys + const THREADS = 4; + const INSERTS_PER_THREAD = 500; + var threads: [THREADS]std.Thread = undefined; + + const Worker = struct { + fn run(c: *Collection, thread_id: usize) void { + var key_buf: [64]u8 = undefined; + var val_buf: [128]u8 = undefined; + for (0..INSERTS_PER_THREAD) |i| { + const klen = std.fmt.bufPrint(&key_buf, "t{d}_k{d}", .{ thread_id, i }) catch continue; + const vlen = std.fmt.bufPrint(&val_buf, "{{\"thread\":{d},\"seq\":{d},\"data\":\"payload\"}}", .{ thread_id, i }) catch continue; + _ = c.insert(klen, vlen) catch continue; + } + } + }; + + for (0..THREADS) |t| { + threads[t] = try std.Thread.spawn(.{}, Worker.run, .{ col, t }); + } + for (0..THREADS) |t| { + threads[t].join(); + } + + // Verify data integrity — spot-check a few keys + for (0..THREADS) |t| { + var key_buf: [64]u8 = undefined; + const klen = std.fmt.bufPrint(&key_buf, "t{d}_k0", .{t}) catch continue; + const doc = col.get(klen); + try std.testing.expect(doc != null); + } +} + +test "bulk insert triggers MVCC GC and versions stay bounded" { + const alloc = std.testing.allocator; + const db = try Database.open(alloc, "./test_bulk_gc_data"); + defer db.close(); + defer std.fs.cwd().deleteTree("./test_bulk_gc_data") catch {}; + + const col = try db.collection("bulkgc"); + + // Insert more than GC_INTERVAL documents + const N = Collection.GC_INTERVAL * 3; + var key_buf: [64]u8 = undefined; + for (0..N) |i| { + const klen = std.fmt.bufPrint(&key_buf, "doc_{d}", .{i}) catch continue; + _ = col.insert(klen, "{\"val\":1}") catch continue; + } + + // GC should have run at least twice (every 500 inserts) + // all_versions should be well under N because GC trims old entries + col.shared_mu.lock(); + const versions_len = col.versions.all_versions.items.len; + col.shared_mu.unlock(); + // After GC, the remaining versions should be <= N (they could be N if + // all are latest, but at minimum GC should not have crashed) + try std.testing.expect(versions_len <= N); + // Verify reads still work + const doc = col.get("doc_0"); + try std.testing.expect(doc != null); +} + +test "concurrent reads and writes do not crash" { + const alloc = std.testing.allocator; + const db = try Database.open(alloc, "./test_rw_data"); + defer db.close(); + defer std.fs.cwd().deleteTree("./test_rw_data") catch {}; + + const col = try db.collection("readwrite"); + + // Pre-populate some data + for (0..100) |i| { + var key_buf: [64]u8 = undefined; + const klen = std.fmt.bufPrint(&key_buf, "pre_{d}", .{i}) catch continue; + _ = col.insert(klen, "{\"pre\":true}") catch continue; + } + + // Spawn writers and readers concurrently + const WriterCtx = struct { + fn run(c: *Collection) void { + var key_buf: [64]u8 = undefined; + var val_buf: [128]u8 = undefined; + for (0..200) |i| { + const klen = std.fmt.bufPrint(&key_buf, "w_{d}", .{i}) catch continue; + const vlen = std.fmt.bufPrint(&val_buf, "{{\"w\":{d}}}", .{i}) catch continue; + _ = c.insert(klen, vlen) catch continue; + } + } + }; + const ReaderCtx = struct { + fn run(c: *Collection) void { + var key_buf: [64]u8 = undefined; + for (0..200) |i| { + const klen = std.fmt.bufPrint(&key_buf, "pre_{d}", .{i % 100}) catch continue; + _ = c.get(klen); + } + } + }; + + var writer = try std.Thread.spawn(.{}, WriterCtx.run, .{col}); + var reader1 = try std.Thread.spawn(.{}, ReaderCtx.run, .{col}); + var reader2 = try std.Thread.spawn(.{}, ReaderCtx.run, .{col}); + + writer.join(); + reader1.join(); + reader2.join(); + + // If we got here without crashing, the test passes + try std.testing.expect(true); +} diff --git a/src/compression.zig b/src/compression.zig index 4b7b981..da8d42e 100644 --- a/src/compression.zig +++ b/src/compression.zig @@ -35,7 +35,7 @@ pub fn compress(src: []const u8, dst: []u8) CompressionError!usize { return emitLiterals(src, 0, src.len, dst, 0); } - var hash_table: [HASH_SIZE]u16 = [_]u16{0} ** HASH_SIZE; + var hash_table: [HASH_SIZE]u32 = [_]u32{0} ** HASH_SIZE; var src_pos: usize = 0; var dst_pos: usize = 0; @@ -47,7 +47,7 @@ pub fn compress(src: []const u8, dst: []u8) CompressionError!usize { // Hash current 4-byte sequence const h = hash4(src, src_pos); const match_pos: usize = hash_table[h]; - hash_table[h] = @intCast(if (src_pos > std.math.maxInt(u16)) std.math.maxInt(u16) else src_pos); + hash_table[h] = @intCast(src_pos); // Check for match: must be within u16 offset range and actually match 4 bytes const offset = src_pos -% match_pos; diff --git a/src/disk_index.zig b/src/disk_index.zig index 58fcbc4..37370a7 100644 --- a/src/disk_index.zig +++ b/src/disk_index.zig @@ -261,15 +261,18 @@ pub const DiskIndex = struct { defer idx_file.close(); const idx_stat = try idx_file.stat(); const idx_mmap = try std.posix.mmap(null, idx_stat.size, std.posix.PROT.READ, .{ .TYPE = .PRIVATE }, idx_file.handle, 0); + errdefer std.posix.munmap(idx_mmap); const n_entries = std.mem.bytesToValue(u32, idx_mmap[0..4]); const entries_start = 4; const entries_end = entries_start + n_entries * @sizeOf(LookupEntry); + if (entries_end > idx_stat.size) return error.CorruptIndex; const lookup: []const LookupEntry = @alignCast(std.mem.bytesAsSlice(LookupEntry, idx_mmap[entries_start..entries_end])); // Open postings file for pread const posts_path = try std.fmt.bufPrint(&buf, "{s}/posts.tdb", .{dir_path}); const posts_fd = try std.fs.cwd().openFile(posts_path, .{}); + errdefer posts_fd.close(); // mmap files table const files_path = try std.fmt.bufPrint(&buf, "{s}/files.tdb", .{dir_path}); @@ -277,6 +280,7 @@ pub const DiskIndex = struct { defer files_file.close(); const files_stat = try files_file.stat(); const files_mmap = try std.posix.mmap(null, files_stat.size, std.posix.PROT.READ, .{ .TYPE = .PRIVATE }, files_file.handle, 0); + errdefer std.posix.munmap(files_mmap); const n_files = std.mem.bytesToValue(u32, files_mmap[0..4]); // mmap frequency table @@ -285,6 +289,7 @@ pub const DiskIndex = struct { defer freq_file.close(); const freq_stat = try freq_file.stat(); const freq_mmap = try std.posix.mmap(null, freq_stat.size, std.posix.PROT.READ, .{ .TYPE = .PRIVATE }, freq_file.handle, 0); + errdefer std.posix.munmap(freq_mmap); return .{ .lookup = lookup, @@ -465,9 +470,10 @@ pub fn extractSparseNgrams( var start: usize = 0; pi = 1; - while (pi < n_pairs) : (pi += 1) { + const max_pairs = @min(n_pairs, 1024); + while (pi < max_pairs) : (pi += 1) { // Is this a local maximum? Both neighbors have lower weight - const is_boundary = (pi == n_pairs - 1) or + const is_boundary = (pi == max_pairs - 1) or (weights[pi] >= weights[pi - 1] and weights[pi] >= weights[pi + 1]); if (is_boundary or pi - start >= 8) { diff --git a/src/doc.zig b/src/doc.zig index 55a9947..a3fec37 100644 --- a/src/doc.zig +++ b/src/doc.zig @@ -76,9 +76,11 @@ pub const Doc = struct { pub fn decode(buf: []const u8) error{TooShort,Corrupt}!struct { doc: Doc, consumed: usize } { if (buf.len < DocHeader.size) return error.TooShort; const hdr: DocHeader = std.mem.bytesToValue(DocHeader, buf[0..DocHeader.size]); - const total = DocHeader.size + hdr.key_len + hdr.val_len; - if (buf.len < total) return error.TooShort; if (hdr.key_len > 1024) return error.Corrupt; + // Sanity check: reject impossibly large values (>64MB) as corruption. + if (hdr.val_len > 64 * 1024 * 1024) return error.Corrupt; + const total = DocHeader.size + @as(usize, hdr.key_len) + @as(usize, hdr.val_len); + if (buf.len < total) return error.TooShort; return .{ .doc = .{ .header = hdr, diff --git a/src/errors.zig b/src/errors.zig index 5689575..1b58ac9 100644 --- a/src/errors.zig +++ b/src/errors.zig @@ -106,9 +106,16 @@ pub fn jsonError(buf: []u8, code: ErrorCode) []const u8 { pub fn jsonErrorDetail(buf: []u8, code: ErrorCode, detail: []const u8) []const u8 { var fbs = std.io.fixedBufferStream(buf); const w = fbs.writer(); - w.print("{{\"error\":{{\"code\":{d},\"message\":\"{s}\",\"detail\":\"{s}\"}}}}", .{ - @intFromEnum(code), code.message(), detail, + w.print("{{\"error\":{{\"code\":{d},\"message\":\"{s}\",\"detail\":\"", .{ + @intFromEnum(code), code.message(), }) catch return "{}"; + // Escape detail to prevent JSON injection + for (detail) |ch| { + if (ch == '"' or ch == '\\') w.writeByte('\\') catch {}; + if (ch == '\n') { w.writeAll("\\n") catch {}; continue; } + w.writeByte(ch) catch {}; + } + w.writeAll("\"}}}}") catch return "{}"; return buf[0..fbs.pos]; } diff --git a/src/fast_index.zig b/src/fast_index.zig index 8c0ab30..aaf3c96 100644 --- a/src/fast_index.zig +++ b/src/fast_index.zig @@ -35,6 +35,9 @@ pub const FastTrigramIndex = struct { allocator: std.mem.Allocator, + /// Mutex for concurrent access. + mu: std.Thread.Mutex = .{}, + /// Map a byte to the reduced 0-63 alphabet. /// a-z -> 0-25, 0-9 -> 26-35, _ -> 36, common punctuation -> 37-63, rest -> 63. fn charMap(c: u8) u6 { @@ -120,9 +123,12 @@ pub const FastTrigramIndex = struct { } pub fn indexFile(self: *FastTrigramIndex, path: []const u8, content: []const u8) !void { + self.mu.lock(); + defer self.mu.unlock(); + // If file was previously indexed, remove it first. if (self.path_to_id.contains(path)) { - self.removeFile(path); + self.removeFileInner(path); } // Allocate or recycle a file_id. @@ -186,6 +192,12 @@ pub const FastTrigramIndex = struct { } pub fn removeFile(self: *FastTrigramIndex, path: []const u8) void { + self.mu.lock(); + defer self.mu.unlock(); + self.removeFileInner(path); + } + + fn removeFileInner(self: *FastTrigramIndex, path: []const u8) void { const file_id = self.path_to_id.get(path) orelse return; // Remove file_id from each trigram's posting list. @@ -226,8 +238,10 @@ pub const FastTrigramIndex = struct { } pub fn candidates(self: *FastTrigramIndex, query: []const u8, alloc: std.mem.Allocator) ?[]const []const u8 { - if (query.len < 3) return null; + self.mu.lock(); + defer self.mu.unlock(); + if (query.len < 3) return null; const tri_count = query.len - 2; // Collect unique trigram keys from query. diff --git a/src/ffi.zig b/src/ffi.zig index b3ef623..63307e7 100644 --- a/src/ffi.zig +++ b/src/ffi.zig @@ -666,13 +666,14 @@ export fn turbodb_branch_diff( var it = br.writes.iterator(); while (it.next()) |entry| { const bw = entry.value_ptr.*; + const bw_key = entry.key_ptr.*; if (bw.deleted) continue; if (!first_file) w.writeByte(',') catch return -1; first_file = false; // Get main version for comparison - const main_doc = col.get(bw.key); + const main_doc = col.get(bw_key); const old_val = if (main_doc) |d| d.value else ""; // Compute line diff @@ -680,9 +681,8 @@ export fn turbodb_branch_diff( defer alloc.free(diffs); w.writeAll("{\"key\":\"") catch return -1; - w.writeAll(bw.key) catch return -1; + w.writeAll(bw_key) catch return -1; w.writeAll("\",\"lines\":[") catch return -1; - for (diffs, 0..) |d, di| { if (di > 0) w.writeByte(',') catch return -1; const kind_str = switch (d.kind) { diff --git a/src/hot_cache.zig b/src/hot_cache.zig index ed99f77..f0e0e67 100644 --- a/src/hot_cache.zig +++ b/src/hot_cache.zig @@ -8,8 +8,9 @@ pub const HotCache = struct { hand: u32, hits: std.atomic.Value(u64), misses: std.atomic.Value(u64), + mutex: std.Thread.Mutex, - const CACHE_SIZE: u32 = 16384; // power of 2 for fast modulo + const CACHE_SIZE: u32 = 4096; // power of 2 for fast modulo const MASK: u32 = CACHE_SIZE - 1; const PROBE_LIMIT: u32 = 4; @@ -35,11 +36,14 @@ pub const HotCache = struct { self.hand = 0; self.hits = std.atomic.Value(u64).init(0); self.misses = std.atomic.Value(u64).init(0); + self.mutex = .{}; return self; } pub fn lookup(self: *HotCache, key_hash: u64) ?struct { page_no: u32, page_off: u16 } { if (key_hash == 0) return null; // 0 is our sentinel + self.mutex.lock(); + defer self.mutex.unlock(); const base = @as(u32, @truncate(key_hash)) & MASK; var i: u32 = 0; while (i < PROBE_LIMIT) : (i += 1) { @@ -58,6 +62,8 @@ pub const HotCache = struct { pub fn insert(self: *HotCache, key_hash: u64, page_no: u32, page_off: u16) void { if (key_hash == 0) return; + self.mutex.lock(); + defer self.mutex.unlock(); const base = @as(u32, @truncate(key_hash)) & MASK; // Try to find existing or empty slot within probe window var i: u32 = 0; @@ -88,6 +94,8 @@ pub const HotCache = struct { pub fn invalidate(self: *HotCache, key_hash: u64) void { if (key_hash == 0) return; + self.mutex.lock(); + defer self.mutex.unlock(); const base = @as(u32, @truncate(key_hash)) & MASK; var i: u32 = 0; while (i < PROBE_LIMIT) : (i += 1) { diff --git a/src/io_engine.zig b/src/io_engine.zig index 1024ca7..9ab6b03 100644 --- a/src/io_engine.zig +++ b/src/io_engine.zig @@ -402,12 +402,14 @@ pub fn MpscRing(comptime T: type, comptime cap: usize) type { buf: [cap]T, head: std.atomic.Value(usize), // consumer reads from here tail: std.atomic.Value(usize), // producers write here (CAS) + ready: [cap]std.atomic.Value(u8), // per-slot readiness flag pub fn init() Self { return .{ .buf = undefined, .head = std.atomic.Value(usize).init(0), .tail = std.atomic.Value(usize).init(0), + .ready = [_]std.atomic.Value(u8){std.atomic.Value(u8).init(0)} ** cap, }; } @@ -418,11 +420,15 @@ pub fn MpscRing(comptime T: type, comptime cap: usize) type { const h = self.head.load(.acquire); const next = (t + 1) % cap; if (next == h) return false; // full + // CAS to claim the slot if (self.tail.cmpxchgWeak(t, next, .release, .monotonic) == null) { + // We own slot t — write data, then publish via ready flag self.buf[t] = item; + self.ready[t].store(1, .release); return true; } // CAS failed, retry. + std.atomic.spinLoopHint(); } } @@ -431,7 +437,10 @@ pub fn MpscRing(comptime T: type, comptime cap: usize) type { const h = self.head.load(.acquire); const t = self.tail.load(.acquire); if (h == t) return null; // empty + // Wait for producer to finish writing (check readiness) + if (self.ready[h].load(.acquire) != 1) return null; const item = self.buf[h]; + self.ready[h].store(0, .release); self.head.store((h + 1) % cap, .release); return item; } diff --git a/src/lsm.zig b/src/lsm.zig index 89fbc86..d94e5cb 100644 --- a/src/lsm.zig +++ b/src/lsm.zig @@ -443,6 +443,7 @@ pub const LSMTree = struct { data_dir_len: usize, alloc: std.mem.Allocator, flush_mu: std.Thread.Mutex, + mem_rw: std.Thread.RwLock, pub const MAX_LEVELS = 7; pub const MEMTABLE_SIZE = 4 * 1024 * 1024; // 4MB @@ -456,6 +457,7 @@ pub const LSMTree = struct { lsm.alloc = alloc; lsm.next_sst_id = 0; lsm.flush_mu = .{}; + lsm.mem_rw = .{}; lsm.data_dir = std.mem.zeroes([256]u8); if (data_dir.len > 256) return error.PathTooLong; @@ -492,7 +494,11 @@ pub const LSMTree = struct { } } - pub fn get(self: *const LSMTree, key_hash: u64) ?BTreeEntry { + pub fn get(self: *LSMTree, key_hash: u64) ?BTreeEntry { + // Acquire shared lock to prevent torn reads during memtable swap. + self.mem_rw.lockShared(); + defer self.mem_rw.unlockShared(); + // 1. Check active memtable. if (memtableSearch(self.active_mem.entries.items, key_hash)) |result| return result.entry; @@ -553,14 +559,23 @@ pub const LSMTree = struct { if (self.active_mem.entries.items.len == 0) return; // Rotate: move active → immutable. - var old = self.active_mem; - self.active_mem = MemTable.init(self.alloc); + // Acquire exclusive lock for the pointer swap so get() cannot see + // a half-swapped state (torn read). + var old: MemTable = undefined; + { + self.mem_rw.lock(); + defer self.mem_rw.unlock(); + old = self.active_mem; + self.active_mem = MemTable.init(self.alloc); + self.immutable_mem = old; + } defer { old.deinit(); + self.mem_rw.lock(); + defer self.mem_rw.unlock(); self.immutable_mem = null; } - self.immutable_mem = old; // Create L0 SSTable. const id = self.next_sst_id; @@ -615,8 +630,8 @@ pub const LSMTree = struct { } // j-1 is the newest entry for this key. const entry = all.items[j - 1]; - // Drop tombstones when compacting to deeper levels (> L1). - if (level > 0 or entry.value != .tombstone) { + // Drop tombstones only at the deepest level (no lower level to shadow). + if (level + 1 < MAX_LEVELS - 1 or entry.value != .tombstone) { try deduped.append(self.alloc, entry); } i = j; diff --git a/src/main.zig b/src/main.zig index f90732c..1cf24cc 100644 --- a/src/main.zig +++ b/src/main.zig @@ -111,8 +111,11 @@ pub fn main() !void { // ── configure auth ──────────────────────────────────────────────────── if (auth_key) |key| { - _ = db.auth.addKey(key, "cli", .admin); - std.log.info("Auth enabled (--auth-key)", .{}); + if (db.auth.addKey(key, "cli", .admin)) |_| { + std.log.info("Auth enabled (--auth-key)", .{}); + } else { + std.log.err("Failed to add auth key: key store full (MAX_KEYS reached)", .{}); + } } // ── replication setup ───────────────────────────────────────────────── diff --git a/src/mvcc.zig b/src/mvcc.zig index 8f359da..ce3b588 100644 --- a/src/mvcc.zig +++ b/src/mvcc.zig @@ -131,10 +131,15 @@ pub const VersionChain = struct { /// Begin a read transaction pinned at the current epoch. pub fn beginRead(self: *VersionChain) ReadTxn { const epoch = self.current_epoch.load(.acquire); - // Pin: ensure min_active_epoch doesn't advance past us. - const min = self.min_active_epoch.load(.acquire); - if (epoch < min) { - self.min_active_epoch.store(epoch, .release); + // Pin: atomically update min_active_epoch only if our epoch is smaller. + // Uses cmpxchgWeak in a loop to avoid the load-then-store race where + // two threads could both read the old min and one's store gets lost. + var min = self.min_active_epoch.load(.acquire); + while (epoch < min) { + if (self.min_active_epoch.cmpxchgWeak(min, epoch, .release, .monotonic)) |actual| { + min = actual; + std.atomic.spinLoopHint(); + } else break; } return ReadTxn{ .epoch = epoch, @@ -152,10 +157,13 @@ pub const VersionChain = struct { /// Garbage-collect versions older than min_active_epoch. /// Returns the number of versions freed. + /// If any keep.append fails (OOM), the original all_versions is restored + /// and 0 is returned — no entries are lost. pub fn gc(self: *VersionChain, alloc: Allocator) u64 { const min_epoch = self.min_active_epoch.load(.acquire); var freed: u64 = 0; var keep: std.ArrayList(VersionEntry) = .empty; + var oom = false; var i: usize = 0; while (i < self.all_versions.items.len) : (i += 1) { @@ -166,10 +174,33 @@ pub const VersionChain = struct { false; if (!is_latest and entry.ver.epoch < min_epoch) { - _ = self.history.remove(chainKey(entry.ver.page_no, entry.ver.page_off)); freed += 1; } else { - keep.append(alloc, .{ .doc_id = entry.doc_id, .ver = entry.ver }) catch {}; + keep.append(alloc, .{ .doc_id = entry.doc_id, .ver = entry.ver }) catch { + oom = true; + break; + }; + } + } + + if (oom) { + // OOM during keep build — discard partial keep list and leave + // the original all_versions intact so no entries are lost. + keep.deinit(alloc); + return 0; + } + + // Only remove history entries for versions we are actually freeing. + i = 0; + while (i < self.all_versions.items.len) : (i += 1) { + const entry = self.all_versions.items[i]; + const is_latest = if (self.latest.get(entry.doc_id)) |cur| + cur.page_no == entry.ver.page_no and cur.page_off == entry.ver.page_off + else + false; + + if (!is_latest and entry.ver.epoch < min_epoch) { + _ = self.history.remove(chainKey(entry.ver.page_no, entry.ver.page_off)); } } diff --git a/src/page.zig b/src/page.zig index 114fbc7..f5933bb 100644 --- a/src/page.zig +++ b/src/page.zig @@ -35,17 +35,19 @@ pub const PageFile = struct { free_head: std.atomic.Value(u32), // head of free-list (0 = empty) next_alloc: std.atomic.Value(u32), // next never-allocated page mu: std.Thread.Mutex, + leaf_shards: [64]std.Thread.Mutex, // per-page-shard locks for leafAppend pub fn open(path: []const u8) !PageFile { var path_buf: [std.fs.max_path_bytes + 1]u8 = undefined; const path_z = try std.fmt.bufPrintZ(&path_buf, "{s}", .{path}); const mm = try mmap.MmapFile.open(path_z, PAGE_SIZE * 16); - const page_count: u32 = @intCast(mm.len / PAGE_SIZE); + const page_count: u32 = @intCast(mm.dataLen() / PAGE_SIZE); return PageFile{ .mm = mm, .free_head = std.atomic.Value(u32).init(0), .next_alloc = std.atomic.Value(u32).init(page_count), .mu = .{}, + .leaf_shards = [1]std.Thread.Mutex{.{}} ** 64, }; } @@ -76,7 +78,7 @@ pub const PageFile = struct { // Extend the file. const pno = self.next_alloc.fetchAdd(1, .seq_cst); const needed = (@as(usize, pno) + 1) * PAGE_SIZE; - if (needed > self.mm.capacity) try self.mm.grow(needed); + if (needed > self.mm.capacity.load(.acquire)) try self.mm.grow(needed); const ph = self.pageHeader(pno); ph.* = std.mem.zeroes(PageHeader); @@ -111,6 +113,9 @@ pub const PageFile = struct { /// Append bytes to a leaf page. Returns offset within the usable area, /// or null if there isn't enough space. pub fn leafAppend(self: *PageFile, pno: u32, data: []const u8) ?u16 { + self.leaf_shards[pno % 64].lock(); + defer self.leaf_shards[pno % 64].unlock(); + const ph = self.pageHeader(pno); const used = ph.used_bytes; if (@as(usize, used) + data.len > PAGE_USABLE) return null; @@ -121,11 +126,27 @@ pub const PageFile = struct { return used; } + /// Reserve `len` bytes on a leaf page and return a writable slice + offset. + /// Caller writes directly into the page, avoiding a double copy. + pub fn leafReserve(self: *PageFile, pno: u32, len: usize) ?struct { buf: []u8, off: u16 } { + self.leaf_shards[pno % 64].lock(); + defer self.leaf_shards[pno % 64].unlock(); + + const ph = self.pageHeader(pno); + const used = ph.used_bytes; + if (@as(usize, used) + len > PAGE_USABLE) return null; + const dest = self.pageData(pno); + ph.used_bytes += @intCast(len); + ph.doc_count += 1; + return .{ .buf = dest[used..][0..len], .off = used }; + } + /// Read `len` bytes from page `pno` at usable-area offset `off`. - pub fn leafRead(self: *PageFile, pno: u32, off: u16, len: usize) []const u8 { + /// Returns null if the read would extend past the page boundary. + pub fn leafRead(self: *PageFile, pno: u32, off: u16, len: usize) ?[]const u8 { + if (@as(usize, off) + len > PAGE_USABLE) return null; const data = self.pageData(pno); - const end = @min(@as(usize, off) + len, PAGE_USABLE); - return data[off..end]; + return data[off..][0..len]; } /// Iterate all documents on a leaf page, calling `cb` for each raw record slice. diff --git a/src/query.zig b/src/query.zig index 1128f44..828a144 100644 --- a/src/query.zig +++ b/src/query.zig @@ -205,7 +205,7 @@ pub fn extractField(json: []const u8, field: []const u8) ?[]const u8 { /// Handles strings, numbers, booleans, null, objects, and arrays. fn extractRawValue(json: []const u8, key: []const u8) ?[]const u8 { // Build needle: "key" - var needle_buf: [256]u8 = undefined; + var needle_buf: [1024]u8 = undefined; const needle = std.fmt.bufPrint(&needle_buf, "\"{s}\"", .{key}) catch return null; var pos: usize = 0; diff --git a/src/replication/integration.zig b/src/replication/integration.zig index 7879235..a1b8a44 100644 --- a/src/replication/integration.zig +++ b/src/replication/integration.zig @@ -71,6 +71,11 @@ pub const ReplicatedDB = struct { receiver: ?PeerReceiver, alloc: std.mem.Allocator, next_txn_id: std.atomic.Value(u64), + executor_ptr: ?*calvin.CalvinExecutor, + + /// Module-level context pointer for the executor callback. + /// Set once during startReplication; asserted to be set only once. + var g_ctx: ?*ReplicatedDB = null; pub fn init( alloc: std.mem.Allocator, @@ -87,6 +92,7 @@ pub const ReplicatedDB = struct { .receiver = null, .alloc = alloc, .next_txn_id = std.atomic.Value(u64).init(1), + .executor_ptr = null, }; if (config.enabled) { @@ -100,24 +106,32 @@ pub const ReplicatedDB = struct { pub fn deinit(self: *ReplicatedDB) void { if (self.receiver) |*r| r.stop(); + if (self.executor_ptr) |ptr| self.alloc.destroy(ptr); if (self.repl_mgr) |*mgr| mgr.deinit(); + if (g_ctx == self) g_ctx = null; } /// Start replication (leader batch loop + replica receiver). pub fn startReplication(self: *ReplicatedDB) !void { if (self.repl_mgr) |*mgr| { - // Set global state for the executor callback - g_db_handle = self.db_handle; - g_ops = self.ops; + // Set module-level context for the executor callback (assert set-once). + std.debug.assert(g_ctx == null); + g_ctx = self; mgr.setBatchReadyHook(self, leaderBatchReady); try mgr.start(); if (!self.config.is_leader) { + // Bug 2 fix: heap-allocate the executor to a stable location + // instead of taking a pointer into the optional field. + const executor_ptr = try self.alloc.create(calvin.CalvinExecutor); + executor_ptr.* = mgr.executor; + self.executor_ptr = executor_ptr; + self.receiver = PeerReceiver.init( self.alloc, self.config.repl_port, - &mgr.executor, + executor_ptr, &executeTransaction, ); const t = try std.Thread.spawn(.{}, PeerReceiver.run, .{&self.receiver.?}); @@ -146,13 +160,18 @@ pub const ReplicatedDB = struct { var data_buf: [4096]u8 = undefined; const data_len = packTxnData(&data_buf, branch_name, col_name, key, value); + if (data_len == 0) return null; const data = self.alloc.alloc(u8, data_len) catch return null; - @memcpy(data, data_buf[0..data_len]); - const ws = self.alloc.alloc(u64, 1) catch return null; + const ws = self.alloc.alloc(u64, 1) catch { + self.alloc.free(data); + return null; + }; ws[0] = key_hash; + @memcpy(data, data_buf[0..data_len]); + const txn_id = self.next_txn_id.fetchAdd(1, .monotonic); const txn = Transaction{ .txn_id = txn_id, @@ -165,7 +184,11 @@ pub const ReplicatedDB = struct { .owns_buffers = true, }; - mgr.submit(txn) catch return null; + mgr.submit(txn) catch { + self.alloc.free(data); + self.alloc.free(ws); + return null; + }; return txn_id; } @@ -190,13 +213,18 @@ pub const ReplicatedDB = struct { var data_buf: [4096]u8 = undefined; const data_len = packTxnData(&data_buf, branch_name, col_name, key, value); + if (data_len == 0) return false; const data = self.alloc.alloc(u8, data_len) catch return false; - @memcpy(data, data_buf[0..data_len]); - const ws = self.alloc.alloc(u64, 1) catch return false; + const ws = self.alloc.alloc(u64, 1) catch { + self.alloc.free(data); + return false; + }; ws[0] = key_hash; + @memcpy(data, data_buf[0..data_len]); + const txn = Transaction{ .txn_id = self.next_txn_id.fetchAdd(1, .monotonic), .txn_type = .put, // update = put with same key @@ -208,7 +236,11 @@ pub const ReplicatedDB = struct { .owns_buffers = true, }; - mgr.submit(txn) catch return false; + mgr.submit(txn) catch { + self.alloc.free(data); + self.alloc.free(ws); + return false; + }; return true; } @@ -232,13 +264,18 @@ pub const ReplicatedDB = struct { var data_buf: [256]u8 = undefined; const data_len = packTxnData(&data_buf, branch_name, col_name, key, ""); + if (data_len == 0) return false; const data = self.alloc.alloc(u8, data_len) catch return false; - @memcpy(data, data_buf[0..data_len]); - const ws = self.alloc.alloc(u64, 1) catch return false; + const ws = self.alloc.alloc(u64, 1) catch { + self.alloc.free(data); + return false; + }; ws[0] = key_hash; + @memcpy(data, data_buf[0..data_len]); + const txn = Transaction{ .txn_id = self.next_txn_id.fetchAdd(1, .monotonic), .txn_type = .delete, @@ -250,7 +287,11 @@ pub const ReplicatedDB = struct { .owns_buffers = true, }; - mgr.submit(txn) catch return false; + mgr.submit(txn) catch { + self.alloc.free(data); + self.alloc.free(ws); + return false; + }; return true; } @@ -276,6 +317,15 @@ pub const ReplicatedDB = struct { // Format: [branch_len:u8][branch][col_name_len:u8][col_name][key_len:u16 LE][key][value] fn packTxnData(buf: []u8, branch_name: []const u8, col_name: []const u8, key: []const u8, value: []const u8) usize { + // Bounds check: 1 (branch_len) + branch + 1 (col_len) + col + 2 (key_len) + key + value + const total = 4 + branch_name.len + col_name.len + key.len + value.len; + if (total > buf.len) return 0; + + // Check that branch_name and col_name lengths fit in u8 + if (branch_name.len > std.math.maxInt(u8)) return 0; + if (col_name.len > std.math.maxInt(u8)) return 0; + if (key.len > std.math.maxInt(u16)) return 0; + var pos: usize = 0; buf[pos] = @intCast(branch_name.len); pos += 1; @@ -317,9 +367,9 @@ fn unpackTxnData(data: []const u8) ?struct { branch: []const u8, col: []const u8 } // ── Transaction executor (called by CalvinExecutor) ────────────────────────── - -var g_db_handle: ?*anyopaque = null; -var g_ops: DBOps = .{}; +// Bug 1 fix: replaced global g_db_handle/g_ops with ReplicatedDB.g_ctx. +// The executor callback signature is fixed (no context parameter), so we use +// a module-level pointer that is set once during startReplication. const BranchCall = struct { branch_buf: [64]u8 = [_]u8{0} ** 64, @@ -355,30 +405,32 @@ const BranchCall = struct { }; fn executeTransaction(txn: *const Transaction) anyerror!void { - const db = g_db_handle orelse return error.DatabaseNotInitialized; + const ctx = ReplicatedDB.g_ctx orelse return error.DatabaseNotInitialized; + const db = ctx.db_handle; + const ops = ctx.ops; const parsed = unpackTxnData(txn.data) orelse return error.InvalidData; switch (txn.txn_type) { .put => { // put = insert or update (same semantics: upsert) if (parsed.branch.len > 0) { - if (g_ops.insert_branch_fn) |f| { + if (ops.insert_branch_fn) |f| { _ = f(db, parsed.branch, parsed.col, parsed.key, parsed.val); return; } } - if (g_ops.insert_fn) |f| { + if (ops.insert_fn) |f| { _ = f(db, parsed.col, parsed.key, parsed.val); } }, .delete => { if (parsed.branch.len > 0) { - if (g_ops.delete_branch_fn) |f| { + if (ops.delete_branch_fn) |f| { _ = f(db, parsed.branch, parsed.col, parsed.key); return; } } - if (g_ops.delete_fn) |f| { + if (ops.delete_fn) |f| { _ = f(db, parsed.col, parsed.key); } }, @@ -386,8 +438,8 @@ fn executeTransaction(txn: *const Transaction) anyerror!void { } } -fn leaderBatchReady(ctx: *anyopaque, batch: *const sequencer.Batch) anyerror!void { - const self: *ReplicatedDB = @ptrCast(@alignCast(ctx)); +fn leaderBatchReady(ctx_ptr: *anyopaque, batch: *const sequencer.Batch) anyerror!void { + const self: *ReplicatedDB = @ptrCast(@alignCast(ctx_ptr)); const buf = try self.alloc.alloc(u8, 1 << 20); defer self.alloc.free(buf); @@ -511,8 +563,13 @@ test "executeTransaction routes branch payloads to branch callbacks" { }; var seen = BranchCall{}; - g_db_handle = @ptrCast(&seen); - g_ops = ops; + // Use a temporary ReplicatedDB to set up g_ctx for the executor callback. + var rdb = try ReplicatedDB.init(std.testing.allocator, @ptrCast(&seen), ops, .{}); + defer rdb.deinit(); + ReplicatedDB.g_ctx = &rdb; + defer { + ReplicatedDB.g_ctx = null; + } var buf: [256]u8 = undefined; const len = packTxnData(&buf, "feature-x", "users", "alice", "{\"name\":\"Alice\"}"); @@ -554,8 +611,10 @@ test "leader batch hook executes queued writes locally" { }); defer rdb.deinit(); - g_db_handle = rdb.db_handle; - g_ops = rdb.ops; + ReplicatedDB.g_ctx = &rdb; + defer { + ReplicatedDB.g_ctx = null; + } try rdb.repl_mgr.?.executor.setLocalPartitions(&.{@as(u16, @intCast(fnv1a("alice") % 256))}); const txn_id = rdb.insert("users", "alice", "{\"name\":\"Alice\"}"); diff --git a/src/replication/router.zig b/src/replication/router.zig index 32d9522..1bd6144 100644 --- a/src/replication/router.zig +++ b/src/replication/router.zig @@ -211,14 +211,15 @@ pub const Router = struct { }; fn parseU16(s: []const u8) ?u16 { + if (s.len == 0) return null; + if (s[0] < '0' or s[0] > '9') return null; var val: u16 = 0; for (s) |c| { if (c >= '0' and c <= '9') { - val = val *% 10 +% (@as(u16, c) - '0'); + val = std.math.mul(u16, val, 10) catch return null; + val = std.math.add(u16, val, @as(u16, c) - '0') catch return null; } else break; } - if (s.len == 0) return null; - if (s[0] < '0' or s[0] > '9') return null; return val; } diff --git a/src/replication/sequencer.zig b/src/replication/sequencer.zig index 99a4925..076459b 100644 --- a/src/replication/sequencer.zig +++ b/src/replication/sequencer.zig @@ -19,9 +19,10 @@ pub const Batch = struct { epoch: u64, sequence_start: u64, transactions: []Transaction, + owns_transactions: bool = true, pub fn deinit(self: *Batch, alloc: std.mem.Allocator) void { - alloc.free(self.transactions); + if (self.owns_transactions) alloc.free(self.transactions); } pub fn deinitDeep(self: *Batch, alloc: std.mem.Allocator) void { @@ -32,7 +33,7 @@ pub const Batch = struct { if (txn.read_set.len > 0) alloc.free(txn.read_set); } } - alloc.free(self.transactions); + if (self.owns_transactions) alloc.free(self.transactions); } }; @@ -262,6 +263,7 @@ test "sequencer: detect write-write conflicts" { .epoch = 0, .sequence_start = 0, .transactions = &txns, + .owns_transactions = false, }; const conflicts = try Sequencer.detectConflicts(&batch, alloc); diff --git a/src/replication/shard.zig b/src/replication/shard.zig index 5c8ec02..bfbf3c7 100644 --- a/src/replication/shard.zig +++ b/src/replication/shard.zig @@ -162,6 +162,8 @@ pub const ShardManager = struct { self.mu.lock(); defer self.mu.unlock(); try self.shard_map.put(self.alloc, entry.partition_id, entry); + // Also update the consistent hash ring so routeKey reflects new partitions. + try self.ring.addNode(entry.node_id); } /// Look up which node owns a partition. diff --git a/src/server.zig b/src/server.zig index 7a61a20..0fd30c1 100644 --- a/src/server.zig +++ b/src/server.zig @@ -17,12 +17,15 @@ const collection = @import("collection.zig"); const Database = collection.Database; const MAX_REQ = 65536; // 64 KiB (initial read) -const MAX_RESP = 131072; // 128 KiB +const MAX_RESP = 65536; // 64 KiB const MAX_BODY = 65536; // 64 KiB -const MAX_BULK = 16 * 1024 * 1024; // 16 MiB for bulk inserts +const MAX_BULK = 256 * 1024 * 1024; // 256 MiB for bulk inserts // Heap-allocated per-connection buffers (threadlocal pointers set in handleConn). // This avoids large threadlocal TLS segments that break in Release mode on macOS. +// Note: connection threads use 2 MiB stacks (up from 256 KiB) because the +// Collection.open call chain is deep and Zig's debug-mode stack probes need +// headroom on macOS ARM64. Production builds should use -Doptimize=ReleaseFast. const ConnBufs = struct { req: [MAX_REQ]u8, resp: [MAX_RESP]u8, @@ -111,11 +114,14 @@ pub const Server = struct { }; // Reject if at connection limit to prevent OOM from thread exhaustion. if (self.active_conns.load(.monotonic) >= MAX_CONNECTIONS) { + // Return HTTP 503 instead of silent RST so clients can retry. + const reject = "HTTP/1.1 503 Service Unavailable\r\nContent-Length: 31\r\nConnection: close\r\n\r\n{\"error\":\"too many connections\"}"; + conn.stream.writeAll(reject) catch {}; conn.stream.close(); _ = self.err_count.fetchAdd(1, .monotonic); continue; } - const t = std.Thread.spawn(.{}, handleConnWrapped, .{self, conn}) catch { + const t = std.Thread.spawn(.{ .stack_size = 2 * 1024 * 1024 }, handleConnWrapped, .{self, conn}) catch { conn.stream.close(); _ = self.err_count.fetchAdd(1, .monotonic); continue; @@ -149,11 +155,13 @@ pub const Server = struct { const stream = std.net.Stream{ .handle = client_fd }; const conn = std.net.Server.Connection{ .stream = stream, .address = std.net.Address.initUnix(path) catch continue }; if (self.active_conns.load(.monotonic) >= MAX_CONNECTIONS) { + const reject = "HTTP/1.1 503 Service Unavailable\r\nContent-Length: 31\r\nConnection: close\r\n\r\n{\"error\":\"too many connections\"}"; + conn.stream.writeAll(reject) catch {}; conn.stream.close(); _ = self.err_count.fetchAdd(1, .monotonic); continue; } - const t = std.Thread.spawn(.{}, handleConnWrapped, .{ self, conn }) catch { + const t = std.Thread.spawn(.{ .stack_size = 2 * 1024 * 1024 }, handleConnWrapped, .{ self, conn }) catch { conn.stream.close(); continue; }; @@ -223,38 +231,57 @@ fn handleConn(srv: *Server, conn: std.net.Server.Connection) void { if (n == 0) return; _ = srv.req_count.fetchAdd(1, .monotonic); - // For bulk inserts: read the full body based on Content-Length. - // The initial read may only contain part of a large body. const initial = bufs.req[0..n]; + + // WebSocket upgrade: switch to persistent framed mode. + if (headerContains(initial, "upgrade", "websocket")) { + handleWebSocket(srv, conn, initial) catch {}; + return; // WS handler owns the connection until close + } + + // Read the full body based on Content-Length if we don't have it yet. const content_length = extractContentLength(initial); - if (content_length > MAX_REQ and content_length <= MAX_BULK) { - // Check this is actually a bulk request before allocating - const is_bulk = std.mem.indexOf(u8, initial[0..@min(n, 256)], "/bulk") != null; - if (is_bulk) { - // Find where headers end - const header_end = if (std.mem.indexOf(u8, initial, "\r\n\r\n")) |p| p + 4 - else if (std.mem.indexOf(u8, initial, "\n\n")) |p| p + 2 - else n; - const total_size = header_end + content_length; - if (total_size <= MAX_BULK) { - const big_buf = std.heap.page_allocator.alloc(u8, total_size) catch { - const resp_len = dispatch(srv, initial, std.heap.page_allocator); - conn.stream.writeAll(bufs.resp[0..resp_len]) catch return; - continue; - }; - defer std.heap.page_allocator.free(big_buf); - @memcpy(big_buf[0..n], initial); - // Read remaining bytes - while (n < total_size) { - const r = conn.stream.read(big_buf[n..total_size]) catch break; - if (r == 0) break; - n += r; - } - const resp_len = dispatch(srv, big_buf[0..n], std.heap.page_allocator); + const is_bulk = std.mem.indexOf(u8, initial[0..@min(n, 256)], "/bulk") != null; + + if (content_length > MAX_REQ and content_length <= MAX_BULK and is_bulk) { + // Large bulk: allocate a big buffer and read the full body. + const header_end = if (std.mem.indexOf(u8, initial, "\r\n\r\n")) |p| p + 4 + else if (std.mem.indexOf(u8, initial, "\n\n")) |p| p + 2 + else n; + const total_size = header_end + content_length; + if (total_size <= MAX_BULK) { + const big_buf = std.heap.page_allocator.alloc(u8, total_size) catch { + const resp_len = dispatch(srv, initial, std.heap.page_allocator); conn.stream.writeAll(bufs.resp[0..resp_len]) catch return; continue; + }; + defer std.heap.page_allocator.free(big_buf); + @memcpy(big_buf[0..n], initial); + // Read remaining bytes + while (n < total_size) { + const r = conn.stream.read(big_buf[n..total_size]) catch break; + if (r == 0) break; + n += r; } + const resp_len = dispatch(srv, big_buf[0..n], std.heap.page_allocator); + conn.stream.writeAll(bufs.resp[0..resp_len]) catch return; + continue; + } + } else if (content_length > 0 and content_length <= MAX_REQ) { + // Body fits in the req buffer but the initial read may not have + // received it all. Keep reading until we have the full body. + const header_end = if (std.mem.indexOf(u8, initial, "\r\n\r\n")) |p| p + 4 + else if (std.mem.indexOf(u8, initial, "\n\n")) |p| p + 2 + else n; + const total_size = @min(header_end + content_length, bufs.req.len); + while (n < total_size) { + const r = conn.stream.read(bufs.req[n..total_size]) catch break; + if (r == 0) break; + n += r; } + const resp_len = dispatch(srv, bufs.req[0..n], std.heap.page_allocator); + conn.stream.writeAll(bufs.resp[0..resp_len]) catch return; + continue; } const resp_len = dispatch(srv, initial, std.heap.page_allocator); @@ -336,7 +363,7 @@ fn dispatch(srv: *Server, raw: []const u8, alloc: std.mem.Allocator) usize { const q_raw = qparam(query, "q") orelse return err(400, "missing q parameter"); var decode_buf: [4096]u8 = undefined; const q = urlDecode(q_raw, &decode_buf) orelse q_raw; - const limit_val: u32 = qparamInt(query, "limit") orelse 50; + const limit_val: u32 = @min(qparamInt(query, "limit") orelse 50, 500); return handleSearch(srv, requestTenant(raw, query), col_name, q, limit_val, alloc); } @@ -346,7 +373,7 @@ fn dispatch(srv: *Server, raw: []const u8, alloc: std.mem.Allocator) usize { const q_raw = qparam(query, "q") orelse return err(400, "missing q parameter"); var decode_buf: [4096]u8 = undefined; const q = urlDecode(q_raw, &decode_buf) orelse q_raw; - const limit_val: u32 = qparamInt(query, "limit") orelse 20; + const limit_val: u32 = @min(qparamInt(query, "limit") orelse 20, 100); return handleDiscoverContext(srv, requestTenant(raw, query), col_name, q, limit_val, alloc); } @@ -453,55 +480,109 @@ fn doInsert(srv: *Server, tenant_id: []const u8, col_name: []const u8, key: []co const col = srv.db.collectionForTenant(tenant_id, col_name) catch return err(500, "open collection failed"); const doc_id = col.insert(key, value) catch return err(500, "insert failed"); srv.recordQueryCost(tenant_id, "insert", 1, value.len, start_ns); + var esc_key_buf: [1024]u8 = undefined; + var esc_col_buf: [1024]u8 = undefined; + var esc_tid_buf: [1024]u8 = undefined; + const esc_key = jsonEscape(key, &esc_key_buf); + const esc_col = jsonEscape(col_name, &esc_col_buf); + const esc_tid = jsonEscape(tenant_id, &esc_tid_buf); var fbs = std.io.fixedBufferStream(getBodyBuf()); std.fmt.format(fbs.writer(), "{{\"doc_id\":{d},\"key\":\"{s}\",\"collection\":\"{s}\",\"tenant\":\"{s}\"}}", - .{ doc_id, key, col_name, tenant_id }) catch {}; + .{ doc_id, esc_key, esc_col, esc_tid }) catch {}; return ok(getBodyBuf()[0..fbs.pos]); } -/// POST /db/:col/bulk — insert multiple documents in one request. -/// Body: NDJSON — one {"key":"...","value":"..."} per line. -/// Response: {"inserted":N,"errors":M,"collection":"...","tenant":"..."} -fn handleBulkInsert(srv: *Server, tenant_id: []const u8, col_name: []const u8, body: []const u8, alloc: std.mem.Allocator) usize { - _ = alloc; - const start_ns = std.time.nanoTimestamp(); - srv.db.recordTenantOperation(tenant_id) catch return err(429, "tenant ops quota exceeded"); - const col = srv.db.collectionForTenant(tenant_id, col_name) catch return err(500, "open collection failed"); +/// SIMD-accelerated scan for the next newline byte starting at `start`. +/// Uses 16-byte vector comparison (SSE2 on x86_64, NEON on aarch64). +fn simdFindNewline(data: []const u8, start: usize) ?usize { + const VEC_LEN = 16; + const Vec = @Vector(VEC_LEN, u8); + const nl_vec: Vec = @splat('\n'); + + var pos = start; + + // Vector scan: compare VEC_LEN bytes per iteration. + while (pos + VEC_LEN <= data.len) { + const chunk: Vec = data[pos..][0..VEC_LEN].*; + const mask: u16 = @bitCast(chunk == nl_vec); + if (mask != 0) return pos + @ctz(mask); + pos += VEC_LEN; + } + + // Scalar tail for remaining bytes. + while (pos < data.len) : (pos += 1) { + if (data[pos] == '\n') return pos; + } + return null; +} + +/// Shared NDJSON parse + insertBatch logic used by both HTTP and WS paths. +const NdjsonResult = struct { inserted: u32, errors: u32, total_bytes: u64 }; +fn parseAndInsertNdjson(col: *collection.Collection, body: []const u8) NdjsonResult { var inserted: u32 = 0; var errors: u32 = 0; var total_bytes: u64 = 0; - // Parse NDJSON: iterate lines, each is a {"key":"...","value":...} object + // Keep outer batch small (32KB stack vs 128KB at 4096) — insertBatch + // handles its own internal chunking so throughput is the same. + const MaxLines = 1024; + var key_buf: [MaxLines][]const u8 = undefined; + var val_buf: [MaxLines][]const u8 = undefined; + var n_lines: usize = 0; + var pos: usize = 0; while (pos < body.len) { - // Find end of line - const line_end = std.mem.indexOfScalarPos(u8, body, pos, '\n') orelse body.len; + const line_end = simdFindNewline(body, pos) orelse body.len; const line = std.mem.trim(u8, body[pos..line_end], " \t\r"); pos = line_end + 1; - if (line.len < 2) continue; // skip empty lines + if (line.len < 2) continue; - // Extract key from this JSON line const key = jsonStr(line, "key") orelse continue; - // Extract value field; fall back to full line for backwards compat. const value = jsonValue(line, "value") orelse line; - _ = col.insert(key, value) catch { - errors += 1; - continue; - }; - inserted += 1; + key_buf[n_lines] = key; + val_buf[n_lines] = value; total_bytes += line.len; + n_lines += 1; + + if (n_lines >= MaxLines) { + const r = col.insertBatch(key_buf[0..n_lines], val_buf[0..n_lines]); + inserted += r.inserted; + errors += r.errors; + n_lines = 0; + } } + if (n_lines > 0) { + const r = col.insertBatch(key_buf[0..n_lines], val_buf[0..n_lines]); + inserted += r.inserted; + errors += r.errors; + } + return .{ .inserted = inserted, .errors = errors, .total_bytes = total_bytes }; +} - srv.recordQueryCost(tenant_id, "bulk_insert", inserted, total_bytes, start_ns); +/// POST /db/:col/bulk — NDJSON batch insert with SIMD line scanning. +fn handleBulkInsert(srv: *Server, tenant_id: []const u8, col_name: []const u8, body: []const u8, alloc: std.mem.Allocator) usize { + _ = alloc; + const start_ns = std.time.nanoTimestamp(); + srv.db.recordTenantOperation(tenant_id) catch return err(429, "tenant ops quota exceeded"); + srv.db.ensureTenantStorageAvailable(tenant_id, body.len) catch return err(429, "tenant storage quota exceeded"); + const col = srv.db.collectionForTenant(tenant_id, col_name) catch return err(500, "open collection failed"); + const r = parseAndInsertNdjson(col, body); + + srv.recordQueryCost(tenant_id, "bulk_insert", r.inserted, r.total_bytes, start_ns); + + var esc_col_buf: [1024]u8 = undefined; + var esc_tid_buf: [1024]u8 = undefined; + const esc_col = jsonEscape(col_name, &esc_col_buf); + const esc_tid = jsonEscape(tenant_id, &esc_tid_buf); var fbs = std.io.fixedBufferStream(getBodyBuf()); std.fmt.format(fbs.writer(), "{{\"inserted\":{d},\"errors\":{d},\"collection\":\"{s}\",\"tenant\":\"{s}\"}}", - .{ inserted, errors, col_name, tenant_id }) catch {}; + .{ r.inserted, r.errors, esc_col, esc_tid }) catch {}; return ok(getBodyBuf()[0..fbs.pos]); } fn handleGet(srv: *Server, tenant_id: []const u8, col_name: []const u8, key: []const u8, as_of: ?i64) usize { @@ -520,14 +601,19 @@ fn handleGet(srv: *Server, tenant_id: []const u8, col_name: []const u8, key: []c var fbs = std.io.fixedBufferStream(resp[HEADER_RESERVE..]); const val = if (d.value.len > 0) d.value else "{}"; const is_json = val.len > 0 and (val[0] == '{' or val[0] == '[' or val[0] == '"'); + var esc_key_buf: [1024]u8 = undefined; + const esc_key = jsonEscape(d.key, &esc_key_buf); const w = fbs.writer(); if (is_json) { - std.fmt.format(w, "{{\"doc_id\":{d},\"key\":\"{s}\",\"version\":{d},\"value\":{s}}}", .{ d.header.doc_id, d.key, d.header.version, val }) catch {}; + std.fmt.format(w, "{{\"doc_id\":{d},\"key\":\"{s}\",\"version\":{d},\"value\":{s}}}", .{ d.header.doc_id, esc_key, d.header.version, val }) catch {}; } else { - std.fmt.format(w, "{{\"doc_id\":{d},\"key\":\"{s}\",\"version\":{d},\"value\":\"", .{ d.header.doc_id, d.key, d.header.version }) catch {}; + std.fmt.format(w, "{{\"doc_id\":{d},\"key\":\"{s}\",\"version\":{d},\"value\":\"", .{ d.header.doc_id, esc_key, d.header.version }) catch {}; for (val) |ch| { - if (ch == '"' or ch == '\\') w.writeByte('\\') catch {}; + if (ch == '"' or ch == '\\') { w.writeByte('\\') catch {}; w.writeByte(ch) catch {}; continue; } if (ch == '\n') { w.writeAll("\\n") catch {}; continue; } + if (ch == '\r') { w.writeAll("\\r") catch {}; continue; } + if (ch == '\t') { w.writeAll("\\t") catch {}; continue; } + if (ch < 0x20) continue; w.writeByte(ch) catch {}; } w.writeAll("\"}") catch {}; @@ -543,7 +629,7 @@ fn handleGet(srv: *Server, tenant_id: []const u8, col_name: []const u8, key: []c // Move body right after headers (memmove if needed) if (hdr_len < HEADER_RESERVE) { - std.mem.copyForwards(u8, resp[hdr_len .. hdr_len + body_len], resp[HEADER_RESERVE .. HEADER_RESERVE + body_len]); + std.mem.copyBackwards(u8, resp[hdr_len .. hdr_len + body_len], resp[HEADER_RESERVE .. HEADER_RESERVE + body_len]); } return hdr_len + body_len; } @@ -572,7 +658,7 @@ fn handleDelete(srv: *Server, tenant_id: []const u8, col_name: []const u8, key: fn handleScan(srv: *Server, tenant_id: []const u8, col_name: []const u8, query_str: []const u8, as_of: ?i64, alloc: std.mem.Allocator) usize { const start_ns = std.time.nanoTimestamp(); - const limit: u32 = qparamInt(query_str, "limit") orelse 20; + const limit: u32 = @min(qparamInt(query_str, "limit") orelse 20, 1000); const offset: u32 = qparamInt(query_str, "offset") orelse 0; srv.db.recordTenantOperation(tenant_id) catch return err(429, "tenant ops quota exceeded"); const col = srv.db.collectionForTenant(tenant_id, col_name) catch return err(500, "open collection failed"); @@ -585,30 +671,55 @@ fn handleScan(srv: *Server, tenant_id: []const u8, col_name: []const u8, query_s for (result.docs) |d| bytes_read += d.key.len + d.value.len; srv.recordQueryCost(tenant_id, "scan", result.docs.len, bytes_read, start_ns); + var esc_tid_buf: [1024]u8 = undefined; + var esc_col_buf: [1024]u8 = undefined; + const esc_tid = jsonEscape(tenant_id, &esc_tid_buf); + const esc_col = jsonEscape(col_name, &esc_col_buf); var fbs = std.io.fixedBufferStream(getBodyBuf()); const w = fbs.writer(); std.fmt.format(w, "{{\"tenant\":\"{s}\",\"collection\":\"{s}\",\"count\":{d},\"docs\":[", - .{ tenant_id, col_name, result.docs.len }) catch {}; + .{ esc_tid, esc_col, result.docs.len }) catch {}; + var truncated = false; for (result.docs, 0..) |d, i| { + const saved_pos = fbs.pos; if (i > 0) w.writeByte(',') catch {}; + var esc_dk_buf: [1024]u8 = undefined; + const esc_dk = jsonEscape(d.key, &esc_dk_buf); const val = if (d.value.len > 0) d.value else "{}"; const is_json = val.len > 0 and (val[0] == '{' or val[0] == '[' or val[0] == '"'); if (is_json) { - std.fmt.format(w, "{{\"doc_id\":{d},\"key\":\"{s}\",\"version\":{d},\"value\":{s}}}", .{ d.header.doc_id, d.key, d.header.version, val }) catch {}; + std.fmt.format(w, "{{\"doc_id\":{d},\"key\":\"{s}\",\"version\":{d},\"value\":{s}}}", .{ d.header.doc_id, esc_dk, d.header.version, val }) catch {}; } else { - std.fmt.format(w, "{{\"doc_id\":{d},\"key\":\"{s}\",\"version\":{d},\"value\":\"", .{ d.header.doc_id, d.key, d.header.version }) catch {}; + std.fmt.format(w, "{{\"doc_id\":{d},\"key\":\"{s}\",\"version\":{d},\"value\":\"", .{ d.header.doc_id, esc_dk, d.header.version }) catch {}; for (val) |ch| { - if (ch == '"' or ch == '\\') w.writeByte('\\') catch {}; + if (ch == '"' or ch == '\\') { w.writeByte('\\') catch {}; w.writeByte(ch) catch {}; continue; } if (ch == '\n') { w.writeAll("\\n") catch {}; continue; } + if (ch == '\r') { w.writeAll("\\r") catch {}; continue; } + if (ch == '\t') { w.writeAll("\\t") catch {}; continue; } + if (ch < 0x20) continue; w.writeByte(ch) catch {}; } w.writeAll("\"}") catch {}; } + // If this doc pushed us near the buffer limit, rewind to discard the + // partial write and stop — this prevents broken JSON in the response. + if (fbs.pos + 64 >= MAX_BODY) { + fbs.pos = saved_pos; + truncated = true; + break; + } } w.writeAll("]}") catch {}; + if (truncated) { + // Overwrite the trailing '}' with ',"truncated":true}' + if (fbs.pos >= 1) { + fbs.pos -= 1; // back over '}' + w.writeAll(",\"truncated\":true}") catch {}; + } + } return ok(getBodyBuf()[0..fbs.pos]); -} +} fn handleDrop(srv: *Server, tenant_id: []const u8, col_name: []const u8) usize { const start_ns = std.time.nanoTimestamp(); srv.db.recordTenantOperation(tenant_id) catch return err(429, "tenant ops quota exceeded"); @@ -638,25 +749,43 @@ fn handleSearch(srv: *Server, tenant_id: []const u8, col_name: []const u8, query std.fmt.format(w, "\",\"hits\":{d},\"candidates\":{d},\"total_docs\":{d},\"total_files\":{d},\"results\":[", .{ result.docs.len, result.candidate_paths.len, col.docCount(), result.total_files }) catch {}; + var truncated = false; for (result.docs, 0..) |d, i| { + const saved_pos = fbs.pos; if (i > 0) w.writeByte(',') catch {}; // Output value as valid JSON — objects/arrays as-is, strings quoted + var esc_dk_buf: [1024]u8 = undefined; + const esc_dk = jsonEscape(d.key, &esc_dk_buf); const val = if (d.value.len > 0) d.value else "{}"; const is_json = val.len > 0 and (val[0] == '{' or val[0] == '[' or val[0] == '"'); if (is_json) { - std.fmt.format(w, "{{\"doc_id\":{d},\"key\":\"{s}\",\"value\":{s}}}", .{ d.header.doc_id, d.key, val }) catch {}; + std.fmt.format(w, "{{\"doc_id\":{d},\"key\":\"{s}\",\"value\":{s}}}", .{ d.header.doc_id, esc_dk, val }) catch {}; } else { w.writeAll("{\"doc_id\":") catch {}; - std.fmt.format(w, "{d},\"key\":\"{s}\",\"value\":\"", .{ d.header.doc_id, d.key }) catch {}; + std.fmt.format(w, "{d},\"key\":\"{s}\",\"value\":\"", .{ d.header.doc_id, esc_dk }) catch {}; for (val) |ch| { - if (ch == '"' or ch == '\\') w.writeByte('\\') catch {}; + if (ch == '"' or ch == '\\') { w.writeByte('\\') catch {}; w.writeByte(ch) catch {}; continue; } if (ch == '\n') { w.writeAll("\\n") catch {}; continue; } + if (ch == '\r') { w.writeAll("\\r") catch {}; continue; } + if (ch == '\t') { w.writeAll("\\t") catch {}; continue; } + if (ch < 0x20) continue; w.writeByte(ch) catch {}; } w.writeAll("\"}") catch {}; } + if (fbs.pos + 64 >= MAX_BODY) { + fbs.pos = saved_pos; + truncated = true; + break; + } } w.writeAll("]}") catch {}; + if (truncated) { + if (fbs.pos >= 1) { + fbs.pos -= 1; + w.writeAll(",\"truncated\":true}") catch {}; + } + } return ok(getBodyBuf()[0..fbs.pos]); } @@ -681,16 +810,19 @@ fn handleDiscoverContext(srv: *Server, tenant_id: []const u8, col_name: []const } w.writeAll("\",\"matching_files\":[") catch {}; for (result.matching_files, 0..) |d, i| { + if (fbs.pos + 128 >= MAX_BODY) break; if (i > 0) w.writeByte(',') catch {}; std.fmt.format(w, "{{\"key\":\"{s}\",\"size\":{d}}}", .{ d.key, d.value.len }) catch {}; } w.writeAll("],\"related_files\":[") catch {}; for (result.related_files, 0..) |d, i| { + if (fbs.pos + 128 >= MAX_BODY) break; if (i > 0) w.writeByte(',') catch {}; std.fmt.format(w, "{{\"key\":\"{s}\"}}", .{d.key}) catch {}; } w.writeAll("],\"test_files\":[") catch {}; for (result.test_files, 0..) |d, i| { + if (fbs.pos + 128 >= MAX_BODY) break; if (i > 0) w.writeByte(',') catch {}; std.fmt.format(w, "{{\"key\":\"{s}\"}}", .{d.key}) catch {}; } @@ -946,13 +1078,28 @@ fn respond(code: u16, status: []const u8, body: []const u8) usize { // ─── mini parsers ──────────────────────────────────────────────────────── fn extractContentLength(raw: []const u8) usize { - // Case-insensitive search for Content-Length header + // Fully case-insensitive search for Content-Length header. const headers = raw[0..@min(raw.len, 2048)]; // only scan headers - const needle = "ontent-length: "; // skip first char for case insensitivity var i: usize = 0; - while (i + needle.len < headers.len) : (i += 1) { - if ((headers[i] == 'C' or headers[i] == 'c') and std.mem.eql(u8, headers[i + 1 .. i + 1 + needle.len], needle)) { - const start = i + 1 + needle.len; + while (i + 16 < headers.len) : (i += 1) { + if ((headers[i] == 'C' or headers[i] == 'c') and + (headers[i + 1] == 'o' or headers[i + 1] == 'O') and + (headers[i + 2] == 'n' or headers[i + 2] == 'N') and + (headers[i + 3] == 't' or headers[i + 3] == 'T') and + (headers[i + 4] == 'e' or headers[i + 4] == 'E') and + (headers[i + 5] == 'n' or headers[i + 5] == 'N') and + (headers[i + 6] == 't' or headers[i + 6] == 'T') and + headers[i + 7] == '-' and + (headers[i + 8] == 'L' or headers[i + 8] == 'l') and + (headers[i + 9] == 'e' or headers[i + 9] == 'E') and + (headers[i + 10] == 'n' or headers[i + 10] == 'N') and + (headers[i + 11] == 'g' or headers[i + 11] == 'G') and + (headers[i + 12] == 't' or headers[i + 12] == 'T') and + (headers[i + 13] == 'h' or headers[i + 13] == 'H') and + headers[i + 14] == ':') + { + var start = i + 15; + while (start < headers.len and (headers[start] == ' ' or headers[start] == '\t')) start += 1; var end = start; while (end < headers.len and headers[end] >= '0' and headers[end] <= '9') : (end += 1) {} if (end > start) { @@ -963,6 +1110,56 @@ fn extractContentLength(raw: []const u8) usize { return 0; } +/// Escape a string for safe JSON interpolation: `"` → `\"`, `\` → `\\`. +/// Writes into the provided stack buffer; returns the escaped slice or the +/// original string unchanged when no escaping is needed. +fn jsonEscape(raw: []const u8, buf: *[1024]u8) []const u8 { + var needs_escape = false; + for (raw) |ch| { + if (ch == '"' or ch == '\\' or ch < 0x20) { needs_escape = true; break; } + } + if (!needs_escape) return raw; + var pos: usize = 0; + for (raw) |ch| { + if (ch == '"' or ch == '\\') { + if (pos + 2 > buf.len) return buf[0..pos]; + buf[pos] = '\\'; + buf[pos + 1] = ch; + pos += 2; + } else if (ch == '\n') { + if (pos + 2 > buf.len) return buf[0..pos]; + buf[pos] = '\\'; + buf[pos + 1] = 'n'; + pos += 2; + } else if (ch == '\r') { + if (pos + 2 > buf.len) return buf[0..pos]; + buf[pos] = '\\'; + buf[pos + 1] = 'r'; + pos += 2; + } else if (ch == '\t') { + if (pos + 2 > buf.len) return buf[0..pos]; + buf[pos] = '\\'; + buf[pos + 1] = 't'; + pos += 2; + } else if (ch < 0x20) { + if (pos + 6 > buf.len) return buf[0..pos]; + const hex = "0123456789abcdef"; + buf[pos] = '\\'; + buf[pos + 1] = 'u'; + buf[pos + 2] = '0'; + buf[pos + 3] = '0'; + buf[pos + 4] = hex[ch >> 4]; + buf[pos + 5] = hex[ch & 0x0f]; + pos += 6; + } else { + if (pos + 1 > buf.len) return buf[0..pos]; + buf[pos] = ch; + pos += 1; + } + } + return buf[0..pos]; +} + fn jsonStr(json: []const u8, key: []const u8) ?[]const u8 { var kbuf: [64]u8 = undefined; const needle = std.fmt.bufPrint(&kbuf, "\"{s}\":", .{key}) catch return null; @@ -972,8 +1169,13 @@ fn jsonStr(json: []const u8, key: []const u8) ?[]const u8 { while (start < json.len and (json[start] == ' ' or json[start] == '\t')) start += 1; if (start >= json.len or json[start] != '"') return null; start += 1; // skip opening quote - const end = std.mem.indexOfScalarPos(u8, json, start, '"') orelse return null; - return json[start..end]; + // Scan for closing quote, skipping escaped quotes + var i = start; + while (i < json.len) : (i += 1) { + if (json[i] == '\\' and i + 1 < json.len) { i += 1; continue; } + if (json[i] == '"') return json[start..i]; + } + return null; } // Extract a JSON value (string, object, array, number, bool, null) for the given key. @@ -996,9 +1198,9 @@ fn jsonValue(json: []const u8, key: []const u8) ?[]const u8 { var i = start + 1; while (i < json.len) : (i += 1) { if (json[i] == '\\' and i + 1 < json.len) { i += 1; continue; } - if (json[i] == '"') break; + if (json[i] == '"') return json[start .. i + 1]; // include both quotes } - return json[start .. i + 1]; // include both quotes + return null; // unterminated string — don't read past buffer } else if (ch == '{' or ch == '[') { // Object or array — find matching close bracket const close: u8 = if (ch == '{') '}' else ']'; @@ -1084,6 +1286,389 @@ fn hexVal(c: u8) ?u4 { return null; } +// ─── WebSocket ─────────────────────────────────────────────────────────── + +/// Read exactly `buf.len` bytes from the stream, looping as needed. +fn wsReadExact(conn: std.net.Server.Connection, buf: []u8) !void { + var filled: usize = 0; + while (filled < buf.len) { + const n = conn.stream.read(buf[filled..]) catch return error.ReadFailed; + if (n == 0) return error.ConnectionClosed; + filled += n; + } +} +/// Case-insensitive header value check. +fn headerContains(raw: []const u8, name: []const u8, value: []const u8) bool { + const headers = raw[0..@min(raw.len, 4096)]; + var i: usize = 0; + while (i + name.len + 2 < headers.len) : (i += 1) { + if (headers[i] == '\n') { + const line_start = i + 1; + if (line_start + name.len + 2 >= headers.len) continue; + var match = true; + for (0..name.len) |j| { + const a = headers[line_start + j]; + const b = name[j]; + const al = if (a >= 'A' and a <= 'Z') a + 32 else a; + const bl = if (b >= 'A' and b <= 'Z') b + 32 else b; + if (al != bl) { match = false; break; } + } + if (!match) continue; + if (headers[line_start + name.len] != ':') continue; + // Found header — check value (case-insensitive) + var vs = line_start + name.len + 1; + while (vs < headers.len and (headers[vs] == ' ' or headers[vs] == '\t')) vs += 1; + const ve = std.mem.indexOfScalarPos(u8, headers, vs, '\r') orelse + std.mem.indexOfScalarPos(u8, headers, vs, '\n') orelse headers.len; + const hval = headers[vs..ve]; + if (hval.len < value.len) continue; + var vmatch = true; + for (0..value.len) |j| { + const a = hval[j]; + const b = value[j]; + const al = if (a >= 'A' and a <= 'Z') a + 32 else a; + const bl = if (b >= 'A' and b <= 'Z') b + 32 else b; + if (al != bl) { vmatch = false; break; } + } + if (vmatch) return true; + } + } + return false; +} + +/// Extract a specific header value from raw HTTP request. +fn headerValue(raw: []const u8, name: []const u8) ?[]const u8 { + const headers = raw[0..@min(raw.len, 4096)]; + var i: usize = 0; + while (i + name.len + 2 < headers.len) : (i += 1) { + if (headers[i] == '\n') { + const ls = i + 1; + if (ls + name.len + 2 >= headers.len) continue; + var match = true; + for (0..name.len) |j| { + const a = headers[ls + j]; + const b = name[j]; + const al = if (a >= 'A' and a <= 'Z') a + 32 else a; + const bl = if (b >= 'A' and b <= 'Z') b + 32 else b; + if (al != bl) { match = false; break; } + } + if (!match) continue; + if (headers[ls + name.len] != ':') continue; + var vs = ls + name.len + 1; + while (vs < headers.len and (headers[vs] == ' ' or headers[vs] == '\t')) vs += 1; + const ve = std.mem.indexOfScalarPos(u8, headers, vs, '\r') orelse + std.mem.indexOfScalarPos(u8, headers, vs, '\n') orelse headers.len; + return headers[vs..ve]; + } + } + return null; +} + +const WS_MAGIC = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"; + +/// Perform WebSocket upgrade handshake, then loop reading WS frames. +/// Each text message is a JSON request dispatched through the normal +/// HTTP handler; the response body is sent back as a WS text frame. +fn handleWebSocket(srv: *Server, conn: std.net.Server.Connection, initial: []const u8) !void { + // Extract Sec-WebSocket-Key + const ws_key = headerValue(initial, "Sec-WebSocket-Key") orelse return error.MissingKey; + + // Compute accept hash: SHA1(key ++ magic), base64-encoded + var sha = std.crypto.hash.Sha1.init(.{}); + sha.update(ws_key); + sha.update(WS_MAGIC); + const digest = sha.finalResult(); + var accept_buf: [28]u8 = undefined; + const accept = std.base64.standard.Encoder.encode(&accept_buf, &digest); + + // Send upgrade response + var resp_buf: [256]u8 = undefined; + const resp_len = std.fmt.bufPrint(&resp_buf, + "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {s}\r\n\r\n", + .{accept}) catch return error.FormatFailed; + conn.stream.writeAll(resp_len) catch return error.WriteFailed; + + // WebSocket frame loop — heap-allocate frame buffer (too large for stack). + const frame_buf = std.heap.page_allocator.alloc(u8, MAX_BULK) catch return error.OutOfMemory; + defer std.heap.page_allocator.free(frame_buf); + while (true) { + // Read frame header (2 bytes minimum) + var hdr: [14]u8 = undefined; + wsReadExact(conn, hdr[0..2]) catch return; + const fin = hdr[0] & 0x80 != 0; + const opcode = hdr[0] & 0x0F; + const masked = hdr[1] & 0x80 != 0; + var payload_len: u64 = hdr[1] & 0x7F; + + if (opcode == 0x8) return; // close frame + if (opcode == 0x9) { // ping → pong + // Resolve the full payload length (may be extended). + var ping_len: u64 = payload_len; + if (ping_len == 126) { + wsReadExact(conn, hdr[2..4]) catch return; + ping_len = std.mem.readInt(u16, hdr[2..4], .big); + } else if (ping_len == 127) { + wsReadExact(conn, hdr[2..10]) catch return; + ping_len = std.mem.readInt(u64, hdr[2..10], .big); + } + // Read and discard mask key if present. + var ping_mask: [4]u8 = .{0, 0, 0, 0}; + if (masked) wsReadExact(conn, &ping_mask) catch return; + // Read the ping payload so the stream stays synchronised. + const pl: usize = @intCast(ping_len); + if (pl > 0 and pl <= frame_buf.len) { + wsReadExact(conn, frame_buf[0..pl]) catch return; + // Unmask payload before echoing it back. + if (masked) { + for (frame_buf[0..pl], 0..) |*b, j| b.* ^= ping_mask[j % 4]; + } + } + // Send pong with the same payload. + var pong_hdr: [2]u8 = .{ 0x8A, @intCast(ping_len & 0x7F) }; + conn.stream.writeAll(&pong_hdr) catch return; + if (pl > 0 and pl <= frame_buf.len) { + conn.stream.writeAll(frame_buf[0..pl]) catch return; + } + continue; + } + + // Extended payload length + if (payload_len == 126) { + wsReadExact(conn, hdr[2..4]) catch return; + payload_len = std.mem.readInt(u16, hdr[2..4], .big); + } else if (payload_len == 127) { + wsReadExact(conn, hdr[2..10]) catch return; + payload_len = std.mem.readInt(u64, hdr[2..10], .big); + } + + if (payload_len > MAX_BULK) { + wsWriteClose(conn, 1009); // message too big + return; + } + + // Read mask key (4 bytes if masked) + var mask: [4]u8 = .{0, 0, 0, 0}; + if (masked) wsReadExact(conn, &mask) catch return; + + // Read payload + const plen: usize = @intCast(payload_len); + const payload = frame_buf[0..plen]; + wsReadExact(conn, payload) catch return; + + // Unmask + if (masked) { + for (payload, 0..) |*b, j| b.* ^= mask[j % 4]; + } + + _ = fin; // TODO: handle fragmented messages + + if (opcode == 0x1) { // text frame + _ = srv.req_count.fetchAdd(1, .monotonic); + + // STREAM mode: fire-and-forget bulk insert (no per-frame response). + if (payload.len >= 7 and std.mem.startsWith(u8, payload, "STREAM ")) { + handleWsStream(srv, conn, payload, frame_buf) catch return; + return; // stream handler owns the connection + } + + // Normal WS dispatch (request-response). + const resp_body = wsDispatch(srv, payload); + wsWriteText(conn, resp_body) catch return; + } + } +} + +/// Extract a `name=value` parameter from a space-separated string. +fn extractParam(msg: []const u8, name: []const u8) ?[]const u8 { + const pos = std.mem.indexOf(u8, msg, name) orelse return null; + const start = pos + name.len; + const end = std.mem.indexOfScalarPos(u8, msg, start, ' ') orelse msg.len; + if (start >= end) return null; + return msg[start..end]; +} + +/// WebSocket streaming bulk insert — fire-and-forget NDJSON frames. +/// Init frame: "STREAM /db/:col/bulk [tenant=X]" +/// Data frames: NDJSON text (no per-frame response sent). +/// Close frame: triggers final summary response. +fn handleWsStream(srv: *Server, conn: std.net.Server.Connection, init_msg: []const u8, frame_buf: []u8) !void { + // Parse collection and tenant from init message. + const after_prefix = init_msg[7..]; // skip "STREAM " + const path_end = std.mem.indexOfScalar(u8, after_prefix, ' ') orelse after_prefix.len; + const path = after_prefix[0..path_end]; + const params = if (path_end < after_prefix.len) after_prefix[path_end + 1 ..] else ""; + + // Extract collection name from /db/:col/bulk (or /db/:col). + if (!std.mem.startsWith(u8, path, "/db/")) { + wsWriteText(conn, "{\"error\":\"invalid stream path\"}") catch {}; + wsWriteClose(conn, 1008); + return; + } + const after_db = path[4..]; + const col_end = std.mem.indexOfScalar(u8, after_db, '/') orelse after_db.len; + const col_name = after_db[0..col_end]; + if (col_name.len == 0) { + wsWriteText(conn, "{\"error\":\"missing collection name\"}") catch {}; + wsWriteClose(conn, 1008); + return; + } + + const tenant_id = extractParam(params, "tenant=") orelse collection.DEFAULT_TENANT; + + const col = srv.db.collectionForTenant(tenant_id, col_name) catch { + wsWriteText(conn, "{\"error\":\"open collection failed\"}") catch {}; + wsWriteClose(conn, 1011); + return; + }; + + // ACK — client waits for this before streaming. + wsWriteText(conn, "{\"status\":\"streaming\"}") catch return error.WriteFailed; + + var total_inserted: u64 = 0; + var total_errors: u64 = 0; + var total_bytes: u64 = 0; + var frame_count: u64 = 0; + const start_ns = std.time.nanoTimestamp(); + + // Ensure billing is recorded even on connection drop. + defer srv.recordQueryCost(tenant_id, "ws_stream", @intCast(total_inserted), @intCast(total_bytes), start_ns); + + // Frame loop — same structure as handleWebSocket but no per-frame response. + while (true) { + var hdr: [14]u8 = undefined; + wsReadExact(conn, hdr[0..2]) catch break; + const opcode = hdr[0] & 0x0F; + const masked = hdr[1] & 0x80 != 0; + var payload_len: u64 = hdr[1] & 0x7F; + + if (opcode == 0x8) break; // close — end of stream + if (opcode == 0x9) { // ping → pong + var ping_len: u64 = payload_len; + if (ping_len == 126) { + wsReadExact(conn, hdr[2..4]) catch break; + ping_len = std.mem.readInt(u16, hdr[2..4], .big); + } else if (ping_len == 127) { + wsReadExact(conn, hdr[2..10]) catch break; + ping_len = std.mem.readInt(u64, hdr[2..10], .big); + } + var ping_mask: [4]u8 = .{ 0, 0, 0, 0 }; + if (masked) wsReadExact(conn, &ping_mask) catch break; + const pl: usize = @intCast(ping_len); + if (pl > 0 and pl <= frame_buf.len) { + wsReadExact(conn, frame_buf[0..pl]) catch break; + if (masked) { + for (frame_buf[0..pl], 0..) |*b, j| b.* ^= ping_mask[j % 4]; + } + } + var pong_hdr: [2]u8 = .{ 0x8A, @intCast(ping_len & 0x7F) }; + conn.stream.writeAll(&pong_hdr) catch break; + if (pl > 0 and pl <= frame_buf.len) { + conn.stream.writeAll(frame_buf[0..pl]) catch break; + } + continue; + } + + // Extended payload length. + if (payload_len == 126) { + wsReadExact(conn, hdr[2..4]) catch break; + payload_len = std.mem.readInt(u16, hdr[2..4], .big); + } else if (payload_len == 127) { + wsReadExact(conn, hdr[2..10]) catch break; + payload_len = std.mem.readInt(u64, hdr[2..10], .big); + } + + if (payload_len > MAX_BULK) { + wsWriteClose(conn, 1009); + break; + } + + var mask: [4]u8 = .{ 0, 0, 0, 0 }; + if (masked) wsReadExact(conn, &mask) catch break; + + const plen: usize = @intCast(payload_len); + const payload = frame_buf[0..plen]; + wsReadExact(conn, payload) catch break; + + if (masked) { + for (payload, 0..) |*b, j| b.* ^= mask[j % 4]; + } + + // Process NDJSON frame (text or binary). + if (opcode == 0x1 or opcode == 0x2) { + srv.db.recordTenantOperation(tenant_id) catch break; + const r = parseAndInsertNdjson(col, payload); + total_inserted += r.inserted; + total_errors += r.errors; + total_bytes += r.total_bytes; + frame_count += 1; + } + } + + // Send final summary before closing. + var summary_buf: [256]u8 = undefined; + const summary = std.fmt.bufPrint(&summary_buf, + "{{\"inserted\":{d},\"errors\":{d},\"frames\":{d}}}", + .{ total_inserted, total_errors, frame_count }) catch "{\"error\":\"format\"}"; + wsWriteText(conn, summary) catch {}; + wsWriteClose(conn, 1000); +} + + +/// Dispatch a WebSocket text message. Supports two formats: +/// 1. Raw HTTP: "METHOD /path\n{body}" — reuses dispatch() +/// 2. JSON ops: {"op":"insert","col":"c","key":"k","value":{}} (future) +fn wsDispatch(srv: *Server, msg: []const u8) []const u8 { + // Wrap as a minimal HTTP request so we can reuse dispatch(). + // Find first \n to split method+path from body. + const nl = std.mem.indexOfScalar(u8, msg, '\n') orelse msg.len; + const req_line = msg[0..nl]; + const body = if (nl < msg.len) msg[nl + 1 ..] else ""; + + // Build HTTP/1.1 request with Content-Length. + // Heap-allocate since WS messages can be up to 16MB. + const needed = req_line.len + 64 + body.len; + const http_buf = std.heap.page_allocator.alloc(u8, needed) catch + return "{\"error\":\"request too large\"}"; + defer std.heap.page_allocator.free(http_buf); + const http_len = std.fmt.bufPrint(http_buf, + "{s} HTTP/1.1\r\nContent-Length: {d}\r\n\r\n{s}", + .{ req_line, body.len, body }) catch return "{\"error\":\"request too large\"}"; + + const resp_len = dispatch(srv, http_len, std.heap.page_allocator); + const resp = getRespBuf()[0..resp_len]; + + // Strip HTTP headers from response, return just the body + if (std.mem.indexOf(u8, resp, "\r\n\r\n")) |p| return resp[p + 4 ..]; + return resp; +} + +fn wsWriteText(conn: std.net.Server.Connection, payload: []const u8) !void { + var hdr: [10]u8 = undefined; + hdr[0] = 0x81; // FIN + text + var hdr_len: usize = 2; + if (payload.len < 126) { + hdr[1] = @intCast(payload.len); + } else if (payload.len < 65536) { + hdr[1] = 126; + std.mem.writeInt(u16, hdr[2..4], @intCast(payload.len), .big); + hdr_len = 4; + } else { + hdr[1] = 127; + std.mem.writeInt(u64, hdr[2..10], @intCast(payload.len), .big); + hdr_len = 10; + } + try conn.stream.writeAll(hdr[0..hdr_len]); + try conn.stream.writeAll(payload); +} + +fn wsWriteClose(conn: std.net.Server.Connection, code: u16) void { + var frame: [4]u8 = undefined; + frame[0] = 0x88; // FIN + close + frame[1] = 2; // payload = 2 bytes (status code) + std.mem.writeInt(u16, frame[2..4], code, .big); + conn.stream.writeAll(&frame) catch {}; +} + test "parse as_of accepts seconds and milliseconds" { try std.testing.expectEqual(@as(?i64, 1_700_000_000_000), parseAsOfTimestamp("1700000000")); try std.testing.expectEqual(@as(?i64, 1_700_000_000_123), parseAsOfTimestamp("1700000000123")); diff --git a/src/storage/epoch.zig b/src/storage/epoch.zig index d265f8d..97f6795 100644 --- a/src/storage/epoch.zig +++ b/src/storage/epoch.zig @@ -68,6 +68,14 @@ pub const EpochManager = struct { self.timeline_mu.lock(); defer self.timeline_mu.unlock(); self.timeline.append(self.allocator, .{ .epoch = epoch, .ts_ms = ts_ms }) catch {}; + // Prune old timeline entries to prevent unbounded memory growth. + // Keep at most 100K entries (~1.6 MiB); drop the oldest half when full. + const TIMELINE_CAP = 100_000; + if (self.timeline.items.len > TIMELINE_CAP) { + const drop = self.timeline.items.len / 2; + std.mem.copyForwards(EpochPoint, self.timeline.items[0..self.timeline.items.len - drop], self.timeline.items[drop..]); + self.timeline.items.len -= drop; + } return epoch; } @@ -96,7 +104,8 @@ pub const EpochManager = struct { /// Pin the current epoch. Returns (reader_slot, snapshot_ts). /// The caller must call unpin(slot) when done. - pub fn pin(self: *EpochManager) struct { slot: usize, ts: u64 } { + /// Returns error.ReadersExhausted when all slots are occupied. + pub fn pin(self: *EpochManager) error{ReadersExhausted}!struct { slot: usize, ts: u64 } { const ts = self.global_ts.load(.acquire); // Find a free slot for (self.slots, 0..) |*s, i| { @@ -106,8 +115,8 @@ pub const EpochManager = struct { return .{ .slot = i, .ts = ts }; } } - // Fallback: return slot 0 (reader count exceeded MAX_READERS — very rare) - return .{ .slot = 0, .ts = ts }; + // All slots occupied — cannot safely pin + return error.ReadersExhausted; } pub fn unpin(self: *EpochManager, slot: usize) void { diff --git a/src/storage/mmap.zig b/src/storage/mmap.zig index d10e842..f6ab543 100644 --- a/src/storage/mmap.zig +++ b/src/storage/mmap.zig @@ -4,22 +4,41 @@ //! zero-copy; writes go through the page cache and are made durable //! by msync(MS_SYNC) at checkpoint time. //! -//! Growth strategy: grow in GROW_CHUNK increments (256 MiB) to amortise -//! ftruncate + remap syscalls. On macOS (no mremap) we munmap and remap. +//! Stability guarantee: we pre-reserve a large virtual address range +//! (MAX_VA_SIZE) at open() time. grow() extends the backing file and +//! maps new pages into the *same* VA range with MAP_FIXED, so the base +//! pointer never moves. at() / slice() are therefore lock-free — no +//! pointer can dangle because the mapping base is stable for the +//! lifetime of the MmapFile. const std = @import("std"); const posix = std.posix; -pub const GROW_CHUNK: usize = 256 * 1024 * 1024; // 256 MiB +pub const GROW_CHUNK: usize = 64 * 1024 * 1024; // 64 MiB pub const PAGE_SIZE: usize = 4096; +/// OS minimum page alignment — on Apple Silicon this is 16384, on x86-64 +/// it is 4096. All mmap pointers must satisfy this alignment. +const OS_PAGE_ALIGN: usize = std.heap.page_size_min; + +/// 4 GiB virtual reservation. Costs zero physical memory — pages are +/// only materialised when backed by ftruncate + MAP_FIXED. +pub const MAX_VA_SIZE: usize = 1024 * 1024 * 1024; // 1 GiB + pub const MmapFile = struct { fd: posix.fd_t, - ptr: [*]align(PAGE_SIZE) u8, - len: usize, // logical data length - capacity: usize, // mapped (file) length; >= len, multiple of PAGE_SIZE - /// Protects ptr/capacity against concurrent grow+read. - /// Writers (grow) take exclusive; readers (at/slice) take shared. - rw_lock: std.Thread.RwLock, + /// Stable base pointer — lives for the entire VA reservation lifetime. + /// Never changes after open(). + ptr: [*]align(OS_PAGE_ALIGN) u8, + /// Logical data length. Written by grow(), read by appendRecords() and + /// external callers. Uses atomic ops to avoid a data race with + /// concurrent readers (no lock needed thanks to the stable VA range). + len: std.atomic.Value(usize), + /// How many bytes of the VA range are currently file-backed and mapped + /// PROT_READ|PROT_WRITE. Atomic so concurrent at()/slice() callers + /// can read without holding grow_mu. + capacity: std.atomic.Value(usize), + /// Serialises concurrent grow() calls. NOT needed for readers. + grow_mu: std.Thread.Mutex, // ── Open / Close ────────────────────────────────────────────────────────── @@ -32,102 +51,133 @@ pub const MmapFile = struct { const existing: usize = @intCast(stat.size); const capacity = alignUp( @max(existing, @max(initial_size, GROW_CHUNK)), - PAGE_SIZE, + OS_PAGE_ALIGN, + ); + + // 1. Reserve the full virtual address range (PROT_NONE, anonymous). + // No physical memory is consumed — this just carves out VA space. + const reserve = try posix.mmap( + null, MAX_VA_SIZE, + posix.PROT.NONE, + .{ .TYPE = .PRIVATE, .ANONYMOUS = true }, + -1, 0, ); + const base_ptr: [*]align(OS_PAGE_ALIGN) u8 = @alignCast(reserve.ptr); + // 2. Extend the backing file to `capacity` if needed. if (@as(usize, @intCast(stat.size)) < capacity) try posix.ftruncate(fd, @intCast(capacity)); - const ptr = try posix.mmap( - null, capacity, + // 3. Map the file content over the start of the reserved range + // with MAP_FIXED so the kernel places it exactly at base_ptr. + _ = try posix.mmap( + base_ptr, capacity, posix.PROT.READ | posix.PROT.WRITE, - .{ .TYPE = .SHARED }, + .{ .TYPE = .SHARED, .FIXED = true }, fd, 0, ); return .{ .fd = fd, - .ptr = ptr.ptr, - .len = existing, - .capacity = capacity, - .rw_lock = .{}, + .ptr = base_ptr, + .len = std.atomic.Value(usize).init(existing), + .capacity = std.atomic.Value(usize).init(capacity), + .grow_mu = .{}, }; } pub fn close(self: *MmapFile) void { - posix.munmap(@alignCast(self.ptr[0..self.capacity])); + // Unmap the entire VA reservation in one call. + posix.munmap(@alignCast(self.ptr[0..MAX_VA_SIZE])); posix.close(self.fd); } // ── Sync / Checkpoint ───────────────────────────────────────────────────── - pub fn syncAsync(self: MmapFile) void { - posix.msync(@alignCast(self.ptr[0..self.capacity]), posix.MSF.ASYNC) catch {}; + pub fn syncAsync(self: *MmapFile) void { + const cap = self.capacity.load(.acquire); + posix.msync(@alignCast(self.ptr[0..cap]), posix.MSF.ASYNC) catch {}; } - pub fn syncSync(self: MmapFile) !void { - try posix.msync(@alignCast(self.ptr[0..self.capacity]), posix.MSF.SYNC); + pub fn syncSync(self: *MmapFile) !void { + const cap = self.capacity.load(.acquire); + try posix.msync(@alignCast(self.ptr[0..cap]), posix.MSF.SYNC); } // ── Grow ────────────────────────────────────────────────────────────────── - /// Extend the logical length. Remaps if needed (macOS: munmap + mmap). - /// Takes an exclusive lock so concurrent readers cannot hold stale pointers. + /// Extend the logical length. If the file-backed region is too small, + /// extends the file and maps the new pages into the pre-reserved VA + /// range with MAP_FIXED. The base pointer never changes. pub fn grow(self: *MmapFile, needed_len: usize) !void { - if (needed_len <= self.capacity) { - self.len = needed_len; + if (needed_len <= self.capacity.load(.acquire)) { + // Capacity is sufficient — just advance the logical length. + self.len.store(needed_len, .release); return; } - // Must remap — take exclusive lock to block readers during munmap/mmap. - self.rw_lock.lock(); - defer self.rw_lock.unlock(); + + // Must extend capacity — serialise with other growers. + self.grow_mu.lock(); + defer self.grow_mu.unlock(); + // Re-check after acquiring lock (another thread may have grown). - if (needed_len <= self.capacity) { - self.len = needed_len; + const cur_cap = self.capacity.load(.acquire); + if (needed_len <= cur_cap) { + self.len.store(needed_len, .release); return; } - const new_cap = alignUp(needed_len + GROW_CHUNK, PAGE_SIZE); - // Extend file + + const old_cap = cur_cap; + const new_cap = alignUp(needed_len + GROW_CHUNK, OS_PAGE_ALIGN); + + std.debug.assert(new_cap <= MAX_VA_SIZE); + + // Extend the backing file. try posix.ftruncate(self.fd, @intCast(new_cap)); - // Remap (macOS has no mremap; unmap then remap) - posix.munmap(@alignCast(self.ptr[0..self.capacity])); - const ptr = try posix.mmap( - null, new_cap, + + // Map the NEW portion of the file into the pre-reserved VA range. + // MAP_FIXED overwrites the PROT_NONE reservation for this sub-range. + const delta = new_cap - old_cap; + _ = try posix.mmap( + @alignCast(self.ptr + old_cap), delta, posix.PROT.READ | posix.PROT.WRITE, - .{ .TYPE = .SHARED }, - self.fd, 0, + .{ .TYPE = .SHARED, .FIXED = true }, + self.fd, @intCast(old_cap), ); - self.ptr = ptr.ptr; - self.capacity = new_cap; - self.len = needed_len; + + self.capacity.store(new_cap, .release); + self.len.store(needed_len, .release); } // ── Typed access ────────────────────────────────────────────────────────── /// Return a pointer to the record at byte offset `off`. - /// Takes a shared lock so grow() cannot remap underneath us. + /// Lock-free: the base pointer is stable for the lifetime of the mapping. pub fn at(self: *MmapFile, comptime T: type, off: usize) *T { - self.rw_lock.lockShared(); - defer self.rw_lock.unlockShared(); + std.debug.assert(off + @sizeOf(T) <= self.capacity.load(.acquire)); return @alignCast(@ptrCast(&self.ptr[off])); } /// Return a slice of T starting at byte offset `off`, `count` elements. - /// Takes a shared lock so grow() cannot remap underneath us. + /// Lock-free: the base pointer is stable for the lifetime of the mapping. pub fn slice(self: *MmapFile, comptime T: type, off: usize, count: usize) []T { - self.rw_lock.lockShared(); - defer self.rw_lock.unlockShared(); + std.debug.assert(off + count * @sizeOf(T) <= self.capacity.load(.acquire)); return @as([*]T, @alignCast(@ptrCast(&self.ptr[off])))[0..count]; } /// Append `count` new T-records after existing data; returns their base offset. pub fn appendRecords(self: *MmapFile, comptime T: type, count: usize) !usize { - const off = alignUp(self.len, @alignOf(T)); + const off = alignUp(self.len.load(.acquire), @alignOf(T)); const need = off + @sizeOf(T) * count; try self.grow(need); return off; } + /// Read the logical data length (atomic, safe from any thread). + pub fn dataLen(self: *const MmapFile) usize { + return self.len.load(.acquire); + } + // ── Helpers ─────────────────────────────────────────────────────────────── inline fn alignUp(v: usize, a: usize) usize { diff --git a/src/storage/parallel_wal.zig b/src/storage/parallel_wal.zig index 1ef8b3e..00f96f8 100644 --- a/src/storage/parallel_wal.zig +++ b/src/storage/parallel_wal.zig @@ -20,6 +20,7 @@ pub const WALSegment = struct { fd: std.fs.File, buf: []align(4096) u8, pos: std.atomic.Value(u32), + completed: std.atomic.Value(u32), // tracks how far writes have actually finished flushed_pos: u32, pub const BUF_SIZE: u32 = 65536; @@ -39,6 +40,7 @@ pub const WALSegment = struct { .fd = fd, .buf = buf, .pos = std.atomic.Value(u32).init(0), + .completed = std.atomic.Value(u32).init(0), .flushed_pos = 0, }; } @@ -56,8 +58,9 @@ pub const WALSegment = struct { // Reserve space atomically — lock-free. const offset = self.pos.fetchAdd(total, .monotonic); if (offset + total > BUF_SIZE) { - // Roll back so other threads see accurate remaining space. - _ = self.pos.fetchSub(total, .monotonic); + // Don't roll back — the segment is full. Other threads that already + // reserved valid offsets can still complete. A rollback via fetchSub + // races with concurrent fetchAdd and can corrupt offsets. return error.SegmentFull; } @@ -70,20 +73,25 @@ pub const WALSegment = struct { // Write payload. @memcpy(self.buf[offset + ENTRY_HEADER_SIZE .. offset + total], data); + // Signal completion: spin until it's our turn, then advance `completed`. + // This ensures the flusher sees a contiguous range of finished writes. + while (self.completed.cmpxchgWeak(offset, offset + total, .release, .monotonic)) |_| { + std.atomic.spinLoopHint(); + } + return offset; } /// Flush dirty bytes to the underlying file. Called by the group-commit /// flusher — NOT per-transaction. pub fn flush(self: *WALSegment) !void { - const current = self.pos.load(.acquire); - if (current <= self.flushed_pos) return; + const safe = self.completed.load(.acquire); + if (safe <= self.flushed_pos) return; - const dirty = self.buf[self.flushed_pos..current]; - const written = try self.fd.write(dirty); - if (written != dirty.len) return error.ShortWrite; + const dirty = self.buf[self.flushed_pos..safe]; + try self.fd.writeAll(dirty); - self.flushed_pos = current; + self.flushed_pos = safe; } /// Fsync the segment file to durable storage. diff --git a/src/storage/wal.zig b/src/storage/wal.zig index f5ee3b1..47e82a4 100644 --- a/src/storage/wal.zig +++ b/src/storage/wal.zig @@ -113,11 +113,15 @@ pub const WAL = struct { cond: std.Thread.Condition, synced_lsn: u64, flushing: bool, + io_err_flag: bool, // Background flusher flush_thread: ?std.Thread, flush_running: std.atomic.Value(bool), + /// Backpressure: block writers if pending buffer exceeds this. + const MAX_WRITE_BUF: usize = 8 * 1024 * 1024; // 8 MiB + // ── Lifecycle ───────────────────────────────────────────────────────────── pub fn open(path: [:0]const u8, allocator: std.mem.Allocator) !WAL { @@ -134,6 +138,7 @@ pub const WAL = struct { .cond = .{}, .synced_lsn = 0, .flushing = false, + .io_err_flag = false, .flush_thread = null, .flush_running = std.atomic.Value(bool).init(false), }; @@ -175,9 +180,30 @@ pub const WAL = struct { self.write_buf = .empty; self.mu.unlock(); - self.file.writeAll(to_write.items) catch {}; + // Write entries to file. If writeAll fails, re-queue the buffer so + // entries are retried on the next flush cycle instead of being lost. + self.file.writeAll(to_write.items) catch |e| { + self.mu.lock(); + // Prepend failed entries before any new writes that arrived. + to_write.appendSlice(self.allocator, self.write_buf.items) catch {}; + self.write_buf.deinit(self.allocator); + self.write_buf = to_write; + self.io_err_flag = true; + self.flushing = false; + self.cond.broadcast(); + self.mu.unlock(); + std.log.err("WAL flushPending write error: {}", .{e}); + return; + }; to_write.deinit(self.allocator); - self.file.sync() catch {}; + + // Data is in the file. fsync for durability — if sync fails the data + // is still in the kernel page cache (and the file), so advance + // synced_lsn either way to avoid permanently stalling the WAL. + self.file.sync() catch |e| { + self.io_err_flag = true; + std.log.err("WAL flushPending sync error (data written, not fsynced): {}", .{e}); + }; self.mu.lock(); self.synced_lsn = target; @@ -201,6 +227,7 @@ pub const WAL = struct { /// Encode an entry into the shared write buffer and return its LSN. /// Thread-safe. Does NOT guarantee durability — call commit(lsn) for that. + /// Blocks if write buffer exceeds MAX_WRITE_BUF to apply backpressure. pub fn write( self: *WAL, txn_id: u64, @@ -226,6 +253,12 @@ pub const WAL = struct { hdr.crc32 = entryChecksum(hdr_bytes, payload); self.mu.lock(); + // Backpressure: wait until flusher drains the buffer below threshold. + // Use condition variable instead of yield-loop so the flusher can wake us. + while (self.write_buf.items.len >= MAX_WRITE_BUF) { + // Release lock and wait for flusher to signal. + self.cond.wait(&self.mu); + } defer self.mu.unlock(); try self.write_buf.appendSlice(self.allocator, std.mem.asBytes(&hdr)); try self.write_buf.appendSlice(self.allocator, payload); @@ -233,6 +266,57 @@ pub const WAL = struct { return lsn; } + /// Batch write: append N entries under a single lock acquisition. + /// All entries share the same op/txn_id/db_tag/flags. + /// More efficient than calling write() in a loop for bulk inserts. + pub fn writeBatch( + self: *WAL, + txn_id: u64, + op: OpCode, + db_tag: u8, + flags: u16, + payloads: []const []const u8, + ) !void { + if (payloads.len == 0) return; + + // Pre-compute total buffer size needed. + var total_size: usize = 0; + for (payloads) |p| { + total_size += HEADER_SIZE + p.len + paddingTo8(HEADER_SIZE + p.len); + } + + self.mu.lock(); + while (self.write_buf.items.len >= MAX_WRITE_BUF) { + self.cond.wait(&self.mu); + } + defer self.mu.unlock(); + + // Reserve all space at once — avoids per-entry ArrayList growth checks. + try self.write_buf.ensureUnusedCapacity(self.allocator, total_size); + + for (payloads) |payload| { + const lsn = self.next_lsn.fetchAdd(1, .monotonic); + const pad = paddingTo8(HEADER_SIZE + payload.len); + + var hdr = EntryHeader{ + .lsn = lsn, + .length = @intCast(payload.len), + .crc32 = 0, + .op_code = @intFromEnum(op), + .db_tag = db_tag, + .flags = flags, + .txn_id = txn_id, + .reserved = 0, + }; + const hdr_bytes = std.mem.asBytes(&hdr); + hdr.crc32 = entryChecksum(hdr_bytes, payload); + + self.write_buf.appendSliceAssumeCapacity(std.mem.asBytes(&hdr)); + self.write_buf.appendSliceAssumeCapacity(payload); + if (pad > 0) self.write_buf.appendNTimesAssumeCapacity(0, pad); + } + } + /// Mark a transaction committed. Returns only after the entry is durable. /// Implements group commit: the first caller flushes for everyone. pub fn commit(self: *WAL, txn_id: u64, db_tag: u8) !void { @@ -266,45 +350,104 @@ pub const WAL = struct { self.mu.unlock(); // ── I/O outside lock ────────────────────────────────────────────────── - var io_err: ?anyerror = null; - self.file.writeAll(to_write.items) catch |e| { io_err = e; }; + // If writeAll fails, re-queue the buffer so entries survive for retry. + self.file.writeAll(to_write.items) catch |e| { + self.mu.lock(); + to_write.appendSlice(self.allocator, self.write_buf.items) catch {}; + self.write_buf.deinit(self.allocator); + self.write_buf = to_write; + self.flushing = false; + self.cond.broadcast(); + self.mu.unlock(); + return e; + }; to_write.deinit(self.allocator); - if (io_err == null) self.file.sync() catch |e| { io_err = e; }; + + // Data written — fsync for durability. If sync fails, still advance + // synced_lsn (data is in page cache) but return the error to caller. + var sync_err: ?anyerror = null; + self.file.sync() catch |e| { sync_err = e; }; // ───────────────────────────────────────────────────────────────────── self.mu.lock(); - if (io_err == null) self.synced_lsn = target; + self.synced_lsn = target; self.flushing = false; self.cond.broadcast(); self.mu.unlock(); - if (io_err) |e| return e; + if (sync_err) |e| return e; } // ── Checkpoint ──────────────────────────────────────────────────────────── - /// Write a checkpoint barrier entry and update checkpoint_lsn. + /// Write a checkpoint barrier entry, flush, then truncate the WAL file. + /// After checkpoint all data is durable in page files, so the WAL can be + /// safely cleared to reclaim disk space. pub fn checkpoint(self: *WAL, db_tag: u8) !void { var p: [8]u8 = undefined; std.mem.writeInt(u64, &p, self.checkpoint_lsn, .little); const lsn = try self.write(0, .checkpoint, db_tag, FLAG_COMMIT, &p); self.checkpoint_lsn = lsn; - // Force flush + + // Swap write buffer under lock, then release so writers can continue + // appending to the fresh buffer while we flush + sync. self.mu.lock(); self.flushing = true; var to_write = self.write_buf; self.write_buf = .empty; + self.cond.broadcast(); // wake writers waiting on full buffer self.mu.unlock(); - self.file.writeAll(to_write.items) catch {}; + + // ── I/O without lock (writers append to new write_buf concurrently) ── + var io_err: ?anyerror = null; + self.file.writeAll(to_write.items) catch |e| { + io_err = e; + }; to_write.deinit(self.allocator); - try self.file.sync(); + + if (io_err) |e| { + self.mu.lock(); + self.io_err_flag = true; + self.flushing = false; + self.cond.broadcast(); + self.mu.unlock(); + std.log.err("WAL checkpoint write error: {}", .{e}); + return e; + } + + self.file.sync() catch |e| { + self.mu.lock(); + self.io_err_flag = true; + self.flushing = false; + self.cond.broadcast(); + self.mu.unlock(); + std.log.err("WAL checkpoint sync error: {}", .{e}); + return e; + }; + + // Reacquire lock for truncation + state update. self.mu.lock(); + defer self.mu.unlock(); + + self.file.seekTo(0) catch |e| { + self.io_err_flag = true; + self.flushing = false; + self.cond.broadcast(); + std.log.err("WAL checkpoint seekTo error: {}", .{e}); + return e; + }; + self.file.setEndPos(0) catch |e| { + self.io_err_flag = true; + self.flushing = false; + self.cond.broadcast(); + std.log.err("WAL checkpoint setEndPos error: {}", .{e}); + return e; + }; + self.synced_lsn = lsn; self.flushing = false; self.cond.broadcast(); - self.mu.unlock(); } - // ── Recovery ────────────────────────────────────────────────────────────── /// Replay WAL from the beginning. apply_fn is called for every committed diff --git a/src/tdb.zig b/src/tdb.zig index f336637..1e2b44c 100644 --- a/src/tdb.zig +++ b/src/tdb.zig @@ -101,7 +101,7 @@ pub fn main() !void { std.debug.print("Usage: tdb word \n", .{}); return; } - try cmdWord(db, col_name, cmd_args[0]); + try cmdWord(db, col_name, cmd_args[0], alloc); } else if (std.mem.eql(u8, command, "insert")) { if (cmd_argc < 2) { std.debug.print("Usage: tdb insert \n", .{}); @@ -214,9 +214,10 @@ fn cmdSearch(db: *Database, col_name: []const u8, query: []const u8, alloc: std. } } -fn cmdWord(db: *Database, col_name: []const u8, word: []const u8) !void { +fn cmdWord(db: *Database, col_name: []const u8, word: []const u8, alloc: std.mem.Allocator) !void { const col = try db.collection(col_name); - const hits = col.searchWord(word); + const hits = try col.searchWord(word, alloc); + defer if (hits.len > 0) alloc.free(hits); if (hits.len == 0) { std.debug.print(" No hits for \"{s}\"\n", .{word}); return; diff --git a/src/trigram.zig b/src/trigram.zig index d5f9c20..9733e93 100644 --- a/src/trigram.zig +++ b/src/trigram.zig @@ -24,6 +24,11 @@ pub const TrigramIndex = struct { doc_count: u64, trigram_count: u64, + /// Maximum posting list length per trigram. Trigrams appearing in more docs + /// than this are too common to be useful for search — stop tracking them to + /// bound memory. 10K docs × 8 bytes = 80KB per trigram, manageable. + const MAX_POSTINGS_PER_TRIGRAM: usize = 10_000; + pub fn init(alloc: std.mem.Allocator) TrigramIndex { return .{ .postings = std.AutoHashMap(u24, std.ArrayList(u64)).init(alloc), @@ -53,6 +58,8 @@ pub const TrigramIndex = struct { self.trigram_count += 1; } const list = gop.value_ptr; + // Skip overly common trigrams to bound memory. + if (list.items.len >= MAX_POSTINGS_PER_TRIGRAM) continue; if (list.items.len == 0 or list.items[list.items.len - 1] != doc_id) { try list.append(self.alloc, doc_id); } @@ -60,16 +67,21 @@ pub const TrigramIndex = struct { self.doc_count += 1; } - /// Remove a document from all posting lists (lazy tombstone). + /// Remove a document from all posting lists (compact in-place). pub fn removeDoc(self: *TrigramIndex, doc_id: u64, value: []const u8) void { if (value.len < 3) return; var i: usize = 0; while (i + 2 < value.len) : (i += 1) { const key = trigramKey(value[i], value[i + 1], value[i + 2]); if (self.postings.getPtr(key)) |list| { - for (list.items) |*id| { - if (id.* == doc_id) id.* = 0; + var write: usize = 0; + for (list.items) |id| { + if (id != doc_id) { + list.items[write] = id; + write += 1; + } } + list.items.len = write; } } if (self.doc_count > 0) self.doc_count -= 1; @@ -119,7 +131,7 @@ pub const TrigramIndex = struct { errdefer result.deinit(alloc); for (seed) |id| { - if (id != 0) try result.append(alloc, id); + try result.append(alloc, id); } // Intersect with remaining diff --git a/src/turboquant.zig b/src/turboquant.zig index 4be3824..1bd6b2b 100644 --- a/src/turboquant.zig +++ b/src/turboquant.zig @@ -226,8 +226,33 @@ pub const TurboQuant = struct { rotated[j] = self.codebook[idx]; } - // Step 2: Inverse rotate — x̃ = Π^T · ỹ - matvecMulTranspose(self.rotation, rotated, out, d); + // Step 2: Inverse rotate — x~ = inverse(rotation) * y~ + if (self.use_fwht) { + // Inverse SRHT: apply 3 rounds in reverse order (FWHT then sign-flip) + const pd: usize = @intCast(self.padded_dims); + var pad_buf: [1024]f32 = undefined; + const buf = if (pd <= 1024) pad_buf[0..pd] else blk: { + break :blk self.allocator.alloc(f32, pd) catch return; + }; + defer if (pd > 1024) self.allocator.free(buf); + + @memcpy(buf[0..d], rotated); + @memset(buf[d..pd], 0); + + // Reverse order: round 2, 1, 0 + var round: usize = 3; + while (round > 0) { + round -= 1; + fwht(buf, pd); + const signs = self.sign_flips[round * pd ..][0..pd]; + for (0..pd) |i| { + buf[i] *= @as(f32, @floatFromInt(signs[i])); + } + } + @memcpy(out[0..d], buf[0..d]); + } else { + matvecMulTranspose(self.rotation, rotated, out, d); + } } /// Asymmetric L2 distance: FP32 query vs quantized database vector. diff --git a/src/wire.zig b/src/wire.zig index 3e0553b..91d796b 100644 --- a/src/wire.zig +++ b/src/wire.zig @@ -23,18 +23,19 @@ const STATUS_NOT_FOUND: u8 = 0x01; const STATUS_ERROR: u8 = 0x02; const HDR: usize = 5; -const MAX_FRAME: usize = 1048576; +const MAX_FRAME: usize = RD_BUF; const RD_BUF: usize = 65536; const WR_BUF: usize = 131072; - +const MAX_WIRE_CONNECTIONS: u32 = 512; pub const WireServer = struct { db: *Database, port: u16, running: std.atomic.Value(bool), activity: activity.ActivityTracker, + conn_count: std.atomic.Value(u32), pub fn init(db: *Database, port: u16) WireServer { - return .{ .db = db, .port = port, .running = std.atomic.Value(bool).init(false), .activity = activity.ActivityTracker.init() }; + return .{ .db = db, .port = port, .running = std.atomic.Value(bool).init(false), .activity = activity.ActivityTracker.init(), .conn_count = std.atomic.Value(u32).init(0) }; } pub fn run(self: *WireServer) !void { @@ -45,7 +46,16 @@ pub const WireServer = struct { std.log.info("TurboDB wire protocol on :{d}", .{self.port}); while (self.running.load(.acquire)) { const conn = listener.accept() catch continue; - const t = std.Thread.spawn(.{}, handleConn, .{ self, conn }) catch continue; + if (self.conn_count.load(.acquire) >= MAX_WIRE_CONNECTIONS) { + conn.stream.close(); + continue; + } + _ = self.conn_count.fetchAdd(1, .acq_rel); + const t = std.Thread.spawn(.{}, handleConn, .{ self, conn }) catch { + _ = self.conn_count.fetchSub(1, .acq_rel); + conn.stream.close(); + continue; + }; t.detach(); } } @@ -75,7 +85,16 @@ pub const WireServer = struct { const client_fd = std.posix.accept(fd, @ptrCast(&client_addr), &addr_len, 0) catch continue; const stream = std.net.Stream{ .handle = client_fd }; const conn = std.net.Server.Connection{ .stream = stream, .address = std.net.Address.initUnix(path) catch continue }; - const t = std.Thread.spawn(.{}, handleConn, .{ self, conn }) catch continue; + if (self.conn_count.load(.acquire) >= MAX_WIRE_CONNECTIONS) { + conn.stream.close(); + continue; + } + _ = self.conn_count.fetchAdd(1, .acq_rel); + const t = std.Thread.spawn(.{}, handleConn, .{ self, conn }) catch { + _ = self.conn_count.fetchSub(1, .acq_rel); + conn.stream.close(); + continue; + }; t.detach(); } } @@ -89,19 +108,27 @@ const Bufs = struct { rd: [RD_BUF]u8, wr: [WR_BUF]u8 }; fn handleConn(srv: *WireServer, conn: std.net.Server.Connection) void { defer conn.stream.close(); + defer _ = srv.conn_count.fetchSub(1, .acq_rel); const bufs = std.heap.page_allocator.create(Bufs) catch return; defer std.heap.page_allocator.destroy(bufs); - // Pre-resolve the most common collection for the fast path - var cached_col: ?*collection_mod.Collection = null; + // Cache only the collection name for fast-path lookups; the pointer + // is re-resolved every time via lookupCollection (O(1) RwLock + HashMap) + // to avoid dangling pointers when another connection drops a collection. var cached_col_name: [128]u8 = undefined; var cached_col_len: usize = 0; var rp: usize = 0; while (true) { - if (rp >= RD_BUF) rp = 0; - const n = conn.stream.read(bufs.rd[rp..]) catch return; + // Clamp read to remaining buffer space. + const remaining = RD_BUF - rp; + if (remaining == 0) { + // Buffer full with an incomplete frame — the frame is larger + // than RD_BUF so it can never be processed. Close connection. + return; + } + const n = conn.stream.read(bufs.rd[rp..RD_BUF]) catch return; if (n == 0) return; rp += n; @@ -116,12 +143,12 @@ fn handleConn(srv: *WireServer, conn: std.net.Server.Connection) void { // ── FAST PATH: inline GET with zero-copy ── if (op == OP_GET) { srv.activity.recordQuery(); - const wn = fastGet(srv, payload, &bufs.wr, &cached_col, &cached_col_name, &cached_col_len); + const wn = fastGet(srv, payload, &bufs.wr, &cached_col_name, &cached_col_len); conn.stream.writeAll(bufs.wr[0..wn]) catch return; } else { - // Invalidate collection cache on writes + // Invalidate collection name cache on writes if (op == OP_INSERT or op == OP_UPDATE or op == OP_DELETE) { - cached_col = null; + cached_col_len = 0; } srv.activity.recordQuery(); const wn = dispatch(srv, op, payload, &bufs.wr); @@ -141,26 +168,23 @@ fn fastGet( srv: *WireServer, p: []const u8, w: *[WR_BUF]u8, - cached_col: *?*collection_mod.Collection, cached_name: *[128]u8, cached_len: *usize, ) usize { const a = parseKey(p) orelse return errResp(w, OP_GET, STATUS_ERROR); - // Fast collection lookup: skip mutex if same collection as last call - const col = blk: { - if (cached_col.*) |cc| { - if (cached_len.* == a.col.len and std.mem.eql(u8, cached_name[0..cached_len.*], a.col)) { - break :blk cc; - } + // Always re-resolve the collection pointer via lookupCollection to + // avoid dangling pointers. The DB's internal path (RwLock read + + // HashMap get) is already O(1), so the overhead is negligible. + // We still cache the name so callers can invalidate on writes. + if (cached_len.* != a.col.len or !std.mem.eql(u8, cached_name[0..cached_len.*], a.col)) { + if (a.col.len <= cached_name.len) { + @memcpy(cached_name[0..a.col.len], a.col); + cached_len.* = a.col.len; } - const c = lookupCollection(srv, a.col) catch return errResp(w, OP_GET, STATUS_ERROR); - @memcpy(cached_name[0..a.col.len], a.col); - cached_len.* = a.col.len; - cached_col.* = c; - break :blk c; - }; + } + const col = lookupCollection(srv, a.col) catch return errResp(w, OP_GET, STATUS_ERROR); const d = col.get(a.key) orelse return errResp(w, OP_GET, STATUS_NOT_FOUND); // Write response: [len:4][op:1][status:1][doc_id:8][ver:1][val_len:4][val:N] @@ -275,7 +299,8 @@ fn doScan(srv: *WireServer, p: []const u8, w: *[WR_BUF]u8) usize { var pos: usize = HDR + 1 + 4; // header + status + count w[4] = OP_SCAN; w[5] = STATUS_OK; - wrU32LE(w[6..10], @intCast(result.docs.len)); + // Count placeholder — will be overwritten after the loop + var written_count: u32 = 0; for (result.docs) |d| { const dsz = 8 + 1 + 2 + d.key.len + 4 + d.value.len; @@ -288,7 +313,11 @@ fn doScan(srv: *WireServer, p: []const u8, w: *[WR_BUF]u8) usize { if (d.value.len > 0) @memcpy(w[pos + 15 + d.key.len ..][0..d.value.len], d.value); pos += dsz; + written_count += 1; } + // Write the actual count of docs serialized (may be less than result.docs.len + // if the write buffer was exhausted) + wrU32LE(w[6..10], written_count); wrU32BE(w, @intCast(pos)); return pos; }