zephyr: work: Create ContextExt

Because Zephyr's workqueue re-scheduling is based on the Context struct,
not the waker, we need a way to get back to our `WorkInfo` struct that
this Context is embedded in.

Create a `ContextExt` that adds some methods to the `Context` to be able
to indicate when the scheduler should reschedule this work.

Signed-off-by: David Brown <[email protected]>

Author: d3zd3z

zephyr/src/kio.rs

@@ -8,15 +8,20 @@
 //! [`futures`]: crate::work::futures
 
 use core::ffi::CStr;
-use core::task::Poll;
+use core::task::{Context, Poll};
 use core::{future::Future, pin::Pin};
 
-use crate::time::NoWait;
+use crate::sys::queue::Queue;
+use crate::sys::sync::Semaphore;
+use crate::time::{NoWait, Timeout};
 use crate::work::futures::WakeInfo;
+use crate::work::Signal;
 use crate::work::{futures::JoinHandle, futures::WorkBuilder, WorkQueue};
 
 pub mod sync;
 
+pub use crate::work::futures::sleep;
+
 /// Run an async future on the given worker thread.
 ///
 /// Arrange to have the given future run on the given worker thread.  The resulting `JoinHandle` has
@@ -94,3 +99,67 @@ impl Future for YieldNow {
         }
     }
 }
+
+/// Extensions on [`Context`] to support scheduling via Zephyr's workqueue system.
+///
+/// All of these are called from within the context of running work, and indicate what _next_
+/// should cause this work to be run again.  If none of these methods are called before the work
+/// exits, the work will be scheduled to run after `Forever`, which is not useful.  There may be
+/// later support for having a `Waker` that can schedule work from another context.
+///
+/// Note that the events to wait on, such as Semaphores or channels, if there are multiple threads
+/// that can wait for them, might cause this worker to run, but not actually be available.  As such,
+/// to maintain the non-blocking requirements of Work, [`Semaphore::take`], and the blocking `send`
+/// and `recv` operations on channels should not be used, even after being woken.
+///
+/// For the timeout [`Forever`] is useful to indicate there is no timeout.  If called with
+/// [`NoWait`], the work will be immediately scheduled. In general, it is better to query the
+/// underlying object directly rather than have the overhead of being rescheduled.
+///
+/// # Safety
+///
+/// The lifetime bounds on the items waited for ensure that these items live at least as long as the
+/// work queue.  Practically, this can only be satisfied by using something with 'static' lifetime,
+/// or embedding the value in the Future itself.
+///
+/// With the Zephyr executor, the `Context` is embedded within a `WakeInfo` struct, which this makes
+/// use of.  If a different executor were to be used, these calls would result in undefined
+/// behavior.
+///
+/// This could be checked at runtime, but it would have runtime cost.
+pub trait ContextExt {
+    /// Indicate the work should next be scheduled based on a semaphore being available for "take".
+    ///
+    /// The work will be scheduled either when the given semaphore becomes available to 'take', or
+    /// after the timeout.
+    fn add_semaphore<'a>(&'a mut self, sem: &'a Semaphore, timeout: impl Into<Timeout>);
+
+    /// Indicate that the work should be scheduled after receiving the given [`Signal`], or the
+    /// timeout occurs.
+    fn add_signal<'a>(&'a mut self, signal: &'a Signal, timeout: impl Into<Timeout>);
+
+    /// Indicate that the work should be scheduled when the given [`Queue`] has data available to
+    /// recv, or the timeout occurs.
+    fn add_queue<'a>(&'a mut self, queue: &'a Queue, timeout: impl Into<Timeout>);
+}
+
+/// Implementation of ContextExt for the Rust [`Context`] type.
+impl<'b> ContextExt for Context<'b> {
+    fn add_semaphore<'a>(&'a mut self, sem: &'a Semaphore, timeout: impl Into<Timeout>) {
+        let info = unsafe { WakeInfo::from_context(self) };
+        info.add_semaphore(sem);
+        info.timeout = timeout.into();
+    }
+
+    fn add_signal<'a>(&'a mut self, signal: &'a Signal, timeout: impl Into<Timeout>) {
+        let info = unsafe { WakeInfo::from_context(self) };
+        info.add_signal(signal);
+        info.timeout = timeout.into();
+    }
+
+    fn add_queue<'a>(&'a mut self, queue: &'a Queue, timeout: impl Into<Timeout>) {
+        let info = unsafe { WakeInfo::from_context(self) };
+        info.add_queue(queue);
+        info.timeout = timeout.into();
+    }
+}

zephyr/src/sync/channel.rs

@@ -51,9 +51,9 @@ use core::mem::MaybeUninit;
 use core::pin::Pin;
 use core::task::Poll;
 
+use crate::kio::ContextExt;
 use crate::sys::queue::Queue;
 use crate::time::{Duration, Forever, NoWait, Timeout};
-use crate::work::futures::WakeInfo;
 
 mod counter;
 
@@ -270,18 +270,14 @@ impl<'a, T: Unpin> Future for SendFuture<'a, T> {
         this.msg = Some(msg);
 
         // Otherwise, schedule to wake up on receipt or timeout.
-        let info = unsafe { WakeInfo::from_context(cx) };
         match &this.sender.flavor {
             SenderFlavor::Unbounded { .. } => {
                 panic!("Implementation error: unbounded queues should never fail");
             }
             SenderFlavor::Bounded(chan) => {
-                unsafe {
-                    info.add_queue(&chan.free);
-                }
+                cx.add_queue(&chan.free, this.timeout);
             }
         }
-        info.timeout = this.timeout;
 
         Poll::Pending
     }
@@ -522,11 +518,7 @@ impl<'a, T> Future for RecvFuture<'a, T> {
         }
 
         // Otherwise, schedule to wakeup on receipt or timeout.
-        let info = unsafe { WakeInfo::from_context(cx) };
-        unsafe {
-            info.add_queue(self.receiver.as_queue());
-        }
-        info.timeout = self.timeout;
+        cx.add_queue(self.receiver.as_queue(), self.timeout);
         self.waited = true;
 
         Poll::Pending

zephyr/src/sys/sync/semaphore.rs

@@ -25,10 +25,9 @@ use core::mem;
 #[cfg(CONFIG_RUST_ALLOC)]
 use zephyr_sys::ETIMEDOUT;
 
+use crate::kio::ContextExt;
 #[cfg(CONFIG_RUST_ALLOC)]
 use crate::time::NoWait;
-#[cfg(CONFIG_RUST_ALLOC)]
-use crate::work::futures::WakeInfo;
 use crate::{
     error::{to_result_void, Result},
     object::{Fixed, StaticKernelObject, Wrapped},
@@ -144,13 +143,7 @@ impl<'a> Future for SemTake<'a> {
         }
 
         // TODO: Clean this up.
-        let info = unsafe { WakeInfo::from_context(cx) };
-        unsafe {
-            // SAFETY: The semaphore must outlive the queued event.  The lifetime ensures that the
-            // Future won't outlive the semaphore.
-            info.add_semaphore(self.sem);
-        }
-        info.timeout = self.timeout;
+        cx.add_semaphore(self.sem, self.timeout);
         self.ran = true;
 
         Poll::Pending

zephyr/src/work.rs

@@ -179,15 +179,14 @@ use core::{
     ptr,
     task::Poll,
 };
-use futures::WakeInfo;
 
 use zephyr_sys::{
     k_poll_signal, k_poll_signal_check, k_poll_signal_init, k_poll_signal_raise,
     k_poll_signal_reset, k_work, k_work_init, k_work_q, k_work_queue_config, k_work_queue_init,
     k_work_queue_start, k_work_submit, k_work_submit_to_queue, ETIMEDOUT,
 };
 
-use crate::{error::to_result_void, object::Fixed, simpletls::StaticTls, sys::thread::ThreadStack, time::Timeout};
+use crate::{error::to_result_void, kio::ContextExt, object::Fixed, simpletls::StaticTls, sys::thread::ThreadStack, time::Timeout};
 
 pub mod futures;
 
@@ -485,13 +484,7 @@ impl<'a> Future for SignalWait<'a> {
             return Poll::Ready(Err(crate::Error(ETIMEDOUT)));
         }
 
-        let info = unsafe { WakeInfo::from_context(cx) };
-        unsafe {
-            // SAFETY: The Signal must outlive the queued event.  The lifetime ensure that the
-            // Future can't outlive the Signal.
-            info.add_signal(self.signal);
-        }
-        info.timeout = self.timeout;
+        cx.add_signal(self.signal, self.timeout);
         self.ran = true;
 
         Poll::Pending

zephyr/src/work/futures.rs

@@ -285,7 +285,7 @@ impl WakeInfo {
     }
 
     /// Add an event that represents waiting for a semaphore to be available for "take".
-    pub unsafe fn add_semaphore<'a>(&'a mut self, sem: &'a Semaphore) {
+    pub fn add_semaphore<'a>(&'a mut self, sem: &'a Semaphore) {
         // SAFETY: Fill with zeroed memory, initializatuon happens in the init function next.
         self.events.push(unsafe { mem::zeroed() });
         let ev = self.events.last().unwrap();
@@ -301,7 +301,7 @@ impl WakeInfo {
     }
 
     /// Add an event that represents waiting for a signal.
-    pub unsafe fn add_signal<'a>(&'a mut self, signal: &'a Signal) {
+    pub fn add_signal<'a>(&'a mut self, signal: &'a Signal) {
         // SAFETY: Fill with zeroed memory, initializatuon happens in the init function next.
         self.events.push(unsafe { mem::zeroed() });
         let ev = self.events.last().unwrap();
@@ -317,7 +317,7 @@ impl WakeInfo {
     }
 
     /// Add an event that represents waiting for a queue to have a message.
-    pub unsafe fn add_queue<'a>(&'a mut self, queue: &'a Queue) {
+    pub fn add_queue<'a>(&'a mut self, queue: &'a Queue) {
         // SAFETY: Fill with zeroed memory, initializatuon happens in the init function next.
         self.events.push(unsafe { mem::zeroed() });
         let ev = self.events.last().unwrap();