final changes

index.zig

@@ -5,6 +5,8 @@ var tree_arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
 const tree_allocator = tree_arena.allocator();
 const temp_allocator = std.heap.c_allocator;
 
+var tree_mutex: std.Thread.Mutex = .{};
+
 const RadixNode = struct {
     edge: []const u8,
     children: std.StringHashMap(*RadixNode),
@@ -50,6 +52,9 @@ fn commonPrefixLen(a: []const u8, b: []const u8) usize {
 }
 
 pub fn insert(key: []const u8) void {
+    tree_mutex.lock();
+    defer tree_mutex.unlock();
+
     ensureRoot();
     if (key.len == 0) return;
 
@@ -116,6 +121,9 @@ pub fn insert(key: []const u8) void {
 }
 
 pub fn delete(key: []const u8) void {
+    tree_mutex.lock();
+    defer tree_mutex.unlock();
+
     ensureRoot();
     if (key.len == 0) return;
 
@@ -218,6 +226,9 @@ pub fn getKeysFromNode(node: *RadixNode, prefix: []const u8) [][]const u8 {
 }
 
 pub fn getKeysByPrefix(prefix: []const u8) []const u8 {
+    tree_mutex.lock();
+    defer tree_mutex.unlock();
+
     ensureRoot();
     const node = searchByPrefix(prefix) orelse return "";
     const keys = getKeysFromNode(node, prefix);
@@ -226,6 +237,9 @@ pub fn getKeysByPrefix(prefix: []const u8) []const u8 {
 }
 
 pub fn getValuesByPrefix(prefix: []const u8) []const u8 {
+    tree_mutex.lock();
+    defer tree_mutex.unlock();
+
     ensureRoot();
     const node = searchByPrefix(prefix) orelse return "";
     const keys = getKeysFromNode(node, prefix);
@@ -241,6 +255,9 @@ pub fn getValuesByPrefix(prefix: []const u8) []const u8 {
 }
 
 pub fn getAllKeys() []const u8 {
+    tree_mutex.lock();
+    defer tree_mutex.unlock();
+
     ensureRoot();
     const keys = getKeysFromNode(root, &[_]u8{});
     if (keys.len == 0) return "";

main.zig

@@ -12,6 +12,7 @@ const redis = @import("redis.zig");
 
 const PORT = 8085;
 var should_exit = std.atomic.Value(bool).init(false);
+var active_connections = std.atomic.Value(u32).init(0);
 var redis_mode = false;
 
 fn handleSignal(sig: c_int) callconv(.c) void {
@@ -101,12 +102,29 @@ pub fn main() !void {
     }
 
     std.debug.print("\nShutdown signal received...\n", .{});
+
+    const max_wait_ms = 5000;
+    const wait_interval_ms = 100;
+    var waited_ms: u32 = 0;
+
+    while (active_connections.load(.seq_cst) > 0 and waited_ms < max_wait_ms) {
+        posix.nanosleep(0, wait_interval_ms * std.time.ns_per_ms);
+        waited_ms += wait_interval_ms;
+    }
+
+    const remaining = active_connections.load(.seq_cst);
+    if (remaining > 0) {
+        std.debug.print("Warning: {d} connections still active after {d}ms, forcing shutdown...\n", .{ remaining, max_wait_ms });
+    }
+
     persistence.flush() catch |err| {
         std.debug.print("Failed to flush persistence: {any}\n", .{err});
     };
 }
 
 pub fn handleConnection(conn: posix.socket_t) !void {
+    _ = active_connections.fetchAdd(1, .seq_cst);
+    defer _ = active_connections.fetchSub(1, .seq_cst);
     defer posix.close(conn);
 
     var requestBuffer: [1024 * 1024]u8 = undefined;
@@ -139,6 +157,8 @@ pub fn handleConnection(conn: posix.socket_t) !void {
 }
 
 pub fn handleRedisConnection(conn: posix.socket_t) !void {
+    _ = active_connections.fetchAdd(1, .seq_cst);
+    defer _ = active_connections.fetchSub(1, .seq_cst);
     defer posix.close(conn);
 
     var requestBuffer: [2 * 1024 * 1024]u8 = undefined;

storage.zig

@@ -4,7 +4,7 @@ const index = @import("index.zig");
 const hashing = @import("hashing.zig");
 const persistence = @import("persistence.zig");
 
-const NUM_SHARDS = 32;
+const NUM_SHARDS = 64;
 const TOTAL_BUCKETS = 1_048_576;
 const BUCKETS_PER_SHARD = TOTAL_BUCKETS / NUM_SHARDS;