Improve reliabiltiy of multithreaded code by using threadlocal mimalloc heaps

2 files changed, 125 insertions, 49 deletions

diff --git a/src/memory_allocator.zig b/src/memory_allocator.zig
index aeb91dfac..9b48d6e66 100644
--- a/src/memory_allocator.zig
+++ b/src/memory_allocator.zig
@@ -57,48 +57,26 @@ const CAllocator = struct {
         return @intToPtr(*[*]u8, @ptrToInt(ptr) - @sizeOf(usize));
     }
 
-    fn alignedAlloc(len: usize, alignment: usize) ?[*]u8 {
-        if (supports_posix_memalign) {
-            // The posix_memalign only accepts alignment values that are a
-            // multiple of the pointer size
-            const eff_alignment = @maximum(alignment, @sizeOf(usize));
-
-            var aligned_ptr: ?*anyopaque = null;
-            if (c.posix_memalign(&aligned_ptr, eff_alignment, len) != 0)
-                return null;
-
-            return @ptrCast([*]u8, aligned_ptr);
-        }
+    const MI_MAX_ALIGN_SIZE = 16;
+    inline fn mi_malloc_satisfies_alignment(alignment: usize, size: usize) bool {
+        return (alignment == @sizeOf(*anyopaque) or (alignment == MI_MAX_ALIGN_SIZE and size > (MI_MAX_ALIGN_SIZE / 2)));
+    }
 
-        // Thin wrapper around regular malloc, overallocate to account for
-        // alignment padding and store the orignal malloc()'ed pointer before
-        // the aligned address.
-        var unaligned_ptr = @ptrCast([*]u8, c.malloc(len + alignment - 1 + @sizeOf(usize)) orelse return null);
-        const unaligned_addr = @ptrToInt(unaligned_ptr);
-        const aligned_addr = mem.alignForward(unaligned_addr + @sizeOf(usize), alignment);
-        var aligned_ptr = unaligned_ptr + (aligned_addr - unaligned_addr);
-        getHeader(aligned_ptr).* = unaligned_ptr;
+    fn alignedAlloc(len: usize, alignment: usize) ?[*]u8 {
+        var ptr = if (mi_malloc_satisfies_alignment(alignment, len))
+            mimalloc.mi_malloc(len)
+        else
+            mimalloc.mi_malloc_aligned(len, alignment);
 
-        return aligned_ptr;
+        return @ptrCast([*]u8, ptr orelse null);
     }
 
     fn alignedFree(ptr: [*]u8) void {
-        if (supports_posix_memalign) {
-            return c.free(ptr);
-        }
-
-        const unaligned_ptr = getHeader(ptr).*;
-        c.free(unaligned_ptr);
+        return c.free(ptr);
     }
 
     fn alignedAllocSize(ptr: [*]u8) usize {
-        if (supports_posix_memalign) {
-            return CAllocator.malloc_size(ptr);
-        }
-
-        const unaligned_ptr = getHeader(ptr).*;
-        const delta = @ptrToInt(ptr) - @ptrToInt(unaligned_ptr);
-        return CAllocator.malloc_size(unaligned_ptr) - delta;
+        return CAllocator.malloc_size(ptr);
     }
 
     fn alloc(
@@ -116,15 +94,7 @@ const CAllocator = struct {
         if (len_align == 0) {
             return ptr[0..len];
         }
-        const full_len = init: {
-            if (CAllocator.supports_malloc_size) {
-                const s = alignedAllocSize(ptr);
-                assert(s >= len);
-                break :init s;
-            }
-            break :init len;
-        };
-        return ptr[0..mem.alignBackwardAnyAlign(full_len, len_align)];
+        return ptr[0..mem.alignBackwardAnyAlign(mimalloc.mi_usable_size(ptr), len_align)];
     }
 
     fn resize(
@@ -140,12 +110,12 @@ const CAllocator = struct {
         if (new_len <= buf.len) {
             return mem.alignAllocLen(buf.len, new_len, len_align);
         }
-        if (CAllocator.supports_malloc_size) {
-            const full_len = alignedAllocSize(buf.ptr);
-            if (new_len <= full_len) {
-                return mem.alignAllocLen(full_len, new_len, len_align);
-            }
+
+        const full_len = alignedAllocSize(buf.ptr);
+        if (new_len <= full_len) {
+            return mem.alignAllocLen(full_len, new_len, len_align);
         }
+
         return null;
     }
 
@@ -157,7 +127,7 @@ const CAllocator = struct {
     ) void {
         _ = buf_align;
         _ = return_address;
-        alignedFree(buf.ptr);
+        mimalloc.mi_free(buf.ptr);
     }
 };
 
diff --git a/src/mimalloc_arena.zig b/src/mimalloc_arena.zig
new file mode 100644
index 000000000..9fa0b586e
--- /dev/null
+++ b/src/mimalloc_arena.zig
@@ -0,0 +1,106 @@
+const mem = @import("std").mem;
+const builtin = @import("std").builtin;
+const std = @import("std");
+
+const mimalloc = @import("./allocators/mimalloc.zig");
+
+const Allocator = mem.Allocator;
+const assert = std.debug.assert;
+
+pub const Arena = struct {
+    heap: *mimalloc.mi_heap_t = undefined,
+
+    pub fn backingAllocator(this: Arena) Allocator {
+        var arena = Arena{ .heap = this.heap.backing() };
+        return arena.allocator();
+    }
+
+    pub fn allocator(this: Arena) Allocator {
+        return Allocator{ .ptr = this.heap, .vtable = &c_allocator_vtable };
+    }
+
+    pub fn deinit(this: *Arena) void {
+        mimalloc.mi_heap_destroy(this.heap);
+    }
+
+    pub fn init() !Arena {
+        return Arena{ .heap = mimalloc.mi_heap_new() orelse return error.OutOfMemory };
+    }
+
+    const MI_MAX_ALIGN_SIZE = 16;
+    inline fn mi_malloc_satisfies_alignment(alignment: usize, size: usize) bool {
+        return (alignment == @sizeOf(*anyopaque) or (alignment == MI_MAX_ALIGN_SIZE and size > (MI_MAX_ALIGN_SIZE / 2)));
+    }
+
+    fn alignedAlloc(heap: *mimalloc.mi_heap_t, len: usize, alignment: usize) ?[*]u8 {
+        // this is the logic that posix_memalign does
+        var ptr = if (mi_malloc_satisfies_alignment(alignment, len))
+            mimalloc.mi_heap_malloc(heap, len)
+        else
+            mimalloc.mi_heap_malloc_aligned(heap, len, alignment);
+
+        return @ptrCast([*]u8, ptr orelse null);
+    }
+
+    fn alignedFree(ptr: [*]u8) void {
+        return mimalloc.mi_free(ptr);
+    }
+
+    fn alloc(
+        arena: *anyopaque,
+        len: usize,
+        alignment: u29,
+        len_align: u29,
+        return_address: usize,
+    ) error{OutOfMemory}![]u8 {
+        _ = return_address;
+        assert(len > 0);
+        assert(std.math.isPowerOfTwo(alignment));
+
+        var ptr = alignedAlloc(@ptrCast(*mimalloc.mi_heap_t, arena), len, alignment) orelse return error.OutOfMemory;
+        if (len_align == 0) {
+            return ptr[0..len];
+        }
+        return ptr[0..mem.alignBackwardAnyAlign(mimalloc.mi_usable_size(ptr), len_align)];
+    }
+
+    fn resize(
+        _: *anyopaque,
+        buf: []u8,
+        buf_align: u29,
+        new_len: usize,
+        len_align: u29,
+        return_address: usize,
+    ) ?usize {
+        _ = buf_align;
+        _ = return_address;
+
+        if (new_len <= buf.len) {
+            return mem.alignAllocLen(buf.len, new_len, len_align);
+        }
+
+        const full_len = mimalloc.mi_usable_size(buf.ptr);
+        if (new_len <= full_len) {
+            return mem.alignAllocLen(full_len, new_len, len_align);
+        }
+
+        return null;
+    }
+
+    fn free(
+        _: *anyopaque,
+        buf: []u8,
+        buf_align: u29,
+        return_address: usize,
+    ) void {
+        _ = buf_align;
+        _ = return_address;
+        mimalloc.mi_free(buf.ptr);
+    }
+};
+
+const c_allocator_vtable = Allocator.VTable{
+    .alloc = Arena.alloc,
+    .resize = Arena.resize,
+    .free = Arena.free,
+};