scheduler: Fix bug in scheduling safety task

When a task results in error, then it is scheduled on safety worker if enabled. Fixed a bug in wake/schedule of safety tasks which occurs when nested wakers/tasks are used.

Author: prabakaranklst

src/kyron/src/scheduler/context.rs

@@ -318,6 +318,9 @@ pub(crate) struct WorkerContext {
     /// Helper flag to check if safety was enabled in runtime builder
     is_safety_enabled: bool,
 
+    /// Flag to check if current task resulted in safety error, then schedule on safety worker from wake function.
+    is_task_safety_error: bool,
+
     wakeup_time: Cell<Option<u64>>,
 }
 
@@ -399,6 +402,7 @@ impl ContextBuilder {
             worker_id: Cell::new(self.worker_id.expect("Worker type must be set in context builder!")),
             handler: RefCell::new(Some(Rc::new(self.handle.expect("Handler type must be set in context builder!")))),
             is_safety_enabled: self.is_with_safety,
+            is_task_safety_error: false,
             wakeup_time: Cell::new(None),
             drivers: Some(self.drivers),
         }
@@ -452,6 +456,32 @@ pub(crate) fn ctx_is_with_safety() -> bool {
         .unwrap_or_default()
 }
 
+#[allow(dead_code)] // mock function is used in tests instead of this one
+///
+/// Sets task safety error flag
+///
+pub(crate) fn ctx_set_task_safety_error(is_error: bool) {
+    CTX.try_with(|ctx| {
+        ctx.borrow_mut().as_mut().expect("Called before CTX init?").is_task_safety_error = is_error;
+    })
+    .unwrap_or_default();
+}
+
+#[allow(dead_code)] // mock function is used in tests instead of this one
+///
+/// Check if current task resulted in safety error and clear the flag.
+///
+pub(crate) fn ctx_check_task_safety_error_and_clear() -> bool {
+    CTX.try_with(|ctx| {
+        let mut binding = ctx.borrow_mut();
+        let ctx = binding.as_mut().expect("Called before CTX init?");
+        let val = ctx.is_task_safety_error;
+        ctx.is_task_safety_error = false;
+        val
+    })
+    .unwrap_or_default()
+}
+
 pub(crate) fn ctx_set_wakeup_time(time: u64) {
     CTX.try_with(|ctx| ctx.borrow().as_ref().expect("Called before CTX init?").wakeup_time.set(Some(time)))
         .unwrap_or_default();

src/kyron/src/scheduler/mod.rs

@@ -17,7 +17,6 @@ pub(crate) mod execution_engine;
 pub(crate) mod join_handle;
 pub(crate) mod workers;
 
-pub mod safety_waker;
 pub mod scheduler_mt;
 pub mod task;
 pub mod waker;

src/kyron/src/scheduler/safety_waker.rs

@@ -13,8 +13,11 @@
 
 use super::task_state::*;
 use crate::core::types::*;
-use crate::scheduler::safety_waker::create_safety_waker;
+#[cfg(not(any(test, feature = "runtime-api-mock")))]
+use crate::scheduler::context::ctx_set_task_safety_error;
 use crate::scheduler::scheduler_mt::SchedulerTrait;
+#[cfg(any(test, feature = "runtime-api-mock"))]
+use crate::testing::mock::ctx_set_task_safety_error;
 use ::core::future::Future;
 use ::core::mem;
 use ::core::ops::{Deref, DerefMut};
@@ -308,9 +311,8 @@ where
                     self.handle_waker.with_mut(|ptr: *mut Option<Waker>| match unsafe { (*ptr).take() } {
                         Some(v) => {
                             if is_safety_err && self.is_with_safety {
-                                unsafe {
-                                    create_safety_waker(v).wake();
-                                }
+                                ctx_set_task_safety_error(true);
+                                v.wake();
                             } else {
                                 v.wake();
                             }

src/kyron/src/scheduler/task/async_task.rs

@@ -14,6 +14,10 @@
 use kyron_foundation::prelude::FoundationAtomicPtr;
 
 use super::task::async_task::*;
+#[cfg(not(any(test, feature = "runtime-api-mock")))]
+use crate::scheduler::context::ctx_check_task_safety_error_and_clear;
+#[cfg(any(test, feature = "runtime-api-mock"))]
+use crate::testing::mock::ctx_check_task_safety_error_and_clear;
 use core::task::{RawWaker, RawWakerVTable, Waker};
 
 fn clone_waker(data: *const ()) -> RawWaker {
@@ -32,16 +36,23 @@ fn wake(data: *const ()) {
     let task_header_ptr = data as *const TaskHeader;
     let task_ref = unsafe { TaskRef::from_raw(task_header_ptr) };
 
-    task_ref.schedule();
-
+    if ctx_check_task_safety_error_and_clear() {
+        task_ref.schedule_safety();
+    } else {
+        task_ref.schedule();
+    }
     drop(task_ref); // wake uses move semantic, so we are owner of data now, so we need to cleanup
 }
 
 fn wake_by_ref(data: *const ()) {
     let task_header_ptr = data as *const TaskHeader;
     let task_ref = unsafe { TaskRef::from_raw(task_header_ptr) };
 
-    task_ref.schedule();
+    if ctx_check_task_safety_error_and_clear() {
+        task_ref.schedule_safety();
+    } else {
+        task_ref.schedule();
+    }
 
     ::core::mem::forget(task_ref); // don't touch refcount from our data since this is done by drop_waker
 }

src/kyron/src/scheduler/waker.rs

@@ -225,6 +225,25 @@ where
     }
 }
 
+// Thread-local storage to track task safety errors
+thread_local! {
+    static THREAD_CTX_SAFETY_ERROR_FLAG: RefCell<bool> = const { RefCell::new(false) };
+}
+
+pub fn ctx_set_task_safety_error(is_error: bool) {
+    THREAD_CTX_SAFETY_ERROR_FLAG.with(|flag| {
+        *flag.borrow_mut() = is_error;
+    });
+}
+
+pub fn ctx_check_task_safety_error_and_clear() -> bool {
+    THREAD_CTX_SAFETY_ERROR_FLAG.with(|flag| {
+        let current = *flag.borrow();
+        *flag.borrow_mut() = false;
+        current
+    })
+}
+
 ///
 /// Spawns a given `future` into runtime and let it execute on dedicated worker using `worker_id`.
 /// This function allocates a `future` dynamically using [`Box`]