scheduler: Fix bug in scheduling safety task

Fixed bugs related to handling and scheduling safety task.
#18
#20
#39

Author: prabakaranklst

src/kyron/BUILD

@@ -100,3 +100,15 @@ rust_binary(
     visibility = ["//visibility:public"],
     deps = _EXAMPLE_DEPS,
 )
+
+rust_binary(
+    name = "safety_task",
+    srcs = [
+        "examples/safety_task.rs",
+    ],
+    proc_macro_deps = [
+        "//src/kyron-macros:runtime_macros",
+    ],
+    visibility = ["//visibility:public"],
+    deps = _EXAMPLE_DEPS,
+)

src/kyron/examples/safety_task.rs

@@ -0,0 +1,69 @@
+//
+// Copyright (c) 2025 Contributors to the Eclipse Foundation
+//
+// See the NOTICE file(s) distributed with this work for additional
+// information regarding copyright ownership.
+//
+// This program and the accompanying materials are made available under the
+// terms of the Apache License Version 2.0 which is available at
+// <https://www.apache.org/licenses/LICENSE-2.0>
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+use kyron::prelude::*;
+use kyron::safety;
+use kyron::spawn_on_dedicated;
+use kyron_foundation::prelude::*;
+
+async fn failing_safety_task() -> Result<(), String> {
+    info!("Worker-N: failing_safety_task");
+    Err("Intentional failure".to_string())
+}
+
+async fn passing_safety_task() -> Result<(), String> {
+    info!("Worker-N: passing_safety_task");
+    Ok(())
+}
+
+async fn passing_non_safety_task() -> Result<(), String> {
+    info!("Dedicated worker (dw1): passing_non_safety_task");
+    Ok(())
+}
+
+fn main() {
+    tracing_subscriber::fmt()
+        .with_target(false) // Optional: Remove module path
+        .with_max_level(Level::DEBUG)
+        .with_thread_ids(true)
+        .with_thread_names(true)
+        .init();
+
+    // Create runtime
+    let (builder, _engine_id) = kyron::runtime::RuntimeBuilder::new().with_engine(
+        ExecutionEngineBuilder::new()
+            .task_queue_size(256)
+            .enable_safety_worker(ThreadParameters::default())
+            .with_dedicated_worker("dw1".into(), ThreadParameters::default())
+            .workers(2),
+    );
+
+    let mut runtime = builder.build().unwrap();
+    // Put programs into runtime and run them
+    runtime.block_on(async move {
+        let handle1 = safety::spawn(failing_safety_task());
+        let handle2 = safety::spawn(passing_safety_task());
+        let handle3 = spawn_on_dedicated(passing_non_safety_task(), "dw1".into());
+
+        info!("=============================== Spawned all tasks ===============================");
+
+        let _ = handle1.await;
+        info!("Safety worker: Since safety task fails, safety worker executes parent task from this statement onwards.");
+        let _ = handle2.await;
+        let _ = handle3.await;
+
+        info!("Safety worker: Program finished running.");
+    });
+
+    info!("Exit.");
+}

src/kyron/src/scheduler/context.rs

@@ -16,6 +16,7 @@ use crate::core::types::FutureBox;
 use crate::futures::reusable_box_future::ReusableBoxFuture;
 use crate::safety::SafetyResult;
 use crate::scheduler::driver::Drivers;
+use crate::scheduler::task::async_task::TaskHeader;
 use ::core::future::Future;
 use core::cell::Cell;
 use core::cell::RefCell;
@@ -318,6 +319,9 @@ pub(crate) struct WorkerContext {
     /// Helper flag to check if safety was enabled in runtime builder
     is_safety_enabled: bool,
 
+    /// The task that is currently run by worker
+    running_task: Cell<Option<*const TaskHeader>>,
+
     wakeup_time: Cell<Option<u64>>,
 }
 
@@ -399,6 +403,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,
+            running_task: Cell::new(None),
             wakeup_time: Cell::new(None),
             drivers: Some(self.drivers),
         }
@@ -452,6 +457,33 @@ pub(crate) fn ctx_is_with_safety() -> bool {
         .unwrap_or_default()
 }
 
+///
+/// Set currently running task of the worker
+///
+pub(crate) fn ctx_set_running_task(task: *const TaskHeader) {
+    CTX.try_with(|ctx| {
+        ctx.borrow().as_ref().expect("Called before CTX init?").running_task.set(Some(task));
+    })
+    .unwrap_or_default();
+}
+
+#[allow(dead_code)]
+///
+/// Check whether there is any safety error for the running task of the worker and clear it
+///
+pub(crate) fn ctx_check_running_task_safety_error() -> bool {
+    CTX.try_with(|ctx| {
+        let mut binding = ctx.borrow_mut();
+        let ctx = binding.as_mut().expect("Called before CTX init?");
+        let mut res = false;
+        if let Some(task) = ctx.running_task.take() {
+            res = unsafe { (*(task as *mut TaskHeader)).check_safety_error_and_clear() };
+        }
+        res
+    })
+    .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/join_handle.rs

@@ -13,6 +13,7 @@
 use kyron_foundation::prelude::*;
 use kyron_foundation::{not_recoverable_error, prelude::CommonErrors};
 
+use crate::scheduler::context::ctx_set_running_task;
 use crate::{
     futures::{FutureInternalReturn, FutureState},
     TaskRef,
@@ -74,34 +75,52 @@ impl<T: Send + 'static> Future for JoinHandle<T> {
                 if was_set {
                     FutureInternalReturn::default()
                 } else {
+                    // Check whether there is safety error for the completed task
+                    let task_hdr = TaskRef::into_raw(self.for_task.clone());
+                    if unsafe { (*task_hdr).get_safety_error() } {
+                        // Change the running task to safety task to schedule the current task into safety worker
+                        ctx_set_running_task(task_hdr);
+                        waker.wake_by_ref();
+                        FutureInternalReturn::polled()
+                    } else {
+                        let mut ret: Result<T, CommonErrors> = Err(CommonErrors::NoData);
+                        let ret_as_ptr = &mut ret as *mut _;
+                        self.for_task.get_return_val(ret_as_ptr as *mut u8);
+
+                        match ret {
+                            Ok(v) => FutureInternalReturn::ready(Ok(v)),
+                            Err(CommonErrors::OperationAborted) => FutureInternalReturn::ready(Err(CommonErrors::OperationAborted)),
+                            Err(e) => {
+                                not_recoverable_error!(with e, "There has been an error in a task that is not recoverable ({})!");
+                            }
+                        }
+                    }
+                }
+            }
+            FutureState::Polled => {
+                let waker = cx.waker();
+
+                // Set the waker, return values tells what have happen and took care about correct synchronization
+                let was_set = self.for_task.set_join_handle_waker(waker.clone());
+
+                if was_set {
+                    FutureInternalReturn::default()
+                } else {
+                    // Safety belows forms AqrRel so waker is really written before we do marking
                     let mut ret: Result<T, CommonErrors> = Err(CommonErrors::NoData);
                     let ret_as_ptr = &mut ret as *mut _;
                     self.for_task.get_return_val(ret_as_ptr as *mut u8);
 
                     match ret {
                         Ok(v) => FutureInternalReturn::ready(Ok(v)),
+                        Err(CommonErrors::NoData) => FutureInternalReturn::polled(),
                         Err(CommonErrors::OperationAborted) => FutureInternalReturn::ready(Err(CommonErrors::OperationAborted)),
                         Err(e) => {
                             not_recoverable_error!(with e, "There has been an error in a task that is not recoverable ({})!");
                         }
                     }
                 }
             }
-            FutureState::Polled => {
-                // Safety belows forms AqrRel so waker is really written before we do marking
-                let mut ret: Result<T, CommonErrors> = Err(CommonErrors::NoData);
-                let ret_as_ptr = &mut ret as *mut _;
-                self.for_task.get_return_val(ret_as_ptr as *mut u8);
-
-                match ret {
-                    Ok(v) => FutureInternalReturn::ready(Ok(v)),
-                    Err(CommonErrors::NoData) => FutureInternalReturn::polled(),
-                    Err(CommonErrors::OperationAborted) => FutureInternalReturn::ready(Err(CommonErrors::OperationAborted)),
-                    Err(e) => {
-                        not_recoverable_error!(with e, "There has been an error in a task that is not recoverable ({})!");
-                    }
-                }
-            }
             FutureState::Finished => {
                 not_recoverable_error!("Future polled after it finished!");
             }
@@ -130,6 +149,7 @@ mod tests {
     use kyron_testing::prelude::*;
 
     #[test]
+    #[cfg(not(miri))]
     fn test_join_handler_ready_task_get_correct_result_to_handle() {
         let scheduler = create_mock_scheduler();
 
@@ -175,6 +195,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg(not(miri))]
     fn test_join_handler_aborted_task_produce_handle_abort_result() {
         let scheduler = create_mock_scheduler();
 
@@ -205,6 +226,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg(not(miri))]
     #[should_panic]
     fn test_join_handler_panics_when_fetched_data_and_repolled() {
         let scheduler = create_mock_scheduler();
@@ -256,6 +278,40 @@ mod tests {
             assert_eq!(poller.poll(), ::core::task::Poll::Ready(Ok(0)));
         }
     }
+
+    #[test]
+    fn test_join_handle_waker_is_set_in_polled_state_also() {
+        let scheduler = create_mock_scheduler();
+
+        {
+            // Data is present before first poll of join handle
+            let task = ArcInternal::new(AsyncTask::new(box_future(test_function::<u32>()), 1, scheduler.clone()));
+
+            let handle = JoinHandle::<u32>::new(TaskRef::new(task.clone()));
+
+            let mut poller = TestingFuturePoller::new(handle);
+
+            let waker_mock1 = TrackableWaker::new();
+            let waker1 = waker_mock1.get_waker();
+
+            let waker_mock2 = TrackableWaker::new();
+            let waker2 = waker_mock2.get_waker();
+
+            let _ = poller.poll_with_waker(&waker1);
+            // Now in polled state, poll again with waker2
+            let _ = poller.poll_with_waker(&waker2);
+            {
+                let waker = noop_waker();
+                let mut cx = Context::from_waker(&waker);
+                task.poll(&mut cx); // task done
+            }
+
+            assert!(!waker_mock1.was_waked());
+            // this should be TRUE
+            assert!(waker_mock2.was_waked());
+            assert_eq!(poller.poll(), ::core::task::Poll::Ready(Ok(0)));
+        }
+    }
 }
 
 #[cfg(test)]

src/kyron/src/scheduler/safety_waker.rs

@@ -57,11 +57,9 @@ static VTABLE: RawWakerVTable = RawWakerVTable::new(clone_waker, wake, wake_by_r
 ///
 /// Waker will store internally a pointer to the ref counted Task.
 ///
-pub(crate) unsafe fn create_safety_waker(waker: Waker) -> Waker {
-    let raw_waker = RawWaker::new(waker.data(), &VTABLE);
-
-    // Forget original as we took over the ownership, so ref count
-    ::core::mem::forget(waker);
+pub(crate) fn create_safety_waker(ptr: TaskRef) -> Waker {
+    let ptr = TaskRef::into_raw(ptr); // Extracts the pointer from TaskRef not decreasing it's reference count. Since we have a clone here, ref cnt was already increased
+    let raw_waker = RawWaker::new(ptr as *const (), &VTABLE);
 
     // Convert RawWaker to Waker
     unsafe { Waker::from_raw(raw_waker) }