allow handler more control in status reporting

build.zig.zon

@@ -4,8 +4,8 @@
 
     .dependencies = .{
         .flexilib = .{
-            .url = "https://git.lerch.org/lobo/flexilib/archive/c44ad2ba84df735421bef23a2ad612968fb50f06.tar.gz",
+            .url = "https://git.lerch.org/lobo/flexilib/archive/3d3dab9c792651477932e2b61c9f4794ac694dcb.tar.gz",
-            .hash = "122051fdfeefdd75653d3dd678c8aa297150c2893f5fad0728e0d953481383690dbc",
+            .hash = "1220fd7a614fe3c9f6006b630bba528e2ec9dca9c66f5ff10f7e471ad2bdd41b6c89",
         },
     },
 }

src/helpers.zig

@@ -84,19 +84,7 @@ pub const Headers = struct {
 pub fn allHeaders(allocator: std.mem.Allocator, context: universal_lambda.Context) !Headers {
     switch (context) {
         .web_request => |res| return Headers.init(allocator, &res.request.headers, false),
-        .flexilib => |ctx| {
+        .flexilib => |ctx| return Headers.init(allocator, &ctx.request.headers, false),
-            var headers = try allocator.create(std.http.Headers);
-            errdefer allocator.destroy(headers);
-            headers.allocator = allocator;
-            headers.list = .{};
-            headers.index = .{};
-            headers.owned = true;
-            errdefer headers.deinit();
-            for (ctx.request.headers) |hdr| {
-                try headers.append(hdr.name_ptr[0..hdr.name_len], hdr.value_ptr[0..hdr.value_len]);
-            }
-            return Headers.init(allocator, headers, true);
-        },
         .none => return headersWithoutContext(allocator),
     }
 }
@@ -164,8 +152,9 @@ test "can get headers" {
     // leaks. There doesn't seem to be a way to ignore leak detection
     if (@import("builtin").os.tag == .wasi) return error.SkipZigTest;
     const allocator = std.testing.allocator;
+    var response = universal_lambda.Response.init(allocator);
     const context = universal_lambda.Context{
-        .none = {},
+        .none = &response,
     };
     var headers = try allHeaders(allocator, context);
     defer headers.deinit();

src/lambda.zig

@@ -3,6 +3,7 @@ const builtin = @import("builtin");
 
 const HandlerFn = @import("universal_lambda.zig").HandlerFn;
 const Context = @import("universal_lambda.zig").Context;
+const UniversalLambdaResponse = @import("universal_lambda.zig").Response;
 
 const log = std.log.scoped(.lambda);
 
@@ -20,7 +21,7 @@ pub fn deinit() void {
 /// If an allocator is not provided, an approrpriate allocator will be selected and used
 /// This function is intended to loop infinitely. If not used in this manner,
 /// make sure to call the deinit() function
-pub fn run(allocator: ?std.mem.Allocator, event_handler: HandlerFn) !void { // TODO: remove inferred error set?
+pub fn run(allocator: ?std.mem.Allocator, event_handler: HandlerFn) !u8 { // TODO: remove inferred error set?
     const lambda_runtime_uri = std.os.getenv("AWS_LAMBDA_RUNTIME_API") orelse test_lambda_runtime_uri.?;
     // TODO: If this is null, go into single use command line mode
 
@@ -70,15 +71,20 @@ pub fn run(allocator: ?std.mem.Allocator, event_handler: HandlerFn) !void { // T
         // Lambda does not have context, just environment variables. API Gateway
         // might be configured to pass in lots of context, but this comes through
         // event data, not context.
-        const event_response = event_handler(req_allocator, event.event_data, .{ .none = {} }) catch |err| {
+        var response = UniversalLambdaResponse.init(allocator.?);
+        response.output_file = std.io.getStdOut();
+        const event_response = event_handler(req_allocator, event.event_data, .{ .none = &response }) catch |err| {
+            response.finish() catch unreachable;
             event.reportError(@errorReturnTrace(), err, lambda_runtime_uri) catch unreachable;
             continue;
         };
+        response.finish() catch unreachable;
         event.postResponse(lambda_runtime_uri, event_response) catch |err| {
             event.reportError(@errorReturnTrace(), err, lambda_runtime_uri) catch unreachable;
             continue;
         };
     }
+    return 0;
 }
 
 const Event = struct {
@@ -409,10 +415,13 @@ fn handler(allocator: std.mem.Allocator, event_data: []const u8, context: Contex
     _ = context;
     return event_data;
 }
+fn thread_run(allocator: ?std.mem.Allocator, event_handler: HandlerFn) !void {
+    _ = try run(allocator, event_handler);
+}
 fn test_run(allocator: std.mem.Allocator, event_handler: HandlerFn) !std.Thread {
     return try std.Thread.spawn(
         .{},
-        run,
+        thread_run,
         .{ allocator, event_handler },
     );
 }

src/universal_lambda.zig

@@ -5,22 +5,56 @@ pub const HandlerFn = *const fn (std.mem.Allocator, []const u8, Context) anyerro
 
 const log = std.log.scoped(.universal_lambda);
 
-const FakeResponse = struct {
+pub const Response = struct {
+    allocator: std.mem.Allocator,
+    headers: std.http.Headers,
+    output_file: ?std.fs.File = null,
+    status: std.http.Status = .ok,
+    reason: ?[]const u8 = null,
     request: struct {
         target: []const u8,
         headers: std.http.Headers,
     },
+    al: std.ArrayList(u8),
+
+    pub fn init(allocator: std.mem.Allocator) Response {
+        return .{
+            .allocator = allocator,
+            .headers = .{ .allocator = allocator },
+            .request = .{
+                .target = "/",
+                .headers = .{ .allocator = allocator },
+            },
+            .al = std.ArrayList(u8).init(allocator),
         };
+    }
+    pub fn write(res: *Response, bytes: []const u8) !usize {
+        return res.al.writer().write(bytes);
+    }
+
+    pub fn writeAll(res: *Response, bytes: []const u8) !void {
+        return res.al.writer().writeAll(bytes);
+    }
+
+    pub fn writer(res: *Response) std.io.Writer {
+        return res.al.writer().writer();
+    }
+
+    pub fn finish(res: *Response) !void {
+        if (res.output_file) |f| {
+            try f.writer().writeAll(res.al.items);
+        }
+        res.al.deinit();
+    }
+};
+
 pub const Context = union(enum) {
     web_request: switch (build_options.build_type) {
-        .exe_run, .cloudflare => *FakeResponse,
+        .exe_run, .cloudflare => *Response,
         else => *std.http.Server.Response,
     },
-    flexilib: struct {
+    flexilib: *flexilib.ZigResponse,
-        request: flexilib.ZigRequest,
+    none: *Response,
-        response: flexilib.ZigResponse,
-    },
-    none: void,
 };
 
 const runFn = blk: {
@@ -42,11 +76,11 @@ fn deinit() void {
 /// If an allocator is not provided, an approrpriate allocator will be selected and used
 /// This function is intended to loop infinitely. If not used in this manner,
 /// make sure to call the deinit() function
-pub fn run(allocator: ?std.mem.Allocator, event_handler: HandlerFn) !void { // TODO: remove inferred error set?
+pub fn run(allocator: ?std.mem.Allocator, event_handler: HandlerFn) !u8 { // TODO: remove inferred error set?
-    try runFn(allocator, event_handler);
+    return try runFn(allocator, event_handler);
 }
 
-fn runExe(allocator: ?std.mem.Allocator, event_handler: HandlerFn) !void {
+fn runExe(allocator: ?std.mem.Allocator, event_handler: HandlerFn) !u8 {
     var arena = std.heap.ArenaAllocator.init(allocator orelse std.heap.page_allocator);
     defer arena.deinit();
 
@@ -56,14 +90,30 @@ fn runExe(allocator: ?std.mem.Allocator, event_handler: HandlerFn) !void {
     // We're setting up an arena allocator. While we could use a gpa and get
     // some additional safety, this is now "production" runtime, and those
     // things are better handled by unit tests
-    const writer = std.io.getStdOut().writer();
+    var response = Response.init(aa);
-    try writer.writeAll(try event_handler(aa, data, .{ .none = {} }));
+
+    // Note here we are throwing out the status and reason. This is to make
+    // the console experience less "webby" and more "consoly", at the potential
+    // cost of data loss for not outputting the http status/reason
+    const output = event_handler(aa, data, .{ .none = &response }) catch |err| {
+        response.output_file = std.io.getStdErr();
+        try response.finish();
+        return err;
+    };
+
+    response.output_file = if (response.status.class() == .success) std.io.getStdOut() else std.io.getStdErr();
+    const writer = response.output_file.?.writer();
+    try response.finish();
+    try writer.writeAll(output);
     try writer.writeAll("\n");
+    // We might have gotten an error message managed directly by the event handler
+    // If that's the case, we will need to report back an error code
+    return if (response.status.class() == .success) 0 else 1;
 }
 
 /// Will create a web server and marshall all requests back to our event handler
 /// To keep things simple, we'll have this on a single thread, at least for now
-fn runStandaloneServer(allocator: ?std.mem.Allocator, event_handler: HandlerFn) !void {
+fn runStandaloneServer(allocator: ?std.mem.Allocator, event_handler: HandlerFn) !u8 {
     const alloc = allocator orelse std.heap.page_allocator;
 
     var arena = std.heap.ArenaAllocator.init(alloc);
@@ -98,6 +148,7 @@ fn runStandaloneServer(allocator: ?std.mem.Allocator, event_handler: HandlerFn)
             }
         };
     }
+    return 0;
 }
 
 fn processRequest(aa: std.mem.Allocator, server: *std.http.Server, event_handler: HandlerFn) !void {