lazy stack: ensure next stack page dynamically if preemption enabled

This patch modifies number of critical places mostly in the scheduler
code to dynamically pre-fault the stack if preemption is enabled. In all
these places we can be statically sure that interrupts are enabled
but not so sure about preemption (maybe in future we can prove it is enabled
at least in some of the cases and replace conditional
ensure_next_stack_page_if_preemptable() with cheaper single-intruction
ensure_next_stack_page()).

The three call sites before irq_lock is taken (interrupts are disabled)
include:
- cpu::schedule() before WITH_LOCK(irq_lock)
- thread::pin() before WITH_LOCK(irq_lock)
- thread::yield() before guard*(irq_lock)

The reasoning above goes like this: the methods were designed
with intention to be used when interrupts are enabled because otherwise
the act of using WITH_LOCK or guard construct would imply that interrupts
would be re-enabled after the block so the code does not care about restoring
interrupts to the proper state it was before. If that was the intension,
these methods would use irq_save_lock_type.

Two other call sites in the scheduler are:
- timer_base destructor before calling timer_base::cancel() which disables interrupts
- thread::wake_with_from_mutex() only called from waiter::wake()
follow the reasoning that interrupts are enabled most of the time
except few places in scheduler and interrupt handler. Also we assume the mutex
is not used when interrupts are disabled. And based on my analysis of all the code
that disables/enables interrupts that seems to be the case.
In any case we put invariants in these two places to verify that interrupts
are enabled indeed.

The last call site in abort() assumes that calling irq_disable() implies
that interrupts are enabled before and we only need to check preemption status.

Signed-off-by: Waldemar Kozaczuk <jwkozaczuk@gmail.com>

core/sched.cc

@@ -224,6 +224,9 @@ void thread::cputime_estimator_get(
 // scheduler on a different CPU would be disastrous.
 void cpu::schedule()
 {
+#if CONF_lazy_stack
+    sched::ensure_next_stack_page_if_preemptable();
+#endif
     WITH_LOCK(irq_lock) {
 #ifdef __aarch64__
         reschedule_from_interrupt(sched::cpu::current(), false, thyst);
@@ -566,6 +569,9 @@ void thread::pin(cpu *target_cpu)
         t.wake();
     }, sched::thread::attr().pin(source_cpu)));
     wakeme->start();
+#if CONF_lazy_stack
+    sched::ensure_next_stack_page_if_preemptable();
+#endif
     WITH_LOCK(irq_lock) {
         trace_sched_migrate(&t, target_cpu->id);
         t.stat_migrations.incr();
@@ -822,6 +828,12 @@ void thread::yield(thread_runtime::duration preempt_after)
 {
     trace_sched_yield();
     auto t = current();
+#if CONF_lazy_stack_invariant
+    assert(arch::irq_enabled());
+#endif
+#if CONF_lazy_stack
+    sched::ensure_next_stack_page_if_preemptable();
+#endif
     std::lock_guard<irq_lock_type> guard(irq_lock);
     // FIXME: drive by IPI
     cpu::current()->handle_incoming_wakeups();
@@ -1258,6 +1270,12 @@ void thread::wake_impl(detached_state* st, unsigned allowed_initial_states_mask)
 
 void thread::wake()
 {
+#if CONF_lazy_stack_invariant
+    assert(arch::irq_enabled());
+#endif
+#if CONF_lazy_stack
+    sched::ensure_next_stack_page_if_preemptable();
+#endif
     WITH_LOCK(rcu_read_lock) {
         wake_impl(_detached_state.get());
     }
@@ -1604,6 +1622,12 @@ timer_base::timer_base(timer_base::client& t)
 
 timer_base::~timer_base()
 {
+#if CONF_lazy_stack_invariant
+    assert(arch::irq_enabled());
+#endif
+#if CONF_lazy_stack
+    sched::ensure_next_stack_page_if_preemptable();
+#endif
     cancel();
 }
 

include/osv/sched.hh

@@ -1043,6 +1043,15 @@ inline bool preemptable()
     return !get_preempt_counter();
 }
 
+#if CONF_lazy_stack
+inline void ensure_next_stack_page_if_preemptable() {
+    if (!preemptable()) {
+        return;
+    }
+    arch::ensure_next_stack_page();
+}
+#endif
+
 inline void preempt()
 {
     if (preemptable()) {
@@ -1350,6 +1359,14 @@ template <class Action>
 inline
 void thread::wake_with_from_mutex(Action action)
 {
+#if CONF_lazy_stack_invariant
+    assert(arch::irq_enabled());
+#endif
+#if CONF_lazy_stack
+    if (preemptable()) {
+        arch::ensure_next_stack_page();
+    }
+#endif
     return do_wake_with(action, (1 << unsigned(status::waiting))
                               | (1 << unsigned(status::sending_lock)));
 }

runtime.cc

@@ -113,6 +113,9 @@ void abort(const char *fmt, ...)
         do {} while (true);
     }
 
+#if CONF_lazy_stack
+    sched::ensure_next_stack_page_if_preemptable();
+#endif
     arch::irq_disable();
 
     static char msg[1024];