Zig 分配器的应用

发布: 2024-06-16   上次更新: 2024-07-14

文章目录

原文地址: https://www.openmymind.net/Leveraging-Zigs-Allocators/

假设我们想为Zig编写一个 HTTP服务器库。这个库的核心可能是线程池,用于处理请求。以简化的方式来看,它可能类似于:

1
2
3
4
5
6
7
fn run(worker: *Worker) void {
    while (queue.pop()) |conn| {
        const action = worker.route(conn.req.url);
        action(conn.req, conn.res) catch { // TODO: 500 };
        worker.write(conn.res);
    }
}

作为这个库的用户,您可能会编写一些动态内容的操作。如果假设在启动时为服务器提供分配器(Allocator),则可以将此分配器传递给动作:

1
2
3
4
5
6
7
8
fn run(worker: *Worker) void {
    const allocator = worker.server.allocator;
    while (queue.pop()) |conn| {
        const action = worker.route(conn.req.url);
        action(allocator, conn.req, conn.res) catch { // TODO: 500 };
        worker.write(conn.res);
    }
}

这允许用户编写如下的操作:

1
2
3
4
5
fn greet(allocator: Allocator, req: *http.Request, res: *http.Response) !void {
    const name = req.query("name") orelse "guest";
    res.status = 200;
    res.body = try std.fmt.allocPrint(allocator, "Hello {s}", .{name});
}

虽然这是一个正确的方向,但存在明显的问题:分配的问候语从未被释放。我们的run函数不能在写回应后就调用allocator.free(conn.res.body),因为在某些情况下,主体可能不需要被释放。我们可以通过使动作必须 write() 回应并因此能够free它所做的任何分配来结构化API,但这将使得添加一些功能变得不可能,比如支持中间件。

最佳和最简单的方法是使用 ArenaAllocator 。其工作原理很简单:当我们deinit时,所有分配都被释放。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
fn run(worker: *Worker) void {
    const allocator = worker.server.allocator;
    while (queue.pop()) |conn| {
        var arena = std.heap.ArenaAllocator.init(allocator);
        defer arena.deinit();
        const action = worker.route(conn.req.url);
        action(arena.allocator(), conn.req, conn.res) catch { // TODO: 500 };
        worker.write(conn.res);
    }
}

std.mem.Allocator 是一个 “接口” ,我们的动作无需更改。 ArenaAllocator 对HTTP服务器来说是一个很好的选择,因为它们与请求绑定,具有明确/可理解的生命周期,并且相对短暂。虽然有可能滥用它们,但可以说:使用更多!

我们可以更进一步并重用相同的Arena。这可能看起来不太有用,但是请看:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
fn run(worker: *Worker) void {
    const allocator = worker.server.allocator;
    var arena = std.heap.ArenaAllocator.init(allocator);
    defer arena.deinit();
    while (queue.pop()) |conn| {
        // 魔法在此处!
        defer _ = arena.reset(.{.retain_with_limit = 8192});
        const action = worker.route(conn.req.url);
        action(arena.allocator(), conn.req, conn.res) catch { // TODO: 500 };
        worker.write(conn.res);
    }
}

我们将Arena移出了循环,但重要的部分在内部:每个请求后,我们重置了Arena并保留最多8K内存。这意味着对于许多请求,我们无需访问底层分配器(worker.server.allocator)。这种方法简化了内存管理。

现在想象一下,如果我们不能用 retain_with_limit 重置 Arena,我们还能进行同样的优化吗?可以,我们可以创建自己的分配器,首先尝试使用固定缓冲区分配器(FixedBufferAllocator),如果分配适配,回退到 Arena 分配器。

这里是 FallbackAllocator 的完整示例:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
const FallbackAllocator = struct {
  primary: Allocator,
  fallback: Allocator,
  fba: *std.heap.FixedBufferAllocator,

  pub fn allocator(self: *FallbackAllocator) Allocator {
    return .{
      .ptr = self,
      .vtable = &.{.alloc = alloc, .resize = resize, .free = free},
    };
  }

  fn alloc(ctx: *anyopaque, len: usize, ptr_align: u8, ra: usize) ?[*]u8 {
    const self: *FallbackAllocator = @ptrCast(@alignCast(ctx));
    return self.primary.rawAlloc(len, ptr_align, ra)
           orelse self.fallback.rawAlloc(len, ptr_align, ra);
  }

  fn resize(ctx: *anyopaque, buf: []u8, buf_align: u8, new_len: usize, ra: usize) bool {
    const self: *FallbackAllocator = @ptrCast(@alignCast(ctx));
    if (self.fba.ownsPtr(buf.ptr)) {
      if (self.primary.rawResize(buf, buf_align, new_len, ra)) {
        return true;
      }
    }
    return self.fallback.rawResize(buf, buf_align, new_len, ra);
  }

  fn free(_: *anyopaque, _: []u8, _: u8, _: usize) void {
    // we noop this since, in our specific case, we know
    // the fallback is an arena, which won't free individual items
  }
};

我们的alloc实现首先尝试使用我们定义的"主"分配器进行分配。如果失败,我们会使用"备用"分配器。作为std.mem.Allocator接口的一部分,我们需要实现的resize方法会确定正在尝试扩展内存的所有者,并然后调用其rawResize方法。为了保持代码简单,我在这里省略了free方法的具体实现——在这种特定情况下是可以接受的,因为我们计划使用"主"分配器作为FixedBufferAllocator,而"备用"分配器则会是ArenaAllocator(因此所有释放操作会在arena的deinitreset时进行)。

接下来我们需要改变我们的run方法以利用这个新的分配器:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
fn run(worker: *Worker) void {
    const allocator = worker.server.allocator; // 这是FixedBufferAllocator底层的内存
    const buf = try allocator.alloc(u8, 8192); // 分配8K字节的内存用于存储数据
    defer allocator.free(buf); // 完成后释放内存

    var fba = std.heap.FixedBufferAllocator.init(buf); // 初始化FixedBufferAllocator

    while (queue.pop()) |conn| {
        defer fba.reset(); // 重置FixedBufferAllocator,准备处理下一个请求

        var arena = std.heap.ArenaAllocator.init(allocator); // 初始化ArenaAllocator用于分配额外内存
        defer arena.deinit();

        var fallback = FallbackAllocator{
            .fba = &fba,
            .primary = fba.allocator(),
            .fallback = arena.allocator(),
        }; // 创建FallbackAllocator,包含FixedBufferAllocator和ArenaAllocator

        const action = worker.route(conn.req.url); // 路由请求到对应的动作处理函数
        action(fallback.allocator(), conn.req, conn.res) catch { // 处理动作执行中的错误 };

        worker.write(conn.res); // 写回响应信息给客户端
    }
}

这种方法实现了类似于在retain_with_limit中重置arena的功能。我们创建了一个可以重复使用的FixedBufferAllocator,用于处理每个请求的8K字节内存需求。由于一个动作可能需要更多的内存,我们仍然需要ArenaAllocator来提供额外的空间。通过将FixedBufferAllocatorArenaAllocator包裹在我们的FallbackAllocator中,我们可以确保任何分配都首先尝试使用(非常快的)FixedBufferAllocator,当其空间用尽时,则会切换到ArenaAllocator

我们通过暴露std.mem.Allocator接口,可以调整如何工作而不破坏greet。这不仅简化了资源管理(例如通过ArenaAllocator),而且通过重复使用分配来提高了性能(类似于我们做的retain_with_limitFixedBufferAllocator的操作)。

这个示例应该能突出显示我认为明确的分配器提供的两个实际优势:

  1. 简化资源管理(通过类似ArenaAllocator的方式)
  2. 通过重用分配来提高性能(例如我们之前在 retain_with_limitFixedBufferAllocator 时所做的一样)

建议/反馈✉️

  1. 关注微信公众号,加微信群与更多人一起畅聊 Zig
  2. 发现内容错误或链接失效?欢迎提交 PR
  3. 想要分享 Zig 使用经验,欢迎给我们供稿