Fixed Azure#2506, creating an AsyncRuntime trait which can be used by customers to replace the async runtime

Author: LarryOsterman

sdk/core/azure_core/src/lib.rs

@@ -15,7 +15,6 @@ pub mod fs;
 pub mod hmac;
 pub mod http;
 pub mod process;
-pub mod task;
 
 #[cfg(feature = "test")]
 pub mod test;
@@ -31,3 +30,10 @@ pub use typespec_client_core::{
 
 #[cfg(feature = "xml")]
 pub use typespec_client_core::xml;
+
+pub use typespec_client_core::get_async_runtime;
+pub use typespec_client_core::set_async_runtime;
+
+pub mod async_runtime {
+    pub use typespec_client_core::async_runtime::SpawnedTask;
+}

sdk/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs

@@ -5,10 +5,11 @@ use super::recoverable_connection::RecoverableConnection;
 use crate::error::{ErrorKind, EventHubsError};
 use async_lock::Mutex as AsyncMutex;
 use azure_core::{
+    async_runtime::SpawnedTask,
     credentials::{AccessToken, TokenCredential},
     error::ErrorKind as AzureErrorKind,
+    get_async_runtime,
     http::Url,
-    task::{new_task_spawner, SpawnedTask},
     Result,
 };
 use azure_core_amqp::AmqpClaimsBasedSecurityApis as _;
@@ -113,8 +114,8 @@ impl Authorizer {
             self.authorization_refresher.get_or_init(|| {
                 debug!("Starting authorization refresh task.");
                 let self_clone = self.clone();
-                let spawner = new_task_spawner();
-                spawner.spawn(Box::pin(self_clone.refresh_tokens_task()))
+                let async_runtime = get_async_runtime();
+                async_runtime.spawn(Box::pin(self_clone.refresh_tokens_task()))
             });
         } else {
             debug!("Token already exists for path: {path}");

sdk/core/azure_core/src/task/mod.rs renamed to sdk/typespec/typespec_client_core/src/async_runtime/mod.rs

@@ -29,12 +29,17 @@
 //! ```
 //!
 //!
-use std::{fmt::Debug, future::Future, pin::Pin, sync::Arc};
+use std::{
+    fmt::Debug,
+    future::Future,
+    pin::Pin,
+    sync::{Arc, OnceLock},
+};
 
-mod standard_spawn;
+mod standard_runtime;
 
 #[cfg(feature = "tokio")]
-mod tokio_spawn;
+mod tokio_runtime;
 
 #[cfg(test)]
 mod tests;
@@ -61,9 +66,11 @@ pub type SpawnedTask = Pin<
 pub type SpawnedTask =
     Pin<Box<dyn Future<Output = std::result::Result<(), Box<dyn std::error::Error>>> + 'static>>;
 
-/// An async command runner.
+/// An Asynchronous Runtime.
 ///
-pub trait TaskSpawner: Send + Sync + Debug {
+/// This trait defines the various
+///
+pub trait AsyncRuntime: Send + Sync + Debug {
     /// Spawn a task that executes a given future and returns the output.
     ///
     /// # Arguments
@@ -76,13 +83,13 @@ pub trait TaskSpawner: Send + Sync + Debug {
     ///
     /// # Example
     /// ```
-    /// use azure_core::task::{new_task_spawner, TaskSpawner};
+    /// use azure_core::async_runtime::{get_async_runtime, TaskSpawner};
     /// use futures::FutureExt;
     ///
     /// #[tokio::main]
     /// async fn main() {
-    ///   let spawner = new_task_spawner();
-    ///   let future = spawner.spawn(async {
+    ///   let async_runtime = get_async_runtime();
+    ///   let future = async_runtime.spawn(async {
     ///     // Simulate some work
     ///     std::thread::sleep(std::time::Duration::from_secs(1));
     ///   }.boxed());
@@ -99,41 +106,86 @@ pub trait TaskSpawner: Send + Sync + Debug {
     /// that can be awaited.
     ///
     fn spawn(&self, f: TaskFuture) -> SpawnedTask;
+
+    fn sleep(&self, duration: std::time::Duration) -> TaskFuture;
 }
 
-/// Creates a new [`TaskSpawner`] to enable running tasks asynchronously.
+static ASYNC_RUNTIME_IMPLEMENTATION: OnceLock<Arc<dyn AsyncRuntime>> = OnceLock::new();
+
+/// Returns an [`AsyncRuntime`] to enable running operations which need to interact with an
+/// asynchronous runtime.
 ///
 ///
 /// The implementation depends on the target architecture and the features enabled:
-/// - If the `tokio` feature is enabled, it uses a tokio based spawner.
-/// - If the `tokio` feature is not enabled and the target architecture is not `wasm32`, it uses a std::thread based spawner.
+/// - If the `tokio` feature is enabled, it uses a tokio based spawner and timer.
+/// - If the `tokio` feature is not enabled and the target architecture is not `wasm32`, it uses a std::thread based spawner and timer.
 ///
 /// # Returns
-///  A new instance of a [`TaskSpawner`] which can be used to spawn background tasks.
+///  A new instance of a [`AsyncRuntime`] which can be used to spawn background tasks.
 ///
 /// # Example
 ///
 /// ```
-/// use azure_core::task::{new_task_spawner, TaskSpawner};
+/// use azure_core::get_async_runtime;
 /// use futures::FutureExt;
 ///
 /// #[tokio::main]
 /// async fn main() {
-///   let spawner = new_task_spawner();
-///   let handle = spawner.spawn(async {
+///   let async_runtime = get_async_runtime();
+///   let handle = async_runtime.spawn(async {
+///     // Simulate some work
+///     std::thread::sleep(std::time::Duration::from_secs(1));
+///   }.boxed());
+/// }
+/// ```
+///
+pub fn get_async_runtime() -> Arc<dyn AsyncRuntime> {
+    ASYNC_RUNTIME_IMPLEMENTATION
+        .get_or_init(|| create_async_runtime())
+        .clone()
+}
+
+/// Sets the current [`AsyncRuntime`] to enable running operations which need to interact with an
+/// asynchronous runtime.
+///
+///
+/// # Returns
+///  Ok if the async runtime was set successfully, or an error if it has already been set.
+///
+/// # Example
+///
+/// ```
+/// use azure_core::async_runtime::{get_async_runtime, AsyncRuntime};
+/// use futures::FutureExt;
+///
+/// async fn main() {
+///   let async_runtime = set_async_runtime();
+///   let handle = async_runtime.spawn(async {
 ///     // Simulate some work
 ///     std::thread::sleep(std::time::Duration::from_secs(1));
 ///   }.boxed());
 /// }
 /// ```
 ///
-pub fn new_task_spawner() -> Arc<dyn TaskSpawner> {
+pub fn set_async_runtime(runtime: Arc<dyn AsyncRuntime>) -> crate::Result<()> {
+    let result = ASYNC_RUNTIME_IMPLEMENTATION.set(runtime);
+    if result.is_err() {
+        Err(crate::Error::message(
+            crate::error::ErrorKind::Other,
+            "Async runtime has already been set.",
+        ))
+    } else {
+        Ok(())
+    }
+}
+
+fn create_async_runtime() -> Arc<dyn AsyncRuntime> {
     #[cfg(not(feature = "tokio"))]
     {
-        Arc::new(standard_spawn::StdSpawner)
+        Arc::new(standard_runtime::StdRuntime)
     }
     #[cfg(feature = "tokio")]
     {
-        Arc::new(tokio_spawn::TokioSpawner) as Arc<dyn TaskSpawner>
+        Arc::new(tokio_runtime::TokioRuntime) as Arc<dyn AsyncRuntime>
     }
 }

sdk/core/azure_core/src/task/standard_spawn.rs renamed to sdk/typespec/typespec_client_core/src/async_runtime/standard_runtime.rs

@@ -1,18 +1,21 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-use super::{SpawnedTask, TaskFuture, TaskSpawner};
+use super::{AsyncRuntime, SpawnedTask, TaskFuture};
+
 #[cfg(not(target_arch = "wasm32"))]
 use futures::{executor::LocalPool, task::SpawnExt};
+
+use std::sync::atomic::{AtomicBool, Ordering};
 #[cfg(not(target_arch = "wasm32"))]
 use std::{
     future,
     future::Future,
     pin::Pin,
     sync::{Arc, Mutex},
-    task::Waker,
-    task::{Context, Poll},
+    task::{Context, Poll, Waker},
     thread,
+    time::Duration,
 };
 #[cfg(not(target_arch = "wasm32"))]
 use tracing::debug;
@@ -78,11 +81,11 @@ impl Future for ThreadJoinFuture {
     }
 }
 
-/// A [`TaskSpawner`] using [`std::thread::spawn`].
+/// An [`AsyncRuntime`] using [`std::thread::spawn`].
 #[derive(Debug)]
-pub struct StdSpawner;
+pub struct StdRuntime;
 
-impl TaskSpawner for StdSpawner {
+impl AsyncRuntime for StdRuntime {
     #[cfg_attr(target_arch = "wasm32", allow(unused_variables))]
     fn spawn(&self, f: TaskFuture) -> SpawnedTask {
         #[cfg(target_arch = "wasm32")]
@@ -143,4 +146,46 @@ impl TaskSpawner for StdSpawner {
             Box::pin(join_future)
         }
     }
+
+    /// Creates a future that resolves after a specified duration of time.
+    ///
+    /// Uses a simple thread based implementation for sleep. A more efficient
+    /// implementation is available by using the `tokio` crate feature.
+    fn sleep(&self, duration: Duration) -> TaskFuture {
+        Box::pin(Sleep {
+            signal: None,
+            duration,
+        })
+    }
+}
+
+#[derive(Debug)]
+pub struct Sleep {
+    signal: Option<Arc<AtomicBool>>,
+    duration: Duration,
+}
+
+impl Future for Sleep {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        if let Some(signal) = &self.signal {
+            if signal.load(Ordering::Acquire) {
+                Poll::Ready(())
+            } else {
+                Poll::Pending
+            }
+        } else {
+            let signal = Arc::new(AtomicBool::new(false));
+            let waker = cx.waker().clone();
+            let duration = self.duration;
+            self.get_mut().signal = Some(signal.clone());
+            thread::spawn(move || {
+                thread::sleep(duration);
+                signal.store(true, Ordering::Release);
+                waker.wake();
+            });
+            Poll::Pending
+        }
+    }
 }