Add param to use alternate stack and restore old signal handler (#276)

* try keep old_sigaction

Signed-off-by: VGalaxies <[email protected]>

* fixup

Signed-off-by: VGalaxies <[email protected]>

* add comments

Signed-off-by: VGalaxies <[email protected]>

* add on_stack param

Signed-off-by: VGalaxies <[email protected]>

* enable on_stack under frame-pointer feature

Signed-off-by: VGalaxies <[email protected]>

* fix variable does not need to be mutable by shadowing

Signed-off-by: VGalaxies <[email protected]>

---------

Signed-off-by: VGalaxies <[email protected]>

Author: VGalaxies

src/profiler.rs

@@ -32,8 +32,12 @@ pub struct Profiler {
     pub(crate) data: Collector<UnresolvedFrames>,
     sample_counter: i32,
 
+    old_sigaction: Option<signal::SigAction>,
     running: bool,
 
+    #[cfg(feature = "frame-pointer")]
+    on_stack: bool,
+
     #[cfg(any(
         target_arch = "x86_64",
         target_arch = "aarch64",
@@ -46,6 +50,10 @@ pub struct Profiler {
 #[derive(Clone)]
 pub struct ProfilerGuardBuilder {
     frequency: c_int,
+
+    #[cfg(feature = "frame-pointer")]
+    on_stack: bool,
+
     #[cfg(any(
         target_arch = "x86_64",
         target_arch = "aarch64",
@@ -60,6 +68,9 @@ impl Default for ProfilerGuardBuilder {
         ProfilerGuardBuilder {
             frequency: 99,
 
+            #[cfg(feature = "frame-pointer")]
+            on_stack: false,
+
             #[cfg(any(
                 target_arch = "x86_64",
                 target_arch = "aarch64",
@@ -76,6 +87,21 @@ impl ProfilerGuardBuilder {
         Self { frequency, ..self }
     }
 
+    #[cfg(feature = "frame-pointer")]
+    /// Sets whether to use an alternate signal stack via `SA_ONSTACK`.
+    ///
+    /// This is only available and only works correctly when the `frame-pointer` feature is enabled.
+    ///
+    /// The `backtrace-rs` unwinder ignores the signal context and unwinds the current stack. Using
+    /// an alternate stack with it would produce meaningless results. The `frame-pointer` unwinder,
+    /// however, uses the provided `ucontext` to correctly walk the original application stack.
+    ///
+    /// This should be enabled when the profiler is used in an environment
+    /// with small stacks (e.g., inside a Go program) to prevent stack overflow.
+    pub fn on_stack(self, on_stack: bool) -> Self {
+        Self { on_stack, ..self }
+    }
+
     #[cfg(any(
         target_arch = "x86_64",
         target_arch = "aarch64",
@@ -126,6 +152,11 @@ impl ProfilerGuardBuilder {
                 Err(Error::CreatingError)
             }
             Ok(profiler) => {
+                #[cfg(feature = "frame-pointer")]
+                {
+                    profiler.on_stack = self.on_stack;
+                }
+
                 #[cfg(any(
                     target_arch = "x86_64",
                     target_arch = "aarch64",
@@ -387,8 +418,12 @@ impl Profiler {
         Ok(Profiler {
             data: Collector::new()?,
             sample_counter: 0,
+            old_sigaction: None,
             running: false,
 
+            #[cfg(feature = "frame-pointer")]
+            on_stack: false,
+
             #[cfg(any(
                 target_arch = "x86_64",
                 target_arch = "aarch64",
@@ -448,24 +483,30 @@ impl Profiler {
         }
     }
 
-    fn register_signal_handler(&self) -> Result<()> {
+    fn register_signal_handler(&mut self) -> Result<()> {
         let handler = signal::SigHandler::SigAction(perf_signal_handler);
-        let sigaction = signal::SigAction::new(
-            handler,
-            // SA_RESTART will only restart a syscall when it's safe to do so,
-            // e.g. when it's a blocking read(2) or write(2). See man 7 signal.
-            signal::SaFlags::SA_SIGINFO | signal::SaFlags::SA_RESTART,
-            signal::SigSet::empty(),
-        );
-        unsafe { signal::sigaction(signal::SIGPROF, &sigaction) }?;
-
+        // SA_RESTART will only restart a syscall when it's safe to do so,
+        // e.g. when it's a blocking read(2) or write(2). See man 7 signal.
+        let flags = signal::SaFlags::SA_SIGINFO | signal::SaFlags::SA_RESTART;
+        #[cfg(feature = "frame-pointer")]
+        let flags = if self.on_stack {
+            // SA_ONSTACK will deliver the signal on an alternate stack. This is crucial
+            // to prevent a stack overflow if the signal arrives at a thread with
+            // a small stack, which is common when use pprof-rs in Go runtimes.
+            flags | signal::SaFlags::SA_ONSTACK
+        } else {
+            flags
+        };
+        let sigaction = signal::SigAction::new(handler, flags, signal::SigSet::empty());
+        let old_action = unsafe { signal::sigaction(signal::SIGPROF, &sigaction) }?;
+        self.old_sigaction = Some(old_action);
         Ok(())
     }
 
-    fn unregister_signal_handler(&self) -> Result<()> {
-        let handler = signal::SigHandler::SigIgn;
-        unsafe { signal::signal(signal::SIGPROF, handler) }?;
-
+    fn unregister_signal_handler(&mut self) -> Result<()> {
+        if let Some(old_action) = self.old_sigaction.take() {
+            unsafe { signal::sigaction(signal::SIGPROF, &old_action) }?;
+        }
         Ok(())
     }