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/src/scheduler/context.rs

@@ -481,6 +481,7 @@ pub(crate) fn ctx_get_drivers() -> Drivers {
         .unwrap()
 }
 
+#[allow(dead_code)] // Mock function is used instead of this if mock runtime feature is enabled
 ///
 /// Sets currently running `task`
 ///
@@ -494,6 +495,7 @@ pub(super) fn ctx_set_running_task(task: TaskRef) {
         });
 }
 
+#[allow(dead_code)] // Mock function is used instead of this if mock runtime feature is enabled
 ///
 /// Clears currently running `task`
 ///
@@ -505,6 +507,7 @@ pub(super) fn ctx_unset_running_task() {
         .map_err(|_| {});
 }
 
+#[allow(dead_code)] // Mock function is used instead of this if mock runtime feature is enabled
 ///
 /// Gets currently running `task id`
 ///
@@ -523,6 +526,24 @@ pub(crate) fn ctx_get_running_task_id() -> Option<TaskId> {
     })
 }
 
+#[allow(dead_code)] // Mock function is used instead of this if mock runtime feature is enabled
+///
+/// Returns `true` if the running task resulted in safety error
+///
+pub(crate) fn ctx_get_task_safety_error() -> bool {
+    CTX.try_with(|ctx| {
+        // This funcation can be called from a thread outside of Kyron runtime through wake()/wake_by_ref(), so we need to check for ctx presence
+        if let Some(cx) = ctx.borrow().as_ref() {
+            cx.running_task.borrow().as_ref().is_some_and(|task| task.get_task_safety_error())
+        } else {
+            false
+        }
+    })
+    .unwrap_or_else(|e| {
+        panic!("Something is really bad here, error {}!", e);
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/kyron/src/scheduler/execution_engine.rs

@@ -506,7 +506,8 @@ mod tests {
     #[test]
     #[cfg(not(miri))] // Provenance issues
     fn create_engine_with_worker_and_verify_ids() {
-        use crate::scheduler::context::{ctx_get_running_task_id, ctx_get_worker_id};
+        use crate::scheduler::context::ctx_get_worker_id;
+        use crate::testing::mock_context::ctx_get_running_task_id;
         let mut engine = ExecutionEngineBuilder::new().workers(1).task_queue_size(16).set_engine_id(1).build();
         let result: Result<bool, ()> = engine
             .run_in_engine(async move {

src/kyron/src/scheduler/join_handle.rs

@@ -65,7 +65,7 @@ impl<T: Send + 'static> Future for JoinHandle<T> {
     ///
     fn poll(self: ::core::pin::Pin<&mut Self>, cx: &mut ::core::task::Context<'_>) -> Poll<Self::Output> {
         let res: FutureInternalReturn<JoinResult<T>> = match self.state {
-            FutureState::New => {
+            FutureState::New | FutureState::Polled => {
                 let waker = cx.waker();
 
                 // Set the waker, return values tells what have happen and took care about correct synchronization
@@ -80,28 +80,14 @@ impl<T: Send + 'static> Future for JoinHandle<T> {
 
                     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!");
             }
@@ -262,6 +248,41 @@ 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 worker_id = create_mock_worker_id(0, 1);
+            let task = ArcInternal::new(AsyncTask::new(box_future(test_function::<u32>()), &worker_id, 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)]
@@ -284,8 +305,9 @@ mod tests {
 
     #[test]
     fn test_join_handler_mt_get_result() {
-        let builder = Builder::new();
-
+        let mut builder = Builder::new();
+        // Limit preemption to avoid loom error "Model exceeded maximum number of branches."
+        builder.preemption_bound = Some(4);
         builder.check(|| {
             let scheduler = create_mock_scheduler();
 
@@ -307,22 +329,17 @@ mod tests {
 
                 let waker_mock = TrackableWaker::new();
                 let waker = waker_mock.get_waker();
-                let mut was_pending = false;
-
                 loop {
                     match poller.poll_with_waker(&waker) {
                         Poll::Ready(v) => {
                             assert_eq!(v, Ok(1234));
-
-                            if was_pending {
-                                assert!(waker_mock.was_waked());
-                            }
+                            // Note:
+                            // Cannot check whether the waker was woken or not since the waker is set in the join handle poll every time if task is not yet done.
+                            // So depending on the interleaving, the task may finish before the waker is set.
 
                             break;
                         }
-                        Poll::Pending => {
-                            was_pending = true;
-                        }
+                        Poll::Pending => {}
                     }
                     loom::hint::spin_loop();
                 }

src/kyron/src/scheduler/safety_waker.rs

@@ -55,11 +55,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) }

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

@@ -13,9 +13,9 @@
 
 use super::task_state::*;
 use crate::core::types::*;
-use crate::scheduler::safety_waker::create_safety_waker;
 use crate::scheduler::scheduler_mt::SchedulerTrait;
 use crate::scheduler::workers::worker_types::WorkerId;
+use ::core::cell::Cell;
 use ::core::future::Future;
 use ::core::mem;
 use ::core::ops::{Deref, DerefMut};
@@ -82,7 +82,7 @@ pub(crate) enum TaskStage<T, ResultType> {
 pub(crate) struct TaskHeader {
     pub(in crate::scheduler) state: TaskState,
     id: TaskId,
-
+    is_safety_error: Cell<bool>, // Flag to indicate whether task resulted in safety error
     vtable: &'static TaskVTable, // API entrypoint to typed task
 }
 
@@ -97,6 +97,7 @@ impl TaskHeader {
         Self {
             state: TaskState::new(),
             id: TaskId::new(worker_id),
+            is_safety_error: Cell::new(false),
             vtable: create_task_vtable::<T, AllocatedFuture, SchedulerType>(),
         }
     }
@@ -111,9 +112,18 @@ impl TaskHeader {
         Self {
             state: TaskState::new(),
             id: TaskId::new(worker_id),
+            is_safety_error: Cell::new(false),
             vtable: create_task_s_vtable::<T, E, AllocatedFuture, SchedulerType>(),
         }
     }
+
+    pub(crate) fn set_safety_error(&self) {
+        self.is_safety_error.set(true);
+    }
+
+    pub(crate) fn get_safety_error(&self) -> bool {
+        self.is_safety_error.get()
+    }
 }
 
 #[derive(PartialEq, Debug)]
@@ -203,13 +213,20 @@ where
     }
 
     pub(crate) fn set_waker(&self, waker: Waker) -> bool {
-        unsafe {
-            self.handle_waker.with_mut(|ptr| {
-                *ptr = Some(waker);
-            })
+        // Safety: Unset join handle flag before setting waker, the flag would have been set previously in the first poll of join handle.
+        // If flag is not cleared, another worker finishing the task will see the flag set and
+        // read the waker to call wake() while it is written here
+        if self.header.state.unset_join_handle() {
+            unsafe {
+                self.handle_waker.with_mut(|ptr| {
+                    *ptr = Some(waker);
+                })
+            };
+
+            return self.header.state.set_join_handle();
         }
 
-        self.header.state.set_waker() // Safety: makes sure storing waker is not reordered behind this operation
+        false
     }
 
     ///
@@ -308,12 +325,10 @@ 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();
-                                }
-                            } else {
-                                v.wake();
+                                // Set saftey error flag which will be checked in wake()/wkae_by_ref() to schedule parent task into safety worker
+                                self.header.set_safety_error();
                             }
+                            v.wake();
                         }
                         None => not_recoverable_error!("We shall never be here if we have HadConnectedJoinHandle set!"),
                     })
@@ -614,6 +629,10 @@ impl TaskRef {
         snapshot.is_completed() || snapshot.is_canceled()
     }
 
+    pub(crate) fn get_task_safety_error(&self) -> bool {
+        unsafe { self.header.as_ref().get_safety_error() }
+    }
+
     pub(crate) fn id(&self) -> TaskId {
         unsafe { self.header.as_ref().id }
     }
@@ -641,7 +660,10 @@ mod tests {
         safety::SafetyResult,
         scheduler::{
             scheduler_mt::SchedulerTrait,
-            task::async_task::{TaskId, TaskRef},
+            task::{
+                async_task::{TaskId, TaskRef},
+                task_context::TaskContextGuard,
+            },
         },
         testing::*,
     };
@@ -778,6 +800,7 @@ mod tests {
         safety_task_ref.set_join_handle_waker(waker.clone()); // Mimic that JoinHandler is set
 
         let mut ctx = Context::from_waker(&waker);
+        let _guard = TaskContextGuard::new(safety_task_ref.clone());
         assert_eq!(safety_task_ref.poll(&mut ctx), TaskPollResult::Done);
 
         let mut result: SafetyResult<bool, bool> = Ok(true);
@@ -801,6 +824,7 @@ mod tests {
         safety_task_ref.set_join_handle_waker(waker.clone()); // Mimic that JoinHandler is set
 
         let mut ctx = Context::from_waker(&waker);
+        let _guard = TaskContextGuard::new(safety_task_ref.clone());
         assert_eq!(safety_task_ref.poll(&mut ctx), TaskPollResult::Done);
 
         let mut result: SafetyResult<bool, bool> = Ok(true);
@@ -829,6 +853,7 @@ mod tests {
         safety_task_ref.set_join_handle_waker(waker.clone()); // Mimic that JoinHandler is set
 
         let mut ctx = Context::from_waker(&waker);
+        let _guard = TaskContextGuard::new(safety_task_ref.clone());
         assert_eq!(safety_task_ref.poll(&mut ctx), TaskPollResult::Done);
 
         let mut result: SafetyResult<bool, bool> = Ok(true);
@@ -857,6 +882,7 @@ mod tests {
         safety_task_ref.set_join_handle_waker(waker.clone()); // Mimic that JoinHandler is set
 
         let mut ctx = Context::from_waker(&waker);
+        let _guard = TaskContextGuard::new(safety_task_ref.clone());
         assert_eq!(safety_task_ref.poll(&mut ctx), TaskPollResult::Done);
 
         let mut result: SafetyResult<bool, bool> = Ok(true);

src/kyron/src/scheduler/task/task_context.rs

@@ -11,12 +11,13 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+#[cfg(not(any(test, feature = "runtime-api-mock")))]
+use crate::scheduler::context::{ctx_get_running_task_id, ctx_get_task_safety_error, ctx_set_running_task, ctx_unset_running_task};
+#[cfg(any(test, feature = "runtime-api-mock"))]
+use crate::testing::mock_context::{ctx_get_running_task_id, ctx_get_task_safety_error, ctx_set_running_task, ctx_unset_running_task};
 use crate::{
     core::types::TaskId,
-    scheduler::{
-        context::{ctx_get_running_task_id, ctx_get_worker_id, ctx_set_running_task, ctx_unset_running_task},
-        workers::worker_types::WorkerId,
-    },
+    scheduler::{context::ctx_get_worker_id, workers::worker_types::WorkerId},
     TaskRef,
 };
 
@@ -33,6 +34,11 @@ impl TaskContext {
     pub fn task_id() -> Option<TaskId> {
         ctx_get_running_task_id()
     }
+
+    /// Check whether the running task resulted in safety error to schedule parent into safety worker
+    pub(crate) fn should_wake_task_into_safety() -> bool {
+        ctx_get_task_safety_error()
+    }
 }
 
 /// A guard that sets the task on creation and unsets it on drop.