From 52505124c3912a37a9d81e59805c93963575a5d0 Mon Sep 17 00:00:00 2001
From: Waldemar Kozaczuk <jwkozaczuk@gmail.com>
Date: Tue, 30 Aug 2022 12:11:29 -0400
Subject: [PATCH] lazy stack: ensure next stack page conditionally if
 interrupts and preemption enabled

This patch modifies the last set of call sites where we do not know statically
if any of interrupts or preemption are enabled. In those cases we
dynamically check if both preemtion and interrupts are enabled
and only then pre-fault the stack.

Most of these places are in the tracepoint and sampler implementation.
This is obviously the most costly operation but even here we are lucky
that it affects the performance only when tracepoints are enabled and
even then the code already saves the state of the interrupts using
the arch::irq_flag_notrace so checking if interrupts are enabled is pretty
cheap. With sampler on other hand, the performance is only affected when
starting or stoping the sampler which is quite rare.

Signed-off-by: Waldemar Kozaczuk <jwkozaczuk@gmail.com>
---
 arch/aarch64/interrupt.cc |  5 +++++
 core/sampler.cc           | 31 +++++++++++++++++++++++++++----
 core/trace.cc             | 17 +++++++++++++++++
 include/osv/trace.hh      |  5 +++++
 4 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/arch/aarch64/interrupt.cc b/arch/aarch64/interrupt.cc
index e26e10eebd..0b244e2e0a 100644
--- a/arch/aarch64/interrupt.cc
+++ b/arch/aarch64/interrupt.cc
@@ -31,6 +31,11 @@ void sgi_interrupt::send(sched::cpu* cpu)
 
 void sgi_interrupt::send_allbutself()
 {
+#if CONF_lazy_stack
+    if (sched::preemptable() && arch::irq_enabled()) {
+        arch::ensure_next_stack_page();
+    }
+#endif
     gic::gic->send_sgi(gic::sgi_filter::SGI_TARGET_ALL_BUT_SELF,
                        0, get_id());
 }
diff --git a/core/sampler.cc b/core/sampler.cc
index 9e37ba4fe7..b9e2b05c34 100644
--- a/core/sampler.cc
+++ b/core/sampler.cc
@@ -32,11 +32,16 @@ class cpu_sampler : public sched::timer_base::client {
     sched::timer_base _timer;
     bool _active;
 
-    void rearm()
+    void arm()
     {
         _timer.set(_config.period);
     }
 
+    void rearm()
+    {
+        _timer.set_with_irq_disabled(_config.period);
+    }
+
 public:
     cpu_sampler()
         : _timer(*this)
@@ -54,7 +59,11 @@ class cpu_sampler : public sched::timer_base::client {
     {
         assert(!_active);
         _active = true;
-        rearm();
+        if (arch::irq_enabled()) {
+            arm();
+        } else {
+            rearm();
+        }
     }
 
     void stop()
@@ -97,7 +106,11 @@ static void start_on_current()
 
     if (prev_active + 1 == _n_cpus) {
         _started = true;
-        _controller.wake();
+        if (arch::irq_enabled()) {
+            _controller.wake();
+        } else {
+            _controller.wake_from_kernel_or_with_irq_disabled();
+        }
     }
 }
 
@@ -110,7 +123,11 @@ static void stop_on_current()
     _sampler->stop();
 
     if (--_active_cpus == 0) {
-        _controller.wake();
+        if (arch::irq_enabled()) {
+            _controller.wake();
+        } else {
+            _controller.wake_from_kernel_or_with_irq_disabled();
+        }
     }
 }
 
@@ -170,6 +187,12 @@ void stop_sampler() throw()
 
     WITH_LOCK(migration_lock) {
         stop_sampler_ipi.send_allbutself();
+#if CONF_lazy_stack_invariant
+        assert(arch::irq_enabled());
+#endif
+#if CONF_lazy_stack
+        sched::ensure_next_stack_page_if_preemptable();
+#endif
         stop_on_current();
     }
 
diff --git a/core/trace.cc b/core/trace.cc
index dc69c80747..c9ed2babbc 100644
--- a/core/trace.cc
+++ b/core/trace.cc
@@ -247,6 +247,13 @@ void tracepoint_base::update()
     WITH_LOCK(trace_control_lock) {
         bool empty;
 
+#if CONF_lazy_stack_invariant
+        assert(arch::irq_enabled());
+        assert(sched::preemptable());
+#endif
+#if CONF_lazy_stack
+        arch::ensure_next_stack_page();
+#endif
         WITH_LOCK(osv::rcu_read_lock) {
             auto& probes = *probes_ptr.read();
 
@@ -376,6 +383,11 @@ extern "C" void __cyg_profile_func_enter(void *this_fn, void *call_site)
     }
     arch::irq_flag_notrace irq;
     irq.save();
+#if CONF_lazy_stack
+    if (sched::preemptable() && irq.enabled()) {
+        arch::ensure_next_stack_page();
+    }
+#endif
     arch::irq_disable_notrace();
     if (func_trace_nesting++ == 0) {
         trace_function_entry(this_fn, call_site);
@@ -391,6 +403,11 @@ extern "C" void __cyg_profile_func_exit(void *this_fn, void *call_site)
     }
     arch::irq_flag_notrace irq;
     irq.save();
+#if CONF_lazy_stack
+    if (sched::preemptable() && irq.enabled()) {
+        arch::ensure_next_stack_page();
+    }
+#endif
     arch::irq_disable_notrace();
     if (func_trace_nesting++ == 0) {
         trace_function_exit(this_fn, call_site);
diff --git a/include/osv/trace.hh b/include/osv/trace.hh
index d735575cb8..01d7202240 100644
--- a/include/osv/trace.hh
+++ b/include/osv/trace.hh
@@ -348,6 +348,11 @@ public:
         if (active) {
             arch::irq_flag_notrace irq;
             irq.save();
+#if CONF_lazy_stack
+            if (sched::preemptable() && irq.enabled()) {
+                arch::ensure_next_stack_page();
+            }
+#endif
             arch::irq_disable_notrace();
             log(as);
             run_probes();