Add no_std support to bevy_tasks (#15464)

# Objective

- Contributes to #15460

## Solution

- Added the following features:
  - `std` (default)
  - `async_executor` (default)
  - `edge_executor`
  - `critical-section`
  - `portable-atomic`
- Added [`edge-executor`](https://crates.io/crates/edge-executor) as a
`no_std` alternative to `async-executor`.
- Updated the `single_threaded_task_pool` to work in `no_std`
environments by gating its reliance on `thread_local`.

## Testing

- Added to `compile-check-no-std` CI command

## Notes

- In previous iterations of this PR, a custom `async-executor`
alternative was vendored in. This raised concerns around maintenance and
testing. In this iteration, an existing version of that same vendoring
is now used, but _only_ in `no_std` contexts. For existing `std`
contexts, the original `async-executor` is used.
- Due to the way statics work, certain `TaskPool` operations have added
restrictions around `Send`/`Sync` in `no_std`. This is because there
isn't a straightforward way to create a thread-local in `no_std`. If
these added constraints pose an issue we can revisit this at a later
date.
- If a user enables both the `async_executor` and `edge_executor`
features, we will default to using `async-executor`. Since enabling
`async_executor` requires `std`, we can safely assume we are in an `std`
context and use the original library.

---------

Co-authored-by: Mike <2180432+hymm@users.noreply.github.com>
Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>

crates/bevy_tasks/Cargo.toml

@@ -9,14 +9,57 @@ license = "MIT OR Apache-2.0"
 keywords = ["bevy"]
 
 [features]
-multi_threaded = ["dep:async-channel", "dep:concurrent-queue"]
+default = ["std", "async_executor"]
+std = [
+  "futures-lite/std",
+  "async-task/std",
+  "spin/std",
+  "edge-executor?/std",
+  "portable-atomic-util?/std",
+]
+multi_threaded = ["std", "dep:async-channel", "dep:concurrent-queue"]
+async_executor = ["std", "dep:async-executor"]
+edge_executor = ["dep:edge-executor"]
+critical-section = [
+  "dep:critical-section",
+  "edge-executor?/critical-section",
+  "portable-atomic?/critical-section",
+]
+portable-atomic = [
+  "dep:portable-atomic",
+  "dep:portable-atomic-util",
+  "edge-executor?/portable-atomic",
+  "async-task/portable-atomic",
+  "spin/portable_atomic",
+]
 
 [dependencies]
-futures-lite = "2.0.1"
+futures-lite = { version = "2.0.1", default-features = false, features = [
-async-executor = "1.11"
+  "alloc",
+] }
+async-task = { version = "4.4.0", default-features = false }
+spin = { version = "0.9.8", default-features = false, features = [
+  "spin_mutex",
+  "rwlock",
+  "once",
+] }
+derive_more = { version = "1", default-features = false, features = [
+  "deref",
+  "deref_mut",
+] }
+
+async-executor = { version = "1.11", optional = true }
+edge-executor = { version = "0.4.1", default-features = false, optional = true }
 async-channel = { version = "2.3.0", optional = true }
 async-io = { version = "2.0.0", optional = true }
 concurrent-queue = { version = "2.0.0", optional = true }
+critical-section = { version = "1.2.0", optional = true }
+portable-atomic = { version = "1", default-features = false, features = [
+  "fallback",
+], optional = true }
+portable-atomic-util = { version = "0.2.4", features = [
+  "alloc",
+], optional = true }
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 wasm-bindgen-futures = "0.4"

crates/bevy_tasks/README.md

@@ -34,6 +34,10 @@ The determining factor for what kind of work should go in each pool is latency r
   await receiving data from somewhere (i.e. disk) and signal other systems when the data is ready
   for consumption. (likely via channels)
 
+## `no_std` Support
+
+To enable `no_std` support in this crate, you will need to disable default features, and enable the `edge_executor` and `critical-section` features. For platforms without full support for Rust atomics, you may also need to enable the `portable-atomic` feature.
+
 [bevy]: https://bevyengine.org
 [rayon]: https://github.com/rayon-rs/rayon
 [async-executor]: https://github.com/stjepang/async-executor

crates/bevy_tasks/src/executor.rs

@@ -0,0 +1,84 @@
+//! Provides a fundamental executor primitive appropriate for the target platform
+//! and feature set selected.
+//! By default, the `async_executor` feature will be enabled, which will rely on
+//! [`async-executor`] for the underlying implementation. This requires `std`,
+//! so is not suitable for `no_std` contexts. Instead, you must use `edge_executor`,
+//! which relies on the alternate [`edge-executor`] backend.
+//!
+//! [`async-executor`]: https://crates.io/crates/async-executor
+//! [`edge-executor`]: https://crates.io/crates/edge-executor
+
+pub use async_task::Task;
+use core::{
+    fmt,
+    panic::{RefUnwindSafe, UnwindSafe},
+};
+use derive_more::{Deref, DerefMut};
+
+#[cfg(feature = "multi_threaded")]
+pub use async_task::FallibleTask;
+
+#[cfg(feature = "async_executor")]
+type ExecutorInner<'a> = async_executor::Executor<'a>;
+
+#[cfg(feature = "async_executor")]
+type LocalExecutorInner<'a> = async_executor::LocalExecutor<'a>;
+
+#[cfg(all(not(feature = "async_executor"), feature = "edge_executor"))]
+type ExecutorInner<'a> = edge_executor::Executor<'a, 64>;
+
+#[cfg(all(not(feature = "async_executor"), feature = "edge_executor"))]
+type LocalExecutorInner<'a> = edge_executor::LocalExecutor<'a, 64>;
+
+/// Wrapper around a multi-threading-aware async executor.
+/// Spawning will generally require tasks to be `Send` and `Sync` to allow multiple
+/// threads to send/receive/advance tasks.
+///
+/// If you require an executor _without_ the `Send` and `Sync` requirements, consider
+/// using [`LocalExecutor`] instead.
+#[derive(Deref, DerefMut, Default)]
+pub struct Executor<'a>(ExecutorInner<'a>);
+
+/// Wrapper around a single-threaded async executor.
+/// Spawning wont generally require tasks to be `Send` and `Sync`, at the cost of
+/// this executor itself not being `Send` or `Sync`. This makes it unsuitable for
+/// global statics.
+///
+/// If need to store an executor in a global static, or send across threads,
+/// consider using [`Executor`] instead.
+#[derive(Deref, DerefMut, Default)]
+pub struct LocalExecutor<'a>(LocalExecutorInner<'a>);
+
+impl Executor<'_> {
+    /// Construct a new [`Executor`]
+    #[allow(dead_code, reason = "not all feature flags require this function")]
+    pub const fn new() -> Self {
+        Self(ExecutorInner::new())
+    }
+}
+
+impl LocalExecutor<'_> {
+    /// Construct a new [`LocalExecutor`]
+    #[allow(dead_code, reason = "not all feature flags require this function")]
+    pub const fn new() -> Self {
+        Self(LocalExecutorInner::new())
+    }
+}
+
+impl UnwindSafe for Executor<'_> {}
+impl RefUnwindSafe for Executor<'_> {}
+
+impl UnwindSafe for LocalExecutor<'_> {}
+impl RefUnwindSafe for LocalExecutor<'_> {}
+
+impl fmt::Debug for Executor<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Executor").finish()
+    }
+}
+
+impl fmt::Debug for LocalExecutor<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("LocalExecutor").finish()
+    }
+}

crates/bevy_tasks/src/iter/adapters.rs

@@ -1,5 +1,7 @@
 use crate::iter::ParallelIterator;
 
+/// Chains two [`ParallelIterator`]s `T` and `U`, first returning
+/// batches from `T`, and then from `U`.
 #[derive(Debug)]
 pub struct Chain<T, U> {
     pub(crate) left: T,
@@ -24,6 +26,7 @@ where
     }
 }
 
+/// Maps a [`ParallelIterator`] `P` using the provided function `F`.
 #[derive(Debug)]
 pub struct Map<P, F> {
     pub(crate) iter: P,
@@ -41,6 +44,7 @@ where
     }
 }
 
+/// Filters a [`ParallelIterator`] `P` using the provided predicate `F`.
 #[derive(Debug)]
 pub struct Filter<P, F> {
     pub(crate) iter: P,
@@ -60,6 +64,7 @@ where
     }
 }
 
+/// Filter-maps a [`ParallelIterator`] `P` using the provided function `F`.
 #[derive(Debug)]
 pub struct FilterMap<P, F> {
     pub(crate) iter: P,
@@ -77,6 +82,7 @@ where
     }
 }
 
+/// Flat-maps a [`ParallelIterator`] `P` using the provided function `F`.
 #[derive(Debug)]
 pub struct FlatMap<P, F> {
     pub(crate) iter: P,
@@ -98,6 +104,7 @@ where
     }
 }
 
+/// Flattens a [`ParallelIterator`] `P`.
 #[derive(Debug)]
 pub struct Flatten<P> {
     pub(crate) iter: P,
@@ -117,6 +124,8 @@ where
     }
 }
 
+/// Fuses a [`ParallelIterator`] `P`, ensuring once it returns [`None`] once, it always
+/// returns [`None`].
 #[derive(Debug)]
 pub struct Fuse<P> {
     pub(crate) iter: Option<P>,
@@ -138,6 +147,7 @@ where
     }
 }
 
+/// Inspects a [`ParallelIterator`] `P` using the provided function `F`.
 #[derive(Debug)]
 pub struct Inspect<P, F> {
     pub(crate) iter: P,
@@ -155,6 +165,7 @@ where
     }
 }
 
+/// Copies a [`ParallelIterator`] `P`'s returned values.
 #[derive(Debug)]
 pub struct Copied<P> {
     pub(crate) iter: P,
@@ -171,6 +182,7 @@ where
     }
 }
 
+/// Clones a [`ParallelIterator`] `P`'s returned values.
 #[derive(Debug)]
 pub struct Cloned<P> {
     pub(crate) iter: P,
@@ -187,6 +199,7 @@ where
     }
 }
 
+/// Cycles a [`ParallelIterator`] `P` indefinitely.
 #[derive(Debug)]
 pub struct Cycle<P> {
     pub(crate) iter: P,

crates/bevy_tasks/src/iter/mod.rs

@@ -1,4 +1,5 @@
 use crate::TaskPool;
+use alloc::vec::Vec;
 
 mod adapters;
 pub use adapters::*;

crates/bevy_tasks/src/lib.rs

@@ -4,9 +4,12 @@
     html_logo_url = "https://bevyengine.org/assets/icon.png",
     html_favicon_url = "https://bevyengine.org/assets/icon.png"
 )]
+#![cfg_attr(not(feature = "std"), no_std)]
 
 extern crate alloc;
 
+mod executor;
+
 mod slice;
 pub use slice::{ParallelSlice, ParallelSliceMut};
 
@@ -37,9 +40,9 @@ mod thread_executor;
 #[cfg(all(not(target_arch = "wasm32"), feature = "multi_threaded"))]
 pub use thread_executor::{ThreadExecutor, ThreadExecutorTicker};
 
-#[cfg(feature = "async-io")]
+#[cfg(all(feature = "async-io", feature = "std"))]
 pub use async_io::block_on;
-#[cfg(not(feature = "async-io"))]
+#[cfg(all(not(feature = "async-io"), feature = "std"))]
 pub use futures_lite::future::block_on;
 pub use futures_lite::future::poll_once;
 
@@ -54,13 +57,17 @@ pub use futures_lite;
 pub mod prelude {
     #[doc(hidden)]
     pub use crate::{
-        block_on,
         iter::ParallelIterator,
         slice::{ParallelSlice, ParallelSliceMut},
         usages::{AsyncComputeTaskPool, ComputeTaskPool, IoTaskPool},
     };
+
+    #[cfg(feature = "std")]
+    #[doc(hidden)]
+    pub use crate::block_on;
 }
 
+#[cfg(feature = "std")]
 use core::num::NonZero;
 
 /// Gets the logical CPU core count available to the current process.
@@ -69,8 +76,18 @@ use core::num::NonZero;
 /// it will return a default value of 1 if it internally errors out.
 ///
 /// This will always return at least 1.
+#[cfg(feature = "std")]
 pub fn available_parallelism() -> usize {
     std::thread::available_parallelism()
         .map(NonZero::<usize>::get)
         .unwrap_or(1)
 }
+
+/// Gets the logical CPU core count available to the current process.
+///
+/// This will always return at least 1.
+#[cfg(not(feature = "std"))]
+pub fn available_parallelism() -> usize {
+    // Without access to std, assume a single thread is available
+    1
+}

crates/bevy_tasks/src/single_threaded_task_pool.rs

@@ -1,12 +1,34 @@
-use alloc::{rc::Rc, sync::Arc};
+use alloc::{string::String, vec::Vec};
 use core::{cell::RefCell, future::Future, marker::PhantomData, mem};
 
 use crate::Task;
 
+#[cfg(feature = "portable-atomic")]
+use portable_atomic_util::Arc;
+
+#[cfg(not(feature = "portable-atomic"))]
+use alloc::sync::Arc;
+
+#[cfg(feature = "std")]
+use crate::executor::LocalExecutor;
+
+#[cfg(not(feature = "std"))]
+use crate::executor::Executor as LocalExecutor;
+
+#[cfg(feature = "std")]
 thread_local! {
-    static LOCAL_EXECUTOR: async_executor::LocalExecutor<'static> = const { async_executor::LocalExecutor::new() };
+    static LOCAL_EXECUTOR: LocalExecutor<'static> = const { LocalExecutor::new() };
 }
 
+#[cfg(not(feature = "std"))]
+static LOCAL_EXECUTOR: LocalExecutor<'static> = const { LocalExecutor::new() };
+
+#[cfg(feature = "std")]
+type ScopeResult<T> = alloc::rc::Rc<RefCell<Option<T>>>;
+
+#[cfg(not(feature = "std"))]
+type ScopeResult<T> = Arc<spin::Mutex<Option<T>>>;
+
 /// Used to create a [`TaskPool`].
 #[derive(Debug, Default, Clone)]
 pub struct TaskPoolBuilder {}
@@ -124,15 +146,13 @@ impl TaskPool {
         // Any usages of the references passed into `Scope` must be accessed through
         // the transmuted reference for the rest of this function.
 
-        let executor = &async_executor::LocalExecutor::new();
+        let executor = &LocalExecutor::new();
         // SAFETY: As above, all futures must complete in this function so we can change the lifetime
-        let executor: &'env async_executor::LocalExecutor<'env> =
+        let executor: &'env LocalExecutor<'env> = unsafe { mem::transmute(executor) };
-            unsafe { mem::transmute(executor) };
 
-        let results: RefCell<Vec<Rc<RefCell<Option<T>>>>> = RefCell::new(Vec::new());
+        let results: RefCell<Vec<ScopeResult<T>>> = RefCell::new(Vec::new());
         // SAFETY: As above, all futures must complete in this function so we can change the lifetime
-        let results: &'env RefCell<Vec<Rc<RefCell<Option<T>>>>> =
+        let results: &'env RefCell<Vec<ScopeResult<T>>> = unsafe { mem::transmute(&results) };
-            unsafe { mem::transmute(&results) };
 
         let mut scope = Scope {
             executor,
@@ -152,7 +172,16 @@ impl TaskPool {
         let results = scope.results.borrow();
         results
             .iter()
-            .map(|result| result.borrow_mut().take().unwrap())
+            .map(|result| {
+                #[cfg(feature = "std")]
+                return result.borrow_mut().take().unwrap();
+
+                #[cfg(not(feature = "std"))]
+                {
+                    let mut lock = result.lock();
+                    lock.take().unwrap()
+                }
+            })
             .collect()
     }
 
@@ -162,29 +191,42 @@ impl TaskPool {
     /// end-user.
     ///
     /// If the provided future is non-`Send`, [`TaskPool::spawn_local`] should be used instead.
-    pub fn spawn<T>(&self, future: impl Future<Output = T> + 'static) -> Task<T>
+    pub fn spawn<T>(
+        &self,
+        future: impl Future<Output = T> + 'static + MaybeSend + MaybeSync,
+    ) -> Task<T>
     where
-        T: 'static,
+        T: 'static + MaybeSend + MaybeSync,
     {
-        #[cfg(target_arch = "wasm32")]
+        #[cfg(all(target_arch = "wasm32", feature = "std"))]
         return Task::wrap_future(future);
 
-        #[cfg(not(target_arch = "wasm32"))]
+        #[cfg(all(not(target_arch = "wasm32"), feature = "std"))]
-        {
+        return LOCAL_EXECUTOR.with(|executor| {
-            LOCAL_EXECUTOR.with(|executor| {
             let task = executor.spawn(future);
             // Loop until all tasks are done
             while executor.try_tick() {}
 
             Task::new(task)
-            })
+        });
-        }
+
+        #[cfg(not(feature = "std"))]
+        return {
+            let task = LOCAL_EXECUTOR.spawn(future);
+            // Loop until all tasks are done
+            while LOCAL_EXECUTOR.try_tick() {}
+
+            Task::new(task)
+        };
     }
 
     /// Spawns a static future on the JS event loop. This is exactly the same as [`TaskPool::spawn`].
-    pub fn spawn_local<T>(&self, future: impl Future<Output = T> + 'static) -> Task<T>
+    pub fn spawn_local<T>(
+        &self,
+        future: impl Future<Output = T> + 'static + MaybeSend + MaybeSync,
+    ) -> Task<T>
     where
-        T: 'static,
+        T: 'static + MaybeSend + MaybeSync,
     {
         self.spawn(future)
     }
@@ -202,9 +244,13 @@ impl TaskPool {
     /// ```
     pub fn with_local_executor<F, R>(&self, f: F) -> R
     where
-        F: FnOnce(&async_executor::LocalExecutor) -> R,
+        F: FnOnce(&LocalExecutor) -> R,
     {
-        LOCAL_EXECUTOR.with(f)
+        #[cfg(feature = "std")]
+        return LOCAL_EXECUTOR.with(f);
+
+        #[cfg(not(feature = "std"))]
+        return f(&LOCAL_EXECUTOR);
     }
 }
 
@@ -213,9 +259,9 @@ impl TaskPool {
 /// For more information, see [`TaskPool::scope`].
 #[derive(Debug)]
 pub struct Scope<'scope, 'env: 'scope, T> {
-    executor: &'scope async_executor::LocalExecutor<'scope>,
+    executor: &'scope LocalExecutor<'scope>,
     // Vector to gather results of all futures spawned during scope run
-    results: &'env RefCell<Vec<Rc<RefCell<Option<T>>>>>,
+    results: &'env RefCell<Vec<ScopeResult<T>>>,
 
     // make `Scope` invariant over 'scope and 'env
     scope: PhantomData<&'scope mut &'scope ()>,
@@ -230,7 +276,7 @@ impl<'scope, 'env, T: Send + 'env> Scope<'scope, 'env, T> {
     /// On the single threaded task pool, it just calls [`Scope::spawn_on_scope`].
     ///
     /// For more information, see [`TaskPool::scope`].
-    pub fn spawn<Fut: Future<Output = T> + 'scope>(&self, f: Fut) {
+    pub fn spawn<Fut: Future<Output = T> + 'scope + MaybeSend>(&self, f: Fut) {
         self.spawn_on_scope(f);
     }
 
@@ -241,7 +287,7 @@ impl<'scope, 'env, T: Send + 'env> Scope<'scope, 'env, T> {
     /// On the single threaded task pool, it just calls [`Scope::spawn_on_scope`].
     ///
     /// For more information, see [`TaskPool::scope`].
-    pub fn spawn_on_external<Fut: Future<Output = T> + 'scope>(&self, f: Fut) {
+    pub fn spawn_on_external<Fut: Future<Output = T> + 'scope + MaybeSend>(&self, f: Fut) {
         self.spawn_on_scope(f);
     }
 
@@ -250,13 +296,41 @@ impl<'scope, 'env, T: Send + 'env> Scope<'scope, 'env, T> {
     /// returned as a part of [`TaskPool::scope`]'s return value.
     ///
     /// For more information, see [`TaskPool::scope`].
-    pub fn spawn_on_scope<Fut: Future<Output = T> + 'scope>(&self, f: Fut) {
+    pub fn spawn_on_scope<Fut: Future<Output = T> + 'scope + MaybeSend>(&self, f: Fut) {
-        let result = Rc::new(RefCell::new(None));
+        let result = ScopeResult::<T>::default();
         self.results.borrow_mut().push(result.clone());
         let f = async move {
             let temp_result = f.await;
+
+            #[cfg(feature = "std")]
             result.borrow_mut().replace(temp_result);
+
+            #[cfg(not(feature = "std"))]
+            {
+                let mut lock = result.lock();
+                *lock = Some(temp_result);
+            }
         };
         self.executor.spawn(f).detach();
     }
 }
+
+#[cfg(feature = "std")]
+mod send_sync_bounds {
+    pub trait MaybeSend {}
+    impl<T> MaybeSend for T {}
+
+    pub trait MaybeSync {}
+    impl<T> MaybeSync for T {}
+}
+
+#[cfg(not(feature = "std"))]
+mod send_sync_bounds {
+    pub trait MaybeSend: Send {}
+    impl<T: Send> MaybeSend for T {}
+
+    pub trait MaybeSync: Sync {}
+    impl<T: Sync> MaybeSync for T {}
+}
+
+use send_sync_bounds::{MaybeSend, MaybeSync};

crates/bevy_tasks/src/slice.rs

@@ -1,4 +1,5 @@
 use super::TaskPool;
+use alloc::vec::Vec;
 
 /// Provides functions for mapping read-only slices across a provided [`TaskPool`].
 pub trait ParallelSlice<T: Sync>: AsRef<[T]> {

crates/bevy_tasks/src/task.rs

@@ -14,11 +14,11 @@ use core::{
 /// Tasks that panic get immediately canceled. Awaiting a canceled task also causes a panic.
 #[derive(Debug)]
 #[must_use = "Tasks are canceled when dropped, use `.detach()` to run them in the background."]
-pub struct Task<T>(async_executor::Task<T>);
+pub struct Task<T>(crate::executor::Task<T>);
 
 impl<T> Task<T> {
     /// Creates a new task from a given `async_executor::Task`
-    pub fn new(task: async_executor::Task<T>) -> Self {
+    pub fn new(task: crate::executor::Task<T>) -> Self {
         Self(task)
     }
 

crates/bevy_tasks/src/task_pool.rs

@@ -2,7 +2,7 @@ use alloc::sync::Arc;
 use core::{future::Future, marker::PhantomData, mem, panic::AssertUnwindSafe};
 use std::thread::{self, JoinHandle};
 
-use async_executor::FallibleTask;
+use crate::executor::FallibleTask;
 use concurrent_queue::ConcurrentQueue;
 use futures_lite::FutureExt;
 
@@ -102,7 +102,7 @@ impl TaskPoolBuilder {
 #[derive(Debug)]
 pub struct TaskPool {
     /// The executor for the pool.
-    executor: Arc<async_executor::Executor<'static>>,
+    executor: Arc<crate::executor::Executor<'static>>,
 
     // The inner state of the pool.
     threads: Vec<JoinHandle<()>>,
@@ -111,7 +111,7 @@ pub struct TaskPool {
 
 impl TaskPool {
     thread_local! {
-        static LOCAL_EXECUTOR: async_executor::LocalExecutor<'static> = const { async_executor::LocalExecutor::new() };
+        static LOCAL_EXECUTOR: crate::executor::LocalExecutor<'static> = const { crate::executor::LocalExecutor::new() };
         static THREAD_EXECUTOR: Arc<ThreadExecutor<'static>> = Arc::new(ThreadExecutor::new());
     }
 
@@ -128,7 +128,7 @@ impl TaskPool {
     fn new_internal(builder: TaskPoolBuilder) -> Self {
         let (shutdown_tx, shutdown_rx) = async_channel::unbounded::<()>();
 
-        let executor = Arc::new(async_executor::Executor::new());
+        let executor = Arc::new(crate::executor::Executor::new());
 
         let num_threads = builder
             .num_threads
@@ -344,9 +344,9 @@ impl TaskPool {
         // transmute the lifetimes to 'env here to appease the compiler as it is unable to validate safety.
         // Any usages of the references passed into `Scope` must be accessed through
         // the transmuted reference for the rest of this function.
-        let executor: &async_executor::Executor = &self.executor;
+        let executor: &crate::executor::Executor = &self.executor;
         // SAFETY: As above, all futures must complete in this function so we can change the lifetime
-        let executor: &'env async_executor::Executor = unsafe { mem::transmute(executor) };
+        let executor: &'env crate::executor::Executor = unsafe { mem::transmute(executor) };
         // SAFETY: As above, all futures must complete in this function so we can change the lifetime
         let external_executor: &'env ThreadExecutor<'env> =
             unsafe { mem::transmute(external_executor) };
@@ -432,7 +432,7 @@ impl TaskPool {
 
     #[inline]
     async fn execute_global_external_scope<'scope, 'ticker, T>(
-        executor: &'scope async_executor::Executor<'scope>,
+        executor: &'scope crate::executor::Executor<'scope>,
         external_ticker: ThreadExecutorTicker<'scope, 'ticker>,
         scope_ticker: ThreadExecutorTicker<'scope, 'ticker>,
         get_results: impl Future<Output = Vec<T>>,
@@ -478,7 +478,7 @@ impl TaskPool {
 
     #[inline]
     async fn execute_global_scope<'scope, 'ticker, T>(
-        executor: &'scope async_executor::Executor<'scope>,
+        executor: &'scope crate::executor::Executor<'scope>,
         scope_ticker: ThreadExecutorTicker<'scope, 'ticker>,
         get_results: impl Future<Output = Vec<T>>,
     ) -> Vec<T> {
@@ -562,7 +562,7 @@ impl TaskPool {
     /// ```
     pub fn with_local_executor<F, R>(&self, f: F) -> R
     where
-        F: FnOnce(&async_executor::LocalExecutor) -> R,
+        F: FnOnce(&crate::executor::LocalExecutor) -> R,
     {
         Self::LOCAL_EXECUTOR.with(f)
     }
@@ -593,7 +593,7 @@ impl Drop for TaskPool {
 /// For more information, see [`TaskPool::scope`].
 #[derive(Debug)]
 pub struct Scope<'scope, 'env: 'scope, T> {
-    executor: &'scope async_executor::Executor<'scope>,
+    executor: &'scope crate::executor::Executor<'scope>,
     external_executor: &'scope ThreadExecutor<'scope>,
     scope_executor: &'scope ThreadExecutor<'scope>,
     spawned: &'scope ConcurrentQueue<FallibleTask<Result<T, Box<(dyn core::any::Any + Send)>>>>,

crates/bevy_tasks/src/thread_executor.rs

@@ -1,7 +1,7 @@
 use core::marker::PhantomData;
 use std::thread::{self, ThreadId};
 
-use async_executor::{Executor, Task};
+use crate::executor::{Executor, Task};
 use futures_lite::Future;
 
 /// An executor that can only be ticked on the thread it was instantiated on. But

crates/bevy_tasks/src/usages.rs

@@ -1,11 +1,20 @@
 use super::TaskPool;
 use core::ops::Deref;
+
+#[cfg(feature = "std")]
 use std::sync::OnceLock;
 
+#[cfg(not(feature = "std"))]
+use spin::Once;
+
 macro_rules! taskpool {
     ($(#[$attr:meta])* ($static:ident, $type:ident)) => {
+        #[cfg(feature = "std")]
         static $static: OnceLock<$type> = OnceLock::new();
 
+        #[cfg(not(feature = "std"))]
+        static $static: Once<$type> = Once::new();
+
         $(#[$attr])*
         #[derive(Debug)]
         pub struct $type(TaskPool);
@@ -13,9 +22,17 @@ macro_rules! taskpool {
         impl $type {
             #[doc = concat!(" Gets the global [`", stringify!($type), "`] instance, or initializes it with `f`.")]
             pub fn get_or_init(f: impl FnOnce() -> TaskPool) -> &'static Self {
+                #[cfg(feature = "std")]
+                {
                     $static.get_or_init(|| Self(f()))
                 }
 
+                #[cfg(not(feature = "std"))]
+                {
+                    $static.call_once(|| Self(f()))
+                }
+            }
+
             #[doc = concat!(" Attempts to get the global [`", stringify!($type), "`] instance, \
                 or returns `None` if it is not initialized.")]
             pub fn try_get() -> Option<&'static Self> {

tools/ci/src/commands/compile_check_no_std.rs

@@ -62,14 +62,6 @@ impl Prepare for CompileCheckNoStdCommand {
             "Please fix compiler errors in output above for bevy_mikktspace no_std compatibility.",
         ));
 
-        commands.push(PreparedCommand::new::<Self>(
-            cmd!(
-                sh,
-                "cargo check -p bevy_mikktspace --no-default-features --features libm --target {target}"
-            ),
-            "Please fix compiler errors in output above for bevy_mikktspace no_std compatibility.",
-        ));
-
         commands.push(PreparedCommand::new::<Self>(
             cmd!(
                 sh,
@@ -94,6 +86,14 @@ impl Prepare for CompileCheckNoStdCommand {
             "Please fix compiler errors in output above for bevy_color no_std compatibility.",
         ));
 
+        commands.push(PreparedCommand::new::<Self>(
+            cmd!(
+                sh,
+                "cargo check -p bevy_tasks --no-default-features --features edge_executor,critical-section --target {target}"
+            ),
+            "Please fix compiler errors in output above for bevy_tasks no_std compatibility.",
+        ));
+
         commands
     }
 }