Create an AsyncRuntime trait (#2664)

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

* CI pipeline fixes

* wasm32 clippy fixes

* Fixed typo

* Minor test cleanup

* PR feedback

* PR feedback

* doctest fixes

* Added text to README.md describing how to set the async runtime used by the SDK

* Update sdk/core/azure_core/CHANGELOG.md

Co-authored-by: Heath Stewart <heaths@outlook.com>

* Policheck feedback

---------

Co-authored-by: Heath Stewart <heaths@outlook.com>

Author: LarryOsterman

sdk/core/azure_core/CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### Features Added
 
+- Added `get_async_runtime()` and `set_async_runtime()` to allow customers to replace
+the asynchronous runtime used by the Azure SDK.
+
 ### Breaking Changes
 
 ### Bugs Fixed
@@ -16,7 +19,7 @@
 
 - Added `#[safe]` attribute helper for `SafeDebug` derive macro to show or hide types and members as appropriate.
 - Added `Page` trait to facilitate the `ItemIterator`.
-- Added `PageIterator` to asynchronously iterator all pages.
+- Added `PageIterator` to asynchronously iterate all pages.
 
 ### Breaking Changes
 

sdk/core/azure_core/README.md

@@ -20,12 +20,12 @@ you can find the [package on crates.io][Package (crates.io)].
 
 The main shared concepts of `azure_core` - and Azure SDK libraries using `azure_core` - include:
 
-- Configuring service clients, e.g. configuring retries, logging (`ClientOptions`).
-- Accessing HTTP response details (`Response<T>`).
-- Paging and asynchronous streams (`Pager<T>`).
-- Errors from service requests in a consistent fashion. (`azure_core::Error`).
-- Customizing requests (`ClientOptions`).
-- Abstractions for representing Azure SDK credentials. (`TokenCredentials`).
+-   Configuring service clients, e.g. configuring retries, logging (`ClientOptions`).
+-   Accessing HTTP response details (`Response<T>`).
+-   Paging and asynchronous streams (`Pager<T>`).
+-   Errors from service requests in a consistent fashion. (`azure_core::Error`).
+-   Customizing requests (`ClientOptions`).
+-   Abstractions for representing Azure SDK credentials. (`TokenCredentials`).
 
 ### Thread safety
 
@@ -34,23 +34,25 @@ We guarantee that all client instance methods are thread-safe and independent of
 ### Additional concepts
 
 <!-- CLIENT COMMON BAR -->
+
 [Client options](#configuring-service-clients-using-clientoptions) |
 [Accessing the response](#accessing-http-response-details-using-responset) |
 [Handling Errors Results](#handling-errors-results) |
 [Consuming Service Methods Returning `Pager<T>`](#consuming-service-methods-returning-pagert)
+
 <!-- CLIENT COMMON BAR -->
 
 ## Features
 
-- `debug`: enables extra information for developers e.g., emitting all fields in `std::fmt::Debug` implementation.
-- `hmac_openssl`: configures HMAC using `openssl`.
-- `hmac_rust`: configures HMAC using pure Rust.
-- `reqwest` (default): enables and sets `reqwest` as the default `HttpClient`. Enables `reqwest`'s `native-tls` feature.
-- `reqwest_deflate` (default): enables deflate compression for `reqwest`.
-- `reqwest_gzip` (default): enables gzip compression for `reqwest`.
-- `reqwest_rustls`: enables `reqwest`'s `rustls-tls-native-roots-no-provider` feature,
-- `tokio`: enables and sets `tokio` as the default async runtime.
-- `xml`: enables XML support.
+-   `debug`: enables extra information for developers e.g., emitting all fields in `std::fmt::Debug` implementation.
+-   `hmac_openssl`: configures HMAC using `openssl`.
+-   `hmac_rust`: configures HMAC using pure Rust.
+-   `reqwest` (default): enables and sets `reqwest` as the default `HttpClient`. Enables `reqwest`'s `native-tls` feature.
+-   `reqwest_deflate` (default): enables deflate compression for `reqwest`.
+-   `reqwest_gzip` (default): enables gzip compression for `reqwest`.
+-   `reqwest_rustls`: enables `reqwest`'s `rustls-tls-native-roots-no-provider` feature,
+-   `tokio`: enables and sets `tokio` as the default async runtime.
+-   `xml`: enables XML support.
 
 ## Examples
 
@@ -244,6 +246,36 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 }
 ```
 
+### Replacing the async runtime
+
+Internally, the Azure SDK uses either the `tokio` async runtime (with the `tokio` feature), or it implements asynchronous functionality using functions in the `std` namespace.
+
+If your application uses a different asynchronous runtime, you can replace the asynchronous runtime used for internal functions by providing your own implementation of the `azure_core::async_runtime::AsyncRuntime` trait.
+
+You provide the implementation by calling the `set_async_runtime()` API:
+
+```rust no_run
+use azure_core::async_runtime::{
+     set_async_runtime, AsyncRuntime, TaskFuture, SpawnedTask};
+use std::sync::Arc;
+use futures::FutureExt;
+
+struct CustomRuntime;
+
+impl AsyncRuntime for CustomRuntime {
+    fn spawn(&self, f: TaskFuture) -> SpawnedTask {
+      unimplemented!("Custom spawn not implemented");
+    }
+    fn sleep(&self, duration: std::time::Duration) -> TaskFuture {
+      unimplemented!("Custom sleep not implemented");
+    }
+  }
+
+  set_async_runtime(Arc::new(CustomRuntime)).expect("Failed to set async runtime");
+```
+
+There can only be one async runtime set in a given process, so attempts to set the async runtime multiple times will fail.
+
 ## Troubleshooting
 
 ### Logging

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;
@@ -24,7 +23,7 @@ pub use constants::*;
 
 // Re-export modules in typespec_client_core such that azure_core-based crates don't need to reference it directly.
 pub use typespec_client_core::{
-    base64, create_enum, create_extensible_enum, date,
+    async_runtime, base64, create_enum, create_extensible_enum, date,
     error::{self, Error, Result},
     fmt, json, sleep, stream, Bytes, Uuid,
 };

sdk/core/azure_core/src/task/tokio_spawn.rs

@@ -5,10 +5,10 @@ use super::recoverable_connection::RecoverableConnection;
 use crate::error::{ErrorKind, EventHubsError};
 use async_lock::Mutex as AsyncMutex;
 use azure_core::{
+    async_runtime::{get_async_runtime, SpawnedTask},
     credentials::{AccessToken, TokenCredential},
     error::ErrorKind as AzureErrorKind,
     http::Url,
-    task::{new_task_spawner, SpawnedTask},
     Result,
 };
 use azure_core_amqp::AmqpClaimsBasedSecurityApis as _;
@@ -113,8 +113,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/eventhubs/azure_messaging_eventhubs/src/common/authorizer.rs

@@ -4,6 +4,9 @@
 
 ### Features Added
 
+-   Added `get_async_runtime()` and `set_async_runtime()` to allow customers to replace
+    the default asynchronous runtime with another.
+
 ### Breaking Changes
 
 ### Bugs Fixed
@@ -28,29 +31,29 @@
 
 ### Breaking Changes
 
-- The `reqwest_rustls` feature enables `rustls-tls-native-roots-no-provider` instead of `rustls-tls-native-roots` to remove the dependency on the `ring` crate.
+-   The `reqwest_rustls` feature enables `rustls-tls-native-roots-no-provider` instead of `rustls-tls-native-roots` to remove the dependency on the `ring` crate.
 
 ### Other Changes
 
-- Deriving `SafeDebug` formats non-exhaustive types by default. Enable `debug` feature to format normal `Debug` output.
-- Updated dependencies.
+-   Deriving `SafeDebug` formats non-exhaustive types by default. Enable `debug` feature to format normal `Debug` output.
+-   Updated dependencies.
 
 ## 0.2.0 (2025-04-08)
 
 ### Breaking Changes
 
-- Consolidated all the `tokio` features into a single feature named `tokio`. Traits remain separate but `tokio` support is enabled with a single feature.
-- Removed `Header` re-export from `http` module. It is still defined in the `http::headers` module.
-- Removed `http-types` dependency and implemented `Method` instead.
-- Removed `Pager`.
-- Removed `parsing` module.
+-   Consolidated all the `tokio` features into a single feature named `tokio`. Traits remain separate but `tokio` support is enabled with a single feature.
+-   Removed `Header` re-export from `http` module. It is still defined in the `http::headers` module.
+-   Removed `http-types` dependency and implemented `Method` instead.
+-   Removed `Pager`.
+-   Removed `parsing` module.
 
 ### Other Changes
 
-- Use `std::sync::LazyLock` added in rustc 1.80 instead of `once_cell::sync::Lazy`.
+-   Use `std::sync::LazyLock` added in rustc 1.80 instead of `once_cell::sync::Lazy`.
 
 ## 0.1.0 (2025-02-18)
 
 ### Features Added
 
-- Initial supported release.
+-   Initial supported release.