feat: added TypedStorage to Context

mookums · mookums · commit b3438fd85210 · 2025-02-17T16:52:29.000-08:00
This allows for you to have values that persist throughout the lifetime
of the Context.
diff --git a/examples/middleware/main.zig b/examples/middleware/main.zig
@@ -24,10 +24,15 @@ fn root_handler(ctx: *const Context, id: i8) !Respond {
         \\ <body>
         \\ <h1>Hello, World!</h1>
         \\ <p>id: {d}</p>
+        \\ <p>stored: {d}</p>
         \\ </body>
         \\ </html>
     ;
-    const body = try std.fmt.allocPrint(ctx.allocator, body_fmt, .{id});
+    const body = try std.fmt.allocPrint(
+        ctx.allocator,
+        body_fmt,
+        .{ id, ctx.storage.get(usize).? },
+    );
 
     // This is the standard response and what you
     // will usually be using. This will send to the
@@ -41,6 +46,7 @@ fn root_handler(ctx: *const Context, id: i8) !Respond {
 
 fn passing_middleware(next: *Next, _: void) !Respond {
     log.info("pass middleware: {s}", .{next.context.request.uri.?});
+    try next.context.storage.put(usize, 100);
     return try next.run();
 }
 
diff --git a/src/core/typed_storage.zig b/src/core/typed_storage.zig
@@ -0,0 +1,74 @@
+const std = @import("std");
+
+pub const TypedStorage = struct {
+    arena: std.heap.ArenaAllocator,
+    storage: std.AutoHashMapUnmanaged(u64, *anyopaque),
+
+    pub fn init(allocator: std.mem.Allocator) TypedStorage {
+        return .{
+            .arena = std.heap.ArenaAllocator.init(allocator),
+            .storage = std.AutoHashMapUnmanaged(u64, *anyopaque){},
+        };
+    }
+
+    pub fn deinit(self: *TypedStorage) void {
+        self.arena.deinit();
+    }
+
+    /// Clears the Storage.
+    pub fn clear(self: *TypedStorage) void {
+        self.storage.clearAndFree(self.arena.allocator());
+        _ = self.arena.reset(.retain_capacity);
+    }
+
+    /// Inserts a value into the Storage.
+    /// It uses the given type as the K.
+    pub fn put(self: *TypedStorage, comptime T: type, value: T) !void {
+        const allocator = self.arena.allocator();
+        const ptr = try allocator.create(T);
+        ptr.* = value;
+        const type_id = comptime std.hash.Wyhash.hash(0, @typeName(T));
+        try self.storage.put(allocator, type_id, @ptrCast(ptr));
+    }
+
+    /// Extracts a value out of the Storage.
+    /// It uses the given type as the K.
+    pub fn get(self: *TypedStorage, comptime T: type) ?T {
+        const type_id = comptime std.hash.Wyhash.hash(0, @typeName(T));
+        const ptr = self.storage.get(type_id) orelse return null;
+        return @as(*T, @ptrCast(@alignCast(ptr))).*;
+    }
+};
+
+const testing = std.testing;
+
+test "TypedStorage: Basic" {
+    var storage = TypedStorage.init(testing.allocator);
+    defer storage.deinit();
+
+    // Test inserting and getting different types
+    try storage.put(u32, 42);
+    try storage.put([]const u8, "hello");
+    try storage.put(f32, 3.14);
+
+    try testing.expectEqual(@as(u32, 42), storage.get(u32).?);
+    try testing.expectEqualStrings("hello", storage.get([]const u8).?);
+    try testing.expectEqual(@as(f32, 3.14), storage.get(f32).?);
+
+    // Test overwriting a value
+    try storage.put(u32, 100);
+    try testing.expectEqual(@as(u32, 100), storage.get(u32).?);
+
+    // Test getting non-existent type
+    try testing.expectEqual(@as(?bool, null), storage.get(bool));
+
+    // Test clearing
+    storage.clear();
+    try testing.expectEqual(@as(?u32, null), storage.get(u32));
+    try testing.expectEqual(@as(?[]const u8, null), storage.get([]const u8));
+    try testing.expectEqual(@as(?f32, null), storage.get(f32));
+
+    // Test inserting after clear
+    try storage.put(u32, 200);
+    try testing.expectEqual(@as(u32, 200), storage.get(u32).?);
+}
diff --git a/src/http/context.zig b/src/http/context.zig
@@ -5,11 +5,13 @@ const Runtime = @import("tardy").Runtime;
 const SecureSocket = @import("../core/secure_socket.zig").SecureSocket;
 
 const Capture = @import("router/routing_trie.zig").Capture;
+
+const TypedStorage = @import("../core/typed_storage.zig").TypedStorage;
 const AnyCaseStringMap = @import("../core/any_case_string_map.zig").AnyCaseStringMap;
 
-// Context is dependent on the server that gets created.
+/// HTTP Context. Contains all of the various information
+/// that will persist throughout the lifetime of this Request/Response.
 pub const Context = struct {
-    const Self = @This();
     allocator: std.mem.Allocator,
     /// Not safe to access unless you are manually sending the headers
     /// and returning the .responded variant of Respond.
@@ -18,6 +20,8 @@ pub const Context = struct {
     /// The Request that triggered this handler.
     request: *const Request,
     response: *Response,
+    /// Storage
+    storage: *TypedStorage,
     /// Socket for this Connection.
     socket: SecureSocket,
     /// Slice of the URL Slug Captures
diff --git a/src/http/server.zig b/src/http/server.zig
@@ -4,6 +4,7 @@ const tag = builtin.os.tag;
 const assert = std.debug.assert;
 const log = std.log.scoped(.@"zzz/http/server");
 
+const TypedStorage = @import("../core/typed_storage.zig").TypedStorage;
 const Pseudoslice = @import("../core/pseudoslice.zig").Pseudoslice;
 const AnyCaseStringMap = @import("../core/any_case_string_map.zig").AnyCaseStringMap;
 
@@ -139,6 +140,7 @@ pub const Provision = struct {
     zc_recv_buffer: ZeroCopy(u8),
     header_buffer: std.ArrayList(u8),
     arena: std.heap.ArenaAllocator,
+    storage: TypedStorage,
     captures: []Capture,
     queries: AnyCaseStringMap,
     request: Request,
@@ -187,14 +189,17 @@ pub const Server = struct {
         respond,
     };
 
-    fn prepare_new_request(state: *State, provision: *Provision, config: ServerConfig) !void {
+    fn prepare_new_request(state: ?*State, provision: *Provision, config: ServerConfig) !void {
+        assert(provision.initalized);
         provision.request.clear();
         provision.response.clear();
+        provision.storage.clear();
         provision.zc_recv_buffer.clear_retaining_capacity();
         provision.header_buffer.clearRetainingCapacity();
         _ = provision.arena.reset(.{ .retain_with_limit = config.connection_arena_bytes_retain });
-        state.* = .{ .request = .header };
         provision.recv_slice = try provision.zc_recv_buffer.get_write_area(config.socket_buffer_bytes);
+
+        if (state) |s| s.* = .{ .request = .header };
     }
 
     pub fn main_frame(
@@ -256,20 +261,12 @@ pub const Server = struct {
                 @panic("attempting to allocate more memory than available. (Captures)");
             };
             provision.queries = AnyCaseStringMap.init(rt.allocator);
+            provision.storage = TypedStorage.init(rt.allocator);
             provision.request = Request.init(rt.allocator);
             provision.response = Response.init(rt.allocator);
             provision.initalized = true;
         }
-
-        defer if (provision.zc_recv_buffer.len > config.list_recv_bytes_retain)
-            provision.zc_recv_buffer.shrink_clear_and_free(config.list_recv_bytes_retain) catch unreachable
-        else
-            provision.zc_recv_buffer.clear_retaining_capacity();
-        defer provision.header_buffer.clearRetainingCapacity();
-        defer _ = provision.arena.reset(.{ .retain_with_limit = config.connection_arena_bytes_retain });
-        defer provision.queries.clearRetainingCapacity();
-        defer provision.request.clear();
-        defer provision.response.clear();
+        defer prepare_new_request(null, provision, config) catch unreachable;
 
         var state: State = .{ .request = .header };
         const buffer = try provision.zc_recv_buffer.get_write_area(config.socket_buffer_bytes);
@@ -387,6 +384,7 @@ pub const Server = struct {
                     .header_buffer = &provision.header_buffer,
                     .request = &provision.request,
                     .response = &provision.response,
+                    .storage = &provision.storage,
                     .socket = secure,
                     .captures = found.captures,
                     .queries = found.queries,
@@ -523,6 +521,7 @@ pub const Server = struct {
                 @panic("attempting to allocate more memory than available. (Captures)");
             };
             provision.queries = AnyCaseStringMap.init(rt.allocator);
+            provision.storage = TypedStorage.init(rt.allocator);
             provision.request = Request.init(rt.allocator);
             provision.response = Response.init(rt.allocator);
         }
diff --git a/src/unit_test.zig b/src/unit_test.zig
@@ -8,6 +8,7 @@ test "zzz unit tests" {
     // Core
     testing.refAllDecls(@import("./core/any_case_string_map.zig"));
     testing.refAllDecls(@import("./core/pseudoslice.zig"));
+    testing.refAllDecls(@import("./core/typed_storage.zig"));
 
     // HTTP
     testing.refAllDecls(@import("./http/context.zig"));
