diff --git a/lib/std/heap/general_purpose_allocator.zig b/lib/std/heap/general_purpose_allocator.zig index 37df19e823..86f07840d8 100644 --- a/lib/std/heap/general_purpose_allocator.zig +++ b/lib/std/heap/general_purpose_allocator.zig @@ -210,6 +210,7 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type { const LargeAlloc = struct { bytes: []u8, + requested_size: if (config.enable_memory_limit) usize else void, 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, @@ -528,13 +529,9 @@ 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; - } + // Recoverable if this is a free. + if (new_size == 0) return @as(usize, 0); - } @panic("Unrecoverable double free"); } else { unreachable; @@ -556,7 +553,29 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type { }); } - 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); + // Do memory limit accounting with requested sizes rather than what backing_allocator returns + // because if we want to return error.OutOfMemory, we have to leave allocation untouched, and + // that is impossible to guarantee after calling backing_allocator.resizeFn. + const prev_req_bytes = self.total_requested_bytes; + if (config.enable_memory_limit) { + const new_req_bytes = prev_req_bytes + new_size - entry.value_ptr.requested_size; + if (new_req_bytes > prev_req_bytes and new_req_bytes > self.requested_memory_limit) { + return error.OutOfMemory; + } + self.total_requested_bytes = new_req_bytes; + } + errdefer if (config.enable_memory_limit) { + self.total_requested_bytes = prev_req_bytes; + }; + + 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 (config.enable_memory_limit) { + entry.value_ptr.requested_size = new_size; + } if (result_len == 0) { if (config.verbose_log) { @@ -599,18 +618,6 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type { const held = self.mutex.acquire(); defer held.release(); - const prev_req_bytes = self.total_requested_bytes; - if (config.enable_memory_limit) { - const new_req_bytes = prev_req_bytes + new_size - old_mem.len; - if (new_req_bytes > prev_req_bytes and new_req_bytes > self.requested_memory_limit) { - return error.OutOfMemory; - } - self.total_requested_bytes = new_req_bytes; - } - errdefer if (config.enable_memory_limit) { - self.total_requested_bytes = prev_req_bytes; - }; - assert(old_mem.len != 0); const aligned_size = math.max(old_mem.len, old_align); @@ -651,19 +658,28 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type { if (!is_used) { if (config.safety) { 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. - if (config.enable_memory_limit) { - self.total_requested_bytes = prev_req_bytes; - } + // Recoverable if this is a free. + if (new_size == 0) return @as(usize, 0); - } @panic("Unrecoverable double free"); } else { unreachable; } } + + // Definitely an in-use small alloc now. + const prev_req_bytes = self.total_requested_bytes; + if (config.enable_memory_limit) { + const new_req_bytes = prev_req_bytes + new_size - old_mem.len; + if (new_req_bytes > prev_req_bytes and new_req_bytes > self.requested_memory_limit) { + return error.OutOfMemory; + } + self.total_requested_bytes = new_req_bytes; + } + errdefer if (config.enable_memory_limit) { + self.total_requested_bytes = prev_req_bytes; + }; + if (new_size == 0) { // Capture stack trace to be the "first free", in case a double free happens. bucket.captureStackTrace(ret_addr, size_class, slot_index, .free); @@ -745,21 +761,15 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type { const held = self.mutex.acquire(); defer held.release(); + if (!self.isAllocationAllowed(len)) { + return error.OutOfMemory; + } + const new_aligned_size = math.max(len, ptr_align); if (new_aligned_size > largest_bucket_object_size) { try self.large_allocations.ensureUnusedCapacity(self.backing_allocator, 1); - const slice = try self.backing_allocator.allocFn(self.backing_allocator, len, ptr_align, len_align, ret_addr); - // The backing allocator may return a memory block bigger than - // `len`, use the effective size for bookkeeping purposes - if (!self.isAllocationAllowed(slice.len)) { - // Free the block so no memory is leaked - const new_len = try self.backing_allocator.resizeFn(self.backing_allocator, slice, ptr_align, 0, 0, ret_addr); - assert(new_len == 0); - return error.OutOfMemory; - } - const gop = self.large_allocations.getOrPutAssumeCapacity(@ptrToInt(slice.ptr)); if (config.retain_metadata and !config.never_unmap) { // Backing allocator may be reusing memory that we're retaining metadata for @@ -768,6 +778,8 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type { assert(!gop.found_existing); // This would mean the kernel double-mapped pages. } gop.value_ptr.bytes = slice; + if (config.enable_memory_limit) + gop.value_ptr.requested_size = len; gop.value_ptr.captureStackTrace(ret_addr, .alloc); if (config.retain_metadata) { gop.value_ptr.freed = false; @@ -782,10 +794,6 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type { return slice; } - if (!self.isAllocationAllowed(len)) { - return error.OutOfMemory; - } - const new_size_class = math.ceilPowerOfTwoAssert(usize, new_aligned_size); const ptr = try self.allocSlot(new_size_class, ret_addr); if (config.verbose_log) { @@ -1183,3 +1191,14 @@ test "double frees" { try std.testing.expect(gpa.large_allocations.contains(@ptrToInt(normal_large.ptr))); try std.testing.expect(!gpa.large_allocations.contains(@ptrToInt(large.ptr))); } + +test "bug 9995 fix, large allocs count requested size not backing size" { + // with AtLeast, buffer likely to be larger than requested, especially when shrinking + var gpa = GeneralPurposeAllocator(.{ .enable_memory_limit = true }){}; + var buf = try gpa.allocator.allocAdvanced(u8, 1, page_size + 1, .at_least); + try std.testing.expect(gpa.total_requested_bytes == page_size + 1); + buf = try gpa.allocator.reallocAtLeast(buf, 1); + try std.testing.expect(gpa.total_requested_bytes == 1); + buf = try gpa.allocator.reallocAtLeast(buf, 2); + try std.testing.expect(gpa.total_requested_bytes == 2); +}