Zig 分配器的應用

假設我們想爲 Zig 編寫一個 HTTP 服務器庫 [1]。這個庫的核心可能是線程池,用於處理請求。以簡化的方式來看,它可能類似於:

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),則可以將此分配器傳遞給動作:

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);
    }
}

這允許用戶編寫如下的操作:

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時,所有分配都被釋放。

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 是一個 " 接口 [2]" ,我們的動作無需更改。 ArenaAllocator 對 HTTP 服務器來說是一個很好的選擇,因爲它們與請求綁定,具有明確 / 可理解的生命週期,並且相對短暫。雖然有可能濫用它們,但可以說:使用更多!

我們可以更進一步並重用相同的 Arena。這可能看起來不太有用,但是請看:

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 的完整示例:

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方法以利用這個新的分配器:

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. 1. 簡化資源管理(通過類似ArenaAllocator的方式)

  2. 2. 通過重用分配來提高性能

建議 / 反饋 ✉️

引用鏈接

[1] HTTP 服務器庫: https://github.com/karlseguin/http.zig
[2] 接口: https://www.openmymind.net/Zig-Interfaces/
[3] 歡迎給我們供稿: https://ziglang.cc/post/2023/09/05/hello-world/

本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源https://mp.weixin.qq.com/s/L_TqtpQG9jg9nhv0QPiCcA