rocketmq_runtime/stack.rs
1// Copyright 2026 The RocketMQ Rust Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15// Normalize before adding lifecycle/select/timeout wrappers. This bounds
16// inline state, not arbitrary user poll frames or recursive user destructors.
17//
18// Unoptimized builds give every by-value move its own stack slot, so each
19// wrapper and forwarding call below this boundary keeps another copy of an
20// inline future and submission can take many times the future's size. Such
21// builds use Tokio's debug boxing boundary instead.
22pub(crate) const MAX_INLINE_SIZE: usize = if cfg!(debug_assertions) { 2 * 1024 } else { 16 * 1024 };
23
24/// Returns how many bytes of the calling thread's stack are committed.
25///
26/// Windows commits a thread stack one page at a time as it grows and keeps
27/// those pages until the thread exits, so the value is the deepest stack use
28/// of the thread so far, rounded up to whole pages. A process entrypoint can
29/// read it just before returning to measure the headroom left on the main
30/// thread, whose size the runtime cannot choose.
31///
32/// Returns `None` on other platforms, or if the stack layout cannot be read.
33pub fn current_thread_committed_stack_bytes() -> Option<usize> {
34 #[cfg(windows)]
35 {
36 win::committed_stack_bytes()
37 }
38 #[cfg(not(windows))]
39 {
40 None
41 }
42}
43
44#[cfg(windows)]
45mod win {
46 use std::ffi::c_void;
47 use std::mem::size_of;
48
49 use windows_sys::Win32::System::Memory::VirtualQuery;
50 use windows_sys::Win32::System::Memory::MEMORY_BASIC_INFORMATION;
51 use windows_sys::Win32::System::Memory::MEM_COMMIT;
52 use windows_sys::Win32::System::Memory::PAGE_GUARD;
53 use windows_sys::Win32::System::Threading::GetCurrentThreadStackLimits;
54
55 pub(super) fn committed_stack_bytes() -> Option<usize> {
56 let mut low = 0_usize;
57 let mut high = 0_usize;
58 // SAFETY: both arguments point to live, writable locals.
59 unsafe { GetCurrentThreadStackLimits(&mut low, &mut high) };
60
61 // The reservation runs from `low` (never committed) through the guard
62 // pages to the committed pages that end at `high`. Walk it upward and
63 // stop at the first committed page that is not a guard page.
64 let mut address = low;
65 while address < high {
66 let mut region = MEMORY_BASIC_INFORMATION::default();
67 // SAFETY: `region` is a live, writable buffer of the length passed.
68 // VirtualQuery reads the page attributes only, never the page
69 // contents, so it cannot trip a guard page.
70 let written = unsafe {
71 VirtualQuery(
72 address as *const c_void,
73 &mut region,
74 size_of::<MEMORY_BASIC_INFORMATION>(),
75 )
76 };
77 if written == 0 || region.RegionSize == 0 {
78 return None;
79 }
80 let base = region.BaseAddress as usize;
81 if region.State == MEM_COMMIT && region.Protect & PAGE_GUARD == 0 {
82 return high.checked_sub(base);
83 }
84 address = base.checked_add(region.RegionSize)?;
85 }
86 None
87 }
88}
89
90#[cfg(test)]
91mod tests {
92 use super::current_thread_committed_stack_bytes;
93
94 #[test]
95 #[cfg(windows)]
96 fn committed_stack_grows_with_stack_use() {
97 let (before, after) = std::thread::Builder::new()
98 .stack_size(1024 * 1024)
99 .spawn(|| {
100 let before = current_thread_committed_stack_bytes().expect("committed stack before use");
101 (before, use_stack(256 * 1024))
102 })
103 .expect("spawn stack test thread")
104 .join()
105 .expect("stack test thread");
106 assert!(before > 0);
107 assert!(
108 after >= before + 200 * 1024,
109 "committed stack {after} did not grow past {before} after using 256 KiB"
110 );
111 }
112
113 #[cfg(windows)]
114 #[inline(never)]
115 fn use_stack(bytes: usize) -> usize {
116 const CHUNK: usize = 16 * 1024;
117 let mut frame = [0_u8; CHUNK];
118 std::hint::black_box(&mut frame);
119 let committed = if bytes > CHUNK {
120 use_stack(bytes - CHUNK)
121 } else {
122 current_thread_committed_stack_bytes().expect("committed stack after use")
123 };
124 // Keep this frame live across the call so it cannot become a tail call.
125 std::hint::black_box(&frame);
126 committed
127 }
128
129 #[test]
130 #[cfg(not(windows))]
131 fn other_platforms_do_not_report_a_committed_stack() {
132 assert_eq!(current_thread_committed_stack_bytes(), None);
133 }
134}