authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-03-06 14:39:09-05:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2025-03-06 14:39:09-05:00
log4cefd1bd1bac7059f7e00162a579f659e5f59412
tree8daffcd79bf2661338b966cdfb136f37c2925ed0
parent1e0739f0c6c7d731771ea3df140a1e5462f27a7e
parent5b03e248b7c88e7ac0266ba33fd25ca8b4cd7f15
signaturebadge-check Signed by PGP key B5690EEEBB952194

Merge pull request #23097 from ziggoon/master

std.heap.PageAllocator updates to fix race condition and utilize NtAllocateVirtualMemory / NtFreeVirtualMemory instead of VirtualAlloc / VirtualFree

3 files changed, 116 insertions(+), 62 deletions(-)

lib/std/heap/PageAllocator.zig+65-62
...@@ -6,9 +6,14 @@ const maxInt = std.math.maxInt;...@@ -6,9 +6,14 @@ const maxInt = std.math.maxInt;
6const assert = std.debug.assert;6const assert = std.debug.assert;
7const native_os = builtin.os.tag;7const native_os = builtin.os.tag;
8const windows = std.os.windows;8const windows = std.os.windows;
9const ntdll = windows.ntdll;
9const posix = std.posix;10const posix = std.posix;
10const page_size_min = std.heap.page_size_min;11const page_size_min = std.heap.page_size_min;
1112
13const SUCCESS = @import("../os/windows/ntstatus.zig").NTSTATUS.SUCCESS;
14const MEM_RESERVE_PLACEHOLDER = windows.MEM_RESERVE_PLACEHOLDER;
15const MEM_PRESERVE_PLACEHOLDER = windows.MEM_PRESERVE_PLACEHOLDER;
16
12pub const vtable: Allocator.VTable = .{17pub const vtable: Allocator.VTable = .{
13 .alloc = alloc,18 .alloc = alloc,
14 .resize = resize,19 .resize = resize,
...@@ -22,51 +27,62 @@ pub fn map(n: usize, alignment: mem.Alignment) ?[*]u8 {...@@ -22,51 +27,62 @@ pub fn map(n: usize, alignment: mem.Alignment) ?[*]u8 {
22 const alignment_bytes = alignment.toByteUnits();27 const alignment_bytes = alignment.toByteUnits();
2328
24 if (native_os == .windows) {29 if (native_os == .windows) {
25 // According to official documentation, VirtualAlloc aligns to page30 var base_addr: ?*anyopaque = null;
26 // boundary, however, empirically it reserves pages on a 64K boundary.31 var size: windows.SIZE_T = n;
27 // Since it is very likely the requested alignment will be honored,32
28 // this logic first tries a call with exactly the size requested,33 var status = ntdll.NtAllocateVirtualMemory(windows.GetCurrentProcess(), @ptrCast(&base_addr), 0, &size, windows.MEM_COMMIT | windows.MEM_RESERVE, windows.PAGE_READWRITE);
29 // before falling back to the loop below.34
30 // https://devblogs.microsoft.com/oldnewthing/?p=4222335 if (status == SUCCESS and mem.isAligned(@intFromPtr(base_addr), alignment_bytes)) {
31 const addr = windows.VirtualAlloc(36 return @ptrCast(base_addr);
32 null,37 }
33 // VirtualAlloc will round the length to a multiple of page size.38
34 // "If the lpAddress parameter is NULL, this value is rounded up to39 if (status == SUCCESS) {
35 // the next page boundary".40 var region_size: windows.SIZE_T = 0;
36 n,41 _ = ntdll.NtFreeVirtualMemory(windows.GetCurrentProcess(), @ptrCast(&base_addr), &region_size, windows.MEM_RELEASE);
37 windows.MEM_COMMIT | windows.MEM_RESERVE,42 }
38 windows.PAGE_READWRITE,
39 ) catch return null;
40
41 if (mem.isAligned(@intFromPtr(addr), alignment_bytes))
42 return @ptrCast(addr);
43
44 // Fallback: reserve a range of memory large enough to find a
45 // sufficiently aligned address, then free the entire range and
46 // immediately allocate the desired subset. Another thread may have won
47 // the race to map the target range, in which case a retry is needed.
48 windows.VirtualFree(addr, 0, windows.MEM_RELEASE);
4943
50 const overalloc_len = n + alignment_bytes - page_size;44 const overalloc_len = n + alignment_bytes - page_size;
51 const aligned_len = mem.alignForward(usize, n, page_size);45 const aligned_len = mem.alignForward(usize, n, page_size);
5246
53 while (true) {47 base_addr = null;
54 const reserved_addr = windows.VirtualAlloc(48 size = overalloc_len;
55 null,49
56 overalloc_len,50 status = ntdll.NtAllocateVirtualMemory(windows.GetCurrentProcess(), @ptrCast(&base_addr), 0, &size, windows.MEM_RESERVE | MEM_RESERVE_PLACEHOLDER, windows.PAGE_NOACCESS);
57 windows.MEM_RESERVE,51
58 windows.PAGE_NOACCESS,52 if (status != SUCCESS) return null;
59 ) catch return null;53
60 const aligned_addr = mem.alignForward(usize, @intFromPtr(reserved_addr), alignment_bytes);54 const placeholder_addr = @intFromPtr(base_addr);
61 windows.VirtualFree(reserved_addr, 0, windows.MEM_RELEASE);55 const aligned_addr = mem.alignForward(usize, placeholder_addr, alignment_bytes);
62 const ptr = windows.VirtualAlloc(56 const prefix_size = aligned_addr - placeholder_addr;
63 @ptrFromInt(aligned_addr),57
64 aligned_len,58 if (prefix_size > 0) {
65 windows.MEM_COMMIT | windows.MEM_RESERVE,59 var prefix_base = base_addr;
66 windows.PAGE_READWRITE,60 var prefix_size_param: windows.SIZE_T = prefix_size;
67 ) catch continue;61 _ = ntdll.NtFreeVirtualMemory(windows.GetCurrentProcess(), @ptrCast(&prefix_base), &prefix_size_param, windows.MEM_RELEASE | MEM_PRESERVE_PLACEHOLDER);
68 return @ptrCast(ptr);
69 }62 }
63
64 const suffix_start = aligned_addr + aligned_len;
65 const suffix_size = (placeholder_addr + overalloc_len) - suffix_start;
66 if (suffix_size > 0) {
67 var suffix_base = @as(?*anyopaque, @ptrFromInt(suffix_start));
68 var suffix_size_param: windows.SIZE_T = suffix_size;
69 _ = ntdll.NtFreeVirtualMemory(windows.GetCurrentProcess(), @ptrCast(&suffix_base), &suffix_size_param, windows.MEM_RELEASE | MEM_PRESERVE_PLACEHOLDER);
70 }
71
72 base_addr = @ptrFromInt(aligned_addr);
73 size = aligned_len;
74
75 status = ntdll.NtAllocateVirtualMemory(windows.GetCurrentProcess(), @ptrCast(&base_addr), 0, &size, windows.MEM_COMMIT | MEM_PRESERVE_PLACEHOLDER, windows.PAGE_READWRITE);
76
77 if (status == SUCCESS) {
78 return @ptrCast(base_addr);
79 }
80
81 base_addr = @as(?*anyopaque, @ptrFromInt(aligned_addr));
82 size = aligned_len;
83 _ = ntdll.NtFreeVirtualMemory(windows.GetCurrentProcess(), @ptrCast(&base_addr), &size, windows.MEM_RELEASE);
84
85 return null;
70 }86 }
7187
72 const aligned_len = mem.alignForward(usize, n, page_size);88 const aligned_len = mem.alignForward(usize, n, page_size);
...@@ -104,26 +120,14 @@ fn alloc(context: *anyopaque, n: usize, alignment: mem.Alignment, ra: usize) ?[*...@@ -104,26 +120,14 @@ fn alloc(context: *anyopaque, n: usize, alignment: mem.Alignment, ra: usize) ?[*
104 return map(n, alignment);120 return map(n, alignment);
105}121}
106122
107fn resize(123fn resize(context: *anyopaque, memory: []u8, alignment: mem.Alignment, new_len: usize, return_address: usize) bool {
108 context: *anyopaque,
109 memory: []u8,
110 alignment: mem.Alignment,
111 new_len: usize,
112 return_address: usize,
113) bool {
114 _ = context;124 _ = context;
115 _ = alignment;125 _ = alignment;
116 _ = return_address;126 _ = return_address;
117 return realloc(memory, new_len, false) != null;127 return realloc(memory, new_len, false) != null;
118}128}
119129
120fn remap(130fn remap(context: *anyopaque, memory: []u8, alignment: mem.Alignment, new_len: usize, return_address: usize) ?[*]u8 {
121 context: *anyopaque,
122 memory: []u8,
123 alignment: mem.Alignment,
124 new_len: usize,
125 return_address: usize,
126) ?[*]u8 {
127 _ = context;131 _ = context;
128 _ = alignment;132 _ = alignment;
129 _ = return_address;133 _ = return_address;
...@@ -139,7 +143,9 @@ fn free(context: *anyopaque, memory: []u8, alignment: mem.Alignment, return_addr...@@ -139,7 +143,9 @@ fn free(context: *anyopaque, memory: []u8, alignment: mem.Alignment, return_addr
139143
140pub fn unmap(memory: []align(page_size_min) u8) void {144pub fn unmap(memory: []align(page_size_min) u8) void {
141 if (native_os == .windows) {145 if (native_os == .windows) {
142 windows.VirtualFree(memory.ptr, 0, windows.MEM_RELEASE);146 var base_addr: ?*anyopaque = memory.ptr;
147 var region_size: windows.SIZE_T = 0;
148 _ = ntdll.NtFreeVirtualMemory(windows.GetCurrentProcess(), @ptrCast(&base_addr), &region_size, windows.MEM_RELEASE);
143 } else {149 } else {
144 const page_aligned_len = mem.alignForward(usize, memory.len, std.heap.pageSize());150 const page_aligned_len = mem.alignForward(usize, memory.len, std.heap.pageSize());
145 posix.munmap(memory.ptr[0..page_aligned_len]);151 posix.munmap(memory.ptr[0..page_aligned_len]);
...@@ -157,13 +163,10 @@ pub fn realloc(uncasted_memory: []u8, new_len: usize, may_move: bool) ?[*]u8 {...@@ -157,13 +163,10 @@ pub fn realloc(uncasted_memory: []u8, new_len: usize, may_move: bool) ?[*]u8 {
157 const old_addr_end = base_addr + memory.len;163 const old_addr_end = base_addr + memory.len;
158 const new_addr_end = mem.alignForward(usize, base_addr + new_len, page_size);164 const new_addr_end = mem.alignForward(usize, base_addr + new_len, page_size);
159 if (old_addr_end > new_addr_end) {165 if (old_addr_end > new_addr_end) {
160 // For shrinking that is not releasing, we will only decommit166 var decommit_addr: ?*anyopaque = @ptrFromInt(new_addr_end);
161 // the pages not needed anymore.167 var decommit_size: windows.SIZE_T = old_addr_end - new_addr_end;
162 windows.VirtualFree(168
163 @ptrFromInt(new_addr_end),169 _ = ntdll.NtAllocateVirtualMemory(windows.GetCurrentProcess(), @ptrCast(&decommit_addr), 0, &decommit_size, windows.MEM_RESET, windows.PAGE_NOACCESS);
164 old_addr_end - new_addr_end,
165 windows.MEM_DECOMMIT,
166 );
167 }170 }
168 return memory.ptr;171 return memory.ptr;
169 }172 }
lib/std/os/windows.zig+34
...@@ -1758,6 +1758,38 @@ pub fn TerminateProcess(hProcess: HANDLE, uExitCode: UINT) TerminateProcessError...@@ -1758,6 +1758,38 @@ pub fn TerminateProcess(hProcess: HANDLE, uExitCode: UINT) TerminateProcessError
1758 }1758 }
1759}1759}
17601760
1761pub const NtAllocateVirtualMemoryError = error{
1762 AccessDenied,
1763 InvalidParameter,
1764 NoMemory,
1765 Unexpected,
1766};
1767
1768pub fn NtAllocateVirtualMemory(hProcess: HANDLE, addr: ?*PVOID, zero_bits: ULONG_PTR, size: ?*SIZE_T, alloc_type: ULONG, protect: ULONG) NtAllocateVirtualMemoryError!void {
1769 return switch (ntdll.NtAllocateVirtualMemory(hProcess, addr, zero_bits, size, alloc_type, protect)) {
1770 .SUCCESS => return,
1771 .ACCESS_DENIED => NtAllocateVirtualMemoryError.AccessDenied,
1772 .INVALID_PARAMETER => NtAllocateVirtualMemoryError.InvalidParameter,
1773 .NO_MEMORY => NtAllocateVirtualMemoryError.NoMemory,
1774 else => |st| unexpectedStatus(st),
1775 };
1776}
1777
1778pub const NtFreeVirtualMemoryError = error{
1779 AccessDenied,
1780 InvalidParameter,
1781 Unexpected,
1782};
1783
1784pub fn NtFreeVirtualMemory(hProcess: HANDLE, addr: ?*PVOID, size: *SIZE_T, free_type: ULONG) NtFreeVirtualMemoryError!void {
1785 return switch (ntdll.NtFreeVirtualMemory(hProcess, addr, size, free_type)) {
1786 .SUCCESS => return,
1787 .ACCESS_DENIED => NtFreeVirtualMemoryError.AccessDenied,
1788 .INVALID_PARAMETER => NtFreeVirtualMemoryError.InvalidParameter,
1789 else => NtFreeVirtualMemoryError.Unexpected,
1790 };
1791}
1792
1761pub const VirtualAllocError = error{Unexpected};1793pub const VirtualAllocError = error{Unexpected};
17621794
1763pub fn VirtualAlloc(addr: ?LPVOID, size: usize, alloc_type: DWORD, flProtect: DWORD) VirtualAllocError!LPVOID {1795pub fn VirtualAlloc(addr: ?LPVOID, size: usize, alloc_type: DWORD, flProtect: DWORD) VirtualAllocError!LPVOID {
...@@ -3539,6 +3571,8 @@ pub const MEM_LARGE_PAGES = 0x20000000;...@@ -3539,6 +3571,8 @@ pub const MEM_LARGE_PAGES = 0x20000000;
3539pub const MEM_PHYSICAL = 0x400000;3571pub const MEM_PHYSICAL = 0x400000;
3540pub const MEM_TOP_DOWN = 0x100000;3572pub const MEM_TOP_DOWN = 0x100000;
3541pub const MEM_WRITE_WATCH = 0x200000;3573pub const MEM_WRITE_WATCH = 0x200000;
3574pub const MEM_RESERVE_PLACEHOLDER = 0x00040000;
3575pub const MEM_PRESERVE_PLACEHOLDER = 0x00000400;
35423576
3543// Protect values3577// Protect values
3544pub const PAGE_EXECUTE = 0x10;3578pub const PAGE_EXECUTE = 0x10;
lib/std/os/windows/ntdll.zig+17
...@@ -5,6 +5,7 @@ const BOOL = windows.BOOL;...@@ -5,6 +5,7 @@ const BOOL = windows.BOOL;
5const DWORD = windows.DWORD;5const DWORD = windows.DWORD;
6const DWORD64 = windows.DWORD64;6const DWORD64 = windows.DWORD64;
7const ULONG = windows.ULONG;7const ULONG = windows.ULONG;
8const ULONG_PTR = windows.ULONG_PTR;
8const NTSTATUS = windows.NTSTATUS;9const NTSTATUS = windows.NTSTATUS;
9const WORD = windows.WORD;10const WORD = windows.WORD;
10const HANDLE = windows.HANDLE;11const HANDLE = windows.HANDLE;
...@@ -358,3 +359,19 @@ pub extern "ntdll" fn NtCreateNamedPipeFile(...@@ -358,3 +359,19 @@ pub extern "ntdll" fn NtCreateNamedPipeFile(
358 OutboundQuota: ULONG,359 OutboundQuota: ULONG,
359 DefaultTimeout: *LARGE_INTEGER,360 DefaultTimeout: *LARGE_INTEGER,
360) callconv(.winapi) NTSTATUS;361) callconv(.winapi) NTSTATUS;
362
363pub extern "ntdll" fn NtAllocateVirtualMemory(
364 ProcessHandle: HANDLE,
365 BaseAddress: ?*PVOID,
366 ZeroBits: ULONG_PTR,
367 RegionSize: ?*SIZE_T,
368 AllocationType: ULONG,
369 PageProtection: ULONG,
370) callconv(.winapi) NTSTATUS;
371
372pub extern "ntdll" fn NtFreeVirtualMemory(
373 ProcessHandle: HANDLE,
374 BaseAddress: ?*PVOID,
375 RegionSize: *SIZE_T,
376 FreeType: ULONG,
377) callconv(.winapi) NTSTATUS;