let GeneralPurposeAllocator retain metadata to report more double frees

lib/std/heap/general_purpose_allocator.zig

@@ -160,9 +160,16 @@ pub const Config = struct {
 
     /// This is a temporary debugging trick you can use to turn segfaults into more helpful
     /// logged error messages with stack trace details. The downside is that every allocation
-    /// will be leaked!
+    /// will be leaked, unless used with retain_metadata!
     never_unmap: bool = false,
 
+    /// This is a temporary debugging aid that retains metadata about allocations indefinitely.
+    /// This allows a greater range of double frees to be reported. All metadata is freed when
+    /// deinit is called. When used with never_unmap, deliberately leaked memory is also freed
+    /// during deinit. Currently should be used with never_unmap to avoid segfaults.
+    /// TODO https://github.com/ziglang/zig/issues/4298 will allow use without never_unmap
+    retain_metadata: bool = false,
+
     /// Enables emitting info messages with the size and address of every allocation.
     verbose_log: bool = false,
 };
@@ -176,6 +183,8 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
         backing_allocator: *Allocator = std.heap.page_allocator,
         buckets: [small_bucket_count]?*BucketHeader = [1]?*BucketHeader{null} ** small_bucket_count,
         large_allocations: LargeAllocTable = .{},
+        empty_buckets: if (config.retain_metadata) ?*BucketHeader else void =
+            if (config.retain_metadata) null else {},
 
         total_requested_bytes: @TypeOf(total_requested_bytes_init) = total_requested_bytes_init,
         requested_memory_limit: @TypeOf(requested_memory_limit_init) = requested_memory_limit_init,
@@ -205,22 +214,34 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
 
         const LargeAlloc = struct {
             bytes: []u8,
-            stack_addresses: [stack_n]usize,
+            stack_addresses: [trace_n][stack_n]usize,
+            freed: if (config.retain_metadata) bool else void,
+            ptr_align: if (config.never_unmap and config.retain_metadata) u29 else void,
 
-            fn dumpStackTrace(self: *LargeAlloc) void {
-                std.debug.dumpStackTrace(self.getStackTrace());
+            const trace_n = if (config.retain_metadata) traces_per_slot else 1;
+
+            fn dumpStackTrace(self: *LargeAlloc, trace_kind: TraceKind) void {
+                std.debug.dumpStackTrace(self.getStackTrace(trace_kind));
             }
 
-            fn getStackTrace(self: *LargeAlloc) std.builtin.StackTrace {
+            fn getStackTrace(self: *LargeAlloc, trace_kind: TraceKind) std.builtin.StackTrace {
+                assert(@enumToInt(trace_kind) < trace_n);
+                const stack_addresses = &self.stack_addresses[@enumToInt(trace_kind)];
                 var len: usize = 0;
-                while (len < stack_n and self.stack_addresses[len] != 0) {
+                while (len < stack_n and stack_addresses[len] != 0) {
                     len += 1;
                 }
                 return .{
-                    .instruction_addresses = &self.stack_addresses,
+                    .instruction_addresses = stack_addresses,
                     .index = len,
                 };
             }
+
+            fn captureStackTrace(self: *LargeAlloc, ret_addr: usize, trace_kind: TraceKind) void {
+                assert(@enumToInt(trace_kind) < trace_n);
+                const stack_addresses = &self.stack_addresses[@enumToInt(trace_kind)];
+                collectStackTrace(ret_addr, stack_addresses);
+            }
         };
         const LargeAllocTable = std.AutoHashMapUnmanaged(usize, LargeAlloc);
 
@@ -348,16 +369,71 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
             }
             var it = self.large_allocations.valueIterator();
             while (it.next()) |large_alloc| {
+                if (config.retain_metadata and large_alloc.freed) continue;
                 log.err("memory address 0x{x} leaked: {s}", .{
-                    @ptrToInt(large_alloc.bytes.ptr), large_alloc.getStackTrace(),
+                    @ptrToInt(large_alloc.bytes.ptr), large_alloc.getStackTrace(.alloc),
                 });
                 leaks = true;
             }
             return leaks;
         }
 
+        fn freeBucket(self: *Self, bucket: *BucketHeader, size_class: usize) void {
+            const bucket_size = bucketSize(size_class);
+            const bucket_slice = @ptrCast([*]align(@alignOf(BucketHeader)) u8, bucket)[0..bucket_size];
+            self.backing_allocator.free(bucket_slice);
+        }
+
+        fn freeRetainedMetadata(self: *Self) void {
+            if (config.retain_metadata) {
+                if (config.never_unmap) {
+                    // free large allocations that were intentionally leaked by never_unmap
+                    var it = self.large_allocations.iterator();
+                    while (it.next()) |large| {
+                        if (large.value_ptr.freed) {
+                            _ = self.backing_allocator.resizeFn(self.backing_allocator, large.value_ptr.bytes,
+                                large.value_ptr.ptr_align, 0, 0, @returnAddress()) catch unreachable;
+                        }
+                    }
+                }
+                // free retained metadata for small allocations
+                if (self.empty_buckets) |first_bucket| {
+                    var bucket = first_bucket;
+                    while (true) {
+                        const prev = bucket.prev;
+                        if (config.never_unmap) {
+                            // free page that was intentionally leaked by never_unmap
+                            self.backing_allocator.free(bucket.page[0..page_size]);
+                        }
+                        // alloc_cursor was set to slot count when bucket added to empty_buckets
+                        self.freeBucket(bucket, @divExact(page_size, bucket.alloc_cursor));
+                        bucket = prev;
+                        if (bucket == first_bucket)
+                            break;
+                    }
+                    self.empty_buckets = null;
+                }
+            }
+        }
+
+        pub usingnamespace if (config.retain_metadata) struct {
+            pub fn flushRetainedMetadata(self: *Self) void {
+                self.freeRetainedMetadata();
+                // also remove entries from large_allocations
+                var it = self.large_allocations.iterator();
+                while (it.next()) |large| {
+                    if (large.value_ptr.freed) {
+                        _ = self.large_allocations.remove(@ptrToInt(large.value_ptr.bytes.ptr));
+                    }
+                }
+            }
+        } else struct {};
+
         pub fn deinit(self: *Self) bool {
             const leaks = if (config.safety) self.detectLeaks() else false;
+            if (config.retain_metadata) {
+                self.freeRetainedMetadata();
+            }
             self.large_allocations.deinit(self.backing_allocator);
             self.* = undefined;
             return leaks;
@@ -373,6 +449,18 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
             std.debug.captureStackTrace(first_trace_addr, &stack_trace);
         }
 
+        fn reportDoubleFree(ret_addr: usize, alloc_stack_trace: StackTrace, free_stack_trace: StackTrace) void {
+            var addresses: [stack_n]usize = [1]usize{0} ** stack_n;
+            var second_free_stack_trace = StackTrace{
+                .instruction_addresses = &addresses,
+                .index = 0,
+            };
+            std.debug.captureStackTrace(ret_addr, &second_free_stack_trace);
+            log.err("Double free detected. Allocation: {s} First free: {s} Second free: {s}", .{
+                alloc_stack_trace, free_stack_trace, second_free_stack_trace,
+            });
+        }
+
         fn allocSlot(self: *Self, size_class: usize, trace_addr: usize) Error![*]u8 {
             const bucket_index = math.log2(size_class);
             const first_bucket = self.buckets[bucket_index] orelse try self.createBucket(
@@ -408,11 +496,10 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
         }
 
         fn searchBucket(
-            self: *Self,
-            bucket_index: usize,
+            bucket_list: ?*BucketHeader,
             addr: usize,
         ) ?*BucketHeader {
-            const first_bucket = self.buckets[bucket_index] orelse return null;
+            const first_bucket = bucket_list orelse return null;
             var bucket = first_bucket;
             while (true) {
                 const in_bucket_range = (addr >= @ptrToInt(bucket.page) and
@@ -422,7 +509,6 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
                 if (bucket == first_bucket) {
                     return null;
                 }
-                self.buckets[bucket_index] = bucket;
             }
         }
 
@@ -444,6 +530,25 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
                 }
             };
 
+            if (config.retain_metadata and entry.value_ptr.freed) {
+                if (config.safety) {
+                    reportDoubleFree(ret_addr,
+                        entry.value_ptr.getStackTrace(.alloc),
+                        entry.value_ptr.getStackTrace(.free)
+                    );
+                    if (new_size == 0) {
+                        // Recoverable. Restore self.total_requested_bytes if needed.
+                        if (config.enable_memory_limit) {
+                            self.total_requested_bytes += old_mem.len;
+                        }
+                        return @as(usize, 0);
+                    }
+                    @panic("Unrecoverable double free");
+                } else {
+                    unreachable;
+                }
+            }
+
             if (config.safety and old_mem.len != entry.value_ptr.bytes.len) {
                 var addresses: [stack_n]usize = [1]usize{0} ** stack_n;
                 var free_stack_trace = StackTrace{
@@ -454,19 +559,25 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
                 log.err("Allocation size {d} bytes does not match free size {d}. Allocation: {s} Free: {s}", .{
                     entry.value_ptr.bytes.len,
                     old_mem.len,
-                    entry.value_ptr.getStackTrace(),
+                    entry.value_ptr.getStackTrace(.alloc),
                     free_stack_trace,
                 });
             }
 
-            const result_len = try self.backing_allocator.resizeFn(self.backing_allocator, old_mem, old_align, new_size, len_align, ret_addr);
+            const result_len = if (config.never_unmap and new_size == 0) 0 else
+                try self.backing_allocator.resizeFn(self.backing_allocator, old_mem, old_align, new_size, len_align, ret_addr);
 
             if (result_len == 0) {
                 if (config.verbose_log) {
                     log.info("large free {d} bytes at {*}", .{ old_mem.len, old_mem.ptr });
                 }
 
-                assert(self.large_allocations.remove(@ptrToInt(old_mem.ptr)));
+                if (!config.retain_metadata) {
+                    assert(self.large_allocations.remove(@ptrToInt(old_mem.ptr)));
+                } else {
+                    entry.value_ptr.freed = true;
+                    entry.value_ptr.captureStackTrace(ret_addr, .free);
+                }
                 return 0;
             }
 
@@ -476,7 +587,7 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
                 });
             }
             entry.value_ptr.bytes = old_mem.ptr[0..result_len];
-            collectStackTrace(ret_addr, &entry.value_ptr.stack_addresses);
+            entry.value_ptr.captureStackTrace(ret_addr, .alloc);
             return result_len;
         }
 
@@ -520,11 +631,24 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
             var bucket_index = math.log2(size_class_hint);
             var size_class: usize = size_class_hint;
             const bucket = while (bucket_index < small_bucket_count) : (bucket_index += 1) {
-                if (self.searchBucket(bucket_index, @ptrToInt(old_mem.ptr))) |bucket| {
+                if (searchBucket(self.buckets[bucket_index], @ptrToInt(old_mem.ptr))) |bucket| {
+                    // move bucket to head of list to optimize search for nearby allocations
+                    self.buckets[bucket_index] = bucket;
                     break bucket;
                 }
                 size_class *= 2;
-            } else {
+            } else blk: {
+                if (config.retain_metadata) {
+                    if (!self.large_allocations.contains(@ptrToInt(old_mem.ptr))) {
+                        // object not in active buckets or a large allocation, so search empty buckets
+                        if (searchBucket(self.empty_buckets, @ptrToInt(old_mem.ptr))) |bucket| {
+                            // bucket is empty so is_used below will always be false and we exit there
+                            break :blk bucket;
+                        } else {
+                            @panic("Invalid free");
+                        }
+                    }
+                }
                 return self.resizeLarge(old_mem, old_align, new_size, len_align, ret_addr);
             };
             const byte_offset = @ptrToInt(old_mem.ptr) - @ptrToInt(bucket.page);
@@ -535,19 +659,10 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
             const is_used = @truncate(u1, used_byte.* >> used_bit_index) != 0;
             if (!is_used) {
                 if (config.safety) {
-                    const alloc_stack_trace = bucketStackTrace(bucket, size_class, slot_index, .alloc);
-                    const free_stack_trace = bucketStackTrace(bucket, size_class, slot_index, .free);
-                    var addresses: [stack_n]usize = [1]usize{0} ** stack_n;
-                    var second_free_stack_trace = StackTrace{
-                        .instruction_addresses = &addresses,
-                        .index = 0,
-                    };
-                    std.debug.captureStackTrace(ret_addr, &second_free_stack_trace);
-                    log.err("Double free detected. Allocation: {s} First free: {s} Second free: {s}", .{
-                        alloc_stack_trace,
-                        free_stack_trace,
-                        second_free_stack_trace,
-                    });
+                    reportDoubleFree(ret_addr,
+                        bucketStackTrace(bucket, size_class, slot_index, .alloc),
+                        bucketStackTrace(bucket, size_class, slot_index, .free)
+                    );
                     if (new_size == 0) {
                         // Recoverable. Restore self.total_requested_bytes if needed, as we
                         // don't return an error value so the errdefer above does not run.
@@ -579,9 +694,26 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
                     if (!config.never_unmap) {
                         self.backing_allocator.free(bucket.page[0..page_size]);
                     }
-                    const bucket_size = bucketSize(size_class);
-                    const bucket_slice = @ptrCast([*]align(@alignOf(BucketHeader)) u8, bucket)[0..bucket_size];
-                    self.backing_allocator.free(bucket_slice);
+                    if (!config.retain_metadata) {
+                        self.freeBucket(bucket, size_class);
+                    } else {
+                        // move alloc_cursor to end so we can tell size_class later
+                        const slot_count = @divExact(page_size, size_class);
+                        bucket.alloc_cursor = @truncate(SlotIndex, slot_count);
+                        if (self.empty_buckets) |prev_bucket| {
+                            // empty_buckets is ordered newest to oldest through prev so that if
+                            // config.never_unmap is false and backing_allocator reuses freed memory
+                            // then searchBuckets will always return the newer, relevant bucket
+                            bucket.prev = prev_bucket;
+                            bucket.next = prev_bucket.next;
+                            prev_bucket.next = bucket;
+                            bucket.next.prev = bucket;
+                        } else {
+                            bucket.prev = bucket;
+                            bucket.next = bucket;
+                        }
+                        self.empty_buckets = bucket;
+                    }
                 } else {
                     @memset(old_mem.ptr, undefined, old_mem.len);
                 }
@@ -644,9 +776,20 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
                 }
 
                 const gop = self.large_allocations.getOrPutAssumeCapacity(@ptrToInt(slice.ptr));
-                assert(!gop.found_existing); // This would mean the kernel double-mapped pages.
+                if (config.retain_metadata and !config.never_unmap) {
+                    // Backing allocator may be reusing memory that we're retaining metadata for
+                    assert(!gop.found_existing or gop.value_ptr.freed);
+                } else {
+                    assert(!gop.found_existing); // This would mean the kernel double-mapped pages.
+                }
                 gop.value_ptr.bytes = slice;
-                collectStackTrace(ret_addr, &gop.value_ptr.stack_addresses);
+                gop.value_ptr.captureStackTrace(ret_addr, .alloc);
+                if (config.retain_metadata) {
+                    gop.value_ptr.freed = false;
+                    if (config.never_unmap) {
+                        gop.value_ptr.ptr_align = ptr_align;
+                    }
+                }
 
                 if (config.verbose_log) {
                     log.info("large alloc {d} bytes at {*}", .{ slice.len, slice.ptr });
@@ -1015,3 +1158,43 @@ test "setting a memory cap" {
     try std.testing.expect(gpa.total_requested_bytes == 1010);
     allocator.free(exact);
 }
+
+test "double frees" {
+    // use a GPA to back a GPA to check for leaks of the latter's metadata
+    var backing_gpa = GeneralPurposeAllocator(.{ .safety = true }){};
+    defer std.testing.expect(!backing_gpa.deinit()) catch @panic("leak");
+
+    const GPA = GeneralPurposeAllocator(.{ .safety = true, .never_unmap = true, .retain_metadata = true });
+    var gpa = GPA{ .backing_allocator = &backing_gpa.allocator };
+    defer std.testing.expect(!gpa.deinit()) catch @panic("leak");
+    const allocator = &gpa.allocator;
+
+    // detect a small allocation double free, even though bucket is emptied
+    const index: usize = 6;
+    const size_class: usize = @as(usize, 1) << 6;
+    const small = try allocator.alloc(u8, size_class);
+    try std.testing.expect(GPA.searchBucket(gpa.buckets[index], @ptrToInt(small.ptr)) != null);
+    allocator.free(small);
+    try std.testing.expect(GPA.searchBucket(gpa.buckets[index], @ptrToInt(small.ptr)) == null);
+    try std.testing.expect(GPA.searchBucket(gpa.empty_buckets, @ptrToInt(small.ptr)) != null);
+
+    // detect a large allocation double free
+    const large = try allocator.alloc(u8, 2 * page_size);
+    try std.testing.expect(gpa.large_allocations.contains(@ptrToInt(large.ptr)));
+    try std.testing.expectEqual(gpa.large_allocations.getEntry(@ptrToInt(large.ptr)).?.value_ptr.bytes, large);
+    allocator.free(large);
+    try std.testing.expect(gpa.large_allocations.contains(@ptrToInt(large.ptr)));
+    try std.testing.expect(gpa.large_allocations.getEntry(@ptrToInt(large.ptr)).?.value_ptr.freed);
+
+    const normal_small = try allocator.alloc(u8, size_class);
+    defer allocator.free(normal_small);
+    const normal_large = try allocator.alloc(u8, 2 * page_size);
+    defer allocator.free(normal_large);
+
+    // check that flushing retained metadata doesn't disturb live allocations
+    gpa.flushRetainedMetadata();
+    try std.testing.expect(gpa.empty_buckets == null);
+    try std.testing.expect(GPA.searchBucket(gpa.buckets[index], @ptrToInt(normal_small.ptr)) != null);
+    try std.testing.expect(gpa.large_allocations.contains(@ptrToInt(normal_large.ptr)));
+    try std.testing.expect(!gpa.large_allocations.contains(@ptrToInt(large.ptr)));
+}