-
Notifications
You must be signed in to change notification settings - Fork 590
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(frontend): fearless recursion on deep plans (#16279)
Signed-off-by: Bugen Zhao <[email protected]>
- Loading branch information
Showing
6 changed files
with
448 additions
and
216 deletions.
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,190 @@ | ||
| // Copyright 2024 RisingWave Labs | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| //! Track the recursion and grow the stack when necessary to enable fearless recursion. | ||
| use std::cell::RefCell; | ||
|
|
||
| // See documentation of `stacker` for the meaning of these constants. | ||
| // TODO: determine good values or make them configurable | ||
| const RED_ZONE: usize = 128 * 1024; // 128KiB | ||
| const STACK_SIZE: usize = 16 * RED_ZONE; // 2MiB | ||
|
|
||
| /// Recursion depth. | ||
| struct Depth { | ||
| /// The current depth. | ||
| current: usize, | ||
| /// The max depth reached so far, not considering the current depth. | ||
| last_max: usize, | ||
| } | ||
|
|
||
| impl Depth { | ||
| const fn new() -> Self { | ||
| Self { | ||
| current: 0, | ||
| last_max: 0, | ||
| } | ||
| } | ||
|
|
||
| fn reset(&mut self) { | ||
| *self = Self::new(); | ||
| } | ||
| } | ||
|
|
||
| /// The tracker for a recursive function. | ||
| pub struct Tracker { | ||
| depth: RefCell<Depth>, | ||
| } | ||
|
|
||
| impl Tracker { | ||
| /// Create a new tracker. | ||
| pub const fn new() -> Self { | ||
| Self { | ||
| depth: RefCell::new(Depth::new()), | ||
| } | ||
| } | ||
|
|
||
| /// Retrieve the current depth of the recursion. Starts from 1 once the | ||
| /// recursive function is called. | ||
| pub fn depth(&self) -> usize { | ||
| self.depth.borrow().current | ||
| } | ||
|
|
||
| /// Check if the current depth reaches the given depth **for the first time**. | ||
| /// | ||
| /// This is useful for logging without any duplication. | ||
| pub fn depth_reaches(&self, depth: usize) -> bool { | ||
| let d = self.depth.borrow(); | ||
| d.current == depth && d.current > d.last_max | ||
| } | ||
|
|
||
| /// Run a recursive function. Grow the stack if necessary. | ||
| fn recurse<T>(&self, f: impl FnOnce() -> T) -> T { | ||
| struct DepthGuard<'a> { | ||
| depth: &'a RefCell<Depth>, | ||
| } | ||
|
|
||
| impl<'a> DepthGuard<'a> { | ||
| fn new(depth: &'a RefCell<Depth>) -> Self { | ||
| depth.borrow_mut().current += 1; | ||
| Self { depth } | ||
| } | ||
| } | ||
|
|
||
| impl<'a> Drop for DepthGuard<'a> { | ||
| fn drop(&mut self) { | ||
| let mut d = self.depth.borrow_mut(); | ||
| d.last_max = d.last_max.max(d.current); // update the last max depth | ||
| d.current -= 1; // restore the current depth | ||
| if d.current == 0 { | ||
| d.reset(); // reset state if the recursion is finished | ||
| } | ||
| } | ||
| } | ||
|
|
||
| let _guard = DepthGuard::new(&self.depth); | ||
|
|
||
| if cfg!(madsim) { | ||
| f() // madsim does not support stack growth | ||
| } else { | ||
| stacker::maybe_grow(RED_ZONE, STACK_SIZE, f) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// The extension trait for a thread-local tracker to run a recursive function. | ||
| #[easy_ext::ext(Recurse)] | ||
| impl std::thread::LocalKey<Tracker> { | ||
| /// Run the given recursive function. Grow the stack if necessary. | ||
| /// | ||
| /// # Fearless Recursion | ||
| /// | ||
| /// This enables fearless recursion in most cases as long as a single frame | ||
| /// does not exceed the [`RED_ZONE`] size. That is, the caller can recurse | ||
| /// as much as it wants without worrying about stack overflow. | ||
| /// | ||
| /// # Tracker | ||
| /// | ||
| /// The caller can retrieve the [`Tracker`] of the current recursion from | ||
| /// the closure argument. This can be useful for checking the depth of the | ||
| /// recursion, logging or throwing an error gracefully if it's too deep. | ||
| /// | ||
| /// Note that different trackers defined in different functions are | ||
| /// independent of each other. If there's a cross-function recursion, the | ||
| /// tracker retrieved from the closure argument only represents the current | ||
| /// function's state. | ||
| /// | ||
| /// # Example | ||
| /// | ||
| /// Define the tracker with [`tracker!`] and call this method on it to run | ||
| /// a recursive function. | ||
| /// | ||
| /// ```ignore | ||
| /// #[inline(never)] | ||
| /// fn sum(x: u64) -> u64 { | ||
| /// tracker!().recurse(|t| { | ||
| /// if t.depth() % 100000 == 0 { | ||
| /// eprintln!("too deep!"); | ||
| /// } | ||
| /// if x == 0 { | ||
| /// return 0; | ||
| /// } | ||
| /// x + sum(x - 1) | ||
| /// }) | ||
| /// } | ||
| /// ``` | ||
| pub fn recurse<T>(&'static self, f: impl FnOnce(&Tracker) -> T) -> T { | ||
| self.with(|t| t.recurse(|| f(t))) | ||
| } | ||
| } | ||
|
|
||
| /// Define the tracker for recursion and return it. | ||
| /// | ||
| /// Call [`Recurse::recurse`] on it to run a recursive function. See | ||
| /// documentation there for usage. | ||
| #[macro_export] | ||
| macro_rules! __recursive_tracker { | ||
| () => {{ | ||
| use $crate::util::recursive::Tracker; | ||
| std::thread_local! { | ||
| static __TRACKER: Tracker = const { Tracker::new() }; | ||
| } | ||
| __TRACKER | ||
| }}; | ||
| } | ||
| pub use __recursive_tracker as tracker; | ||
|
|
||
| #[cfg(all(test, not(madsim)))] | ||
| mod tests { | ||
| use super::*; | ||
|
|
||
| #[test] | ||
| fn test_fearless_recursion() { | ||
| const X: u64 = 1919810; | ||
| const EXPECTED: u64 = 1842836177955; | ||
|
|
||
| #[inline(never)] | ||
| fn sum(x: u64) -> u64 { | ||
| tracker!().recurse(|t| { | ||
| if x == 0 { | ||
| assert_eq!(t.depth(), X as usize + 1); | ||
| return 0; | ||
| } | ||
| x + sum(x - 1) | ||
| }) | ||
| } | ||
|
|
||
| assert_eq!(sum(X), EXPECTED); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.