From b69bc44aaf9e69429711baef1bc575170120bd2e Mon Sep 17 00:00:00 2001
From: Emil Lerch <emil@lerch.org>
Date: Tue, 1 Aug 2023 19:42:48 -0700
Subject: [PATCH] add allocator parameter

---
 src/lambda.zig      | 16 +++++++++-------
 src/sample-main.zig |  2 +-
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/lambda.zig b/src/lambda.zig
index 6ce043f..f8cf7a8 100644
--- a/src/lambda.zig
+++ b/src/lambda.zig
@@ -2,7 +2,9 @@ const std = @import("std");
 
 const HandlerFn = *const fn (std.mem.Allocator, []const u8) anyerror![]const u8;
 
-pub fn run(event_handler: HandlerFn) !void { // TODO: remove inferred error set?
+/// Starts the lambda framework. Handler will be called when an event is processing
+/// If an allocator is not provided, an approrpriate allocator will be selected and used
+pub fn run(allocator: ?std.mem.Allocator, event_handler: HandlerFn) !void { // TODO: remove inferred error set?
     const prefix = "http://";
     const postfix = "/2018-06-01/runtime/invocation";
     const lambda_runtime_uri = std.os.getenv("AWS_LAMBDA_RUNTIME_API").?;
@@ -10,20 +12,20 @@ pub fn run(event_handler: HandlerFn) !void { // TODO: remove inferred error set?
 
     var gpa = std.heap.GeneralPurposeAllocator(.{}){};
     defer _ = gpa.deinit();
-    const allocator = gpa.allocator();
+    const alloc = allocator orelse gpa.allocator();
 
-    const url = try std.fmt.allocPrint(allocator, "{s}{s}{s}/next", .{ prefix, lambda_runtime_uri, postfix });
-    defer allocator.free(url);
+    const url = try std.fmt.allocPrint(alloc, "{s}{s}{s}/next", .{ prefix, lambda_runtime_uri, postfix });
+    defer alloc.free(url);
     const uri = try std.Uri.parse(url);
 
-    var client: std.http.Client = .{ .allocator = allocator };
+    var client: std.http.Client = .{ .allocator = alloc };
     defer client.deinit();
-    var empty_headers = std.http.Headers.init(allocator);
+    var empty_headers = std.http.Headers.init(alloc);
     defer empty_headers.deinit();
     std.log.info("Bootstrap initializing with event url: {s}", .{url});
 
     while (true) {
-        var req_alloc = std.heap.ArenaAllocator.init(allocator);
+        var req_alloc = std.heap.ArenaAllocator.init(alloc);
         defer req_alloc.deinit();
         const req_allocator = req_alloc.allocator();
         var req = try client.request(.GET, uri, empty_headers, .{});
diff --git a/src/sample-main.zig b/src/sample-main.zig
index 5fec78e..555d85a 100644
--- a/src/sample-main.zig
+++ b/src/sample-main.zig
@@ -2,7 +2,7 @@ const std = @import("std");
 const lambda = @import("lambda.zig");
 
 pub fn main() anyerror!void {
-    try lambda.run(handler);
+    try lambda.run(null, handler);
 }
 
 fn handler(allocator: std.mem.Allocator, event_data: []const u8) ![]const u8 {