Skip to main content

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}