Implement EventHubs token refresh (#2476)

* Added preliminary support to refresh tokens

* Convert AMQP to use async-trait package; Working token refresh logic

* Moved credentials to connection manager since they are inherently per-connection; removed credentials from authorize_path since they're not needed at that point

* Working connection manager token refresh tests


Co-authored-by: Heath Stewart <[email protected]>
Co-authored-by: Copilot <[email protected]>

* Final PR feedback

---------

Co-authored-by: Heath Stewart <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

Cargo.lock

@@ -15,6 +15,8 @@ pub mod fs;
 pub mod hmac;
 pub mod http;
 pub mod process;
+pub mod task;
+
 #[cfg(feature = "test")]
 pub mod test;
 

sdk/core/azure_core/src/lib.rs

@@ -21,9 +21,7 @@ impl Executor for StdExecutor {
             let output = cmd.output();
             tx.send(output)
         });
-        let output = rx
-            .await
-            .map_err(|err| io::Error::new(io::ErrorKind::Other, err))??;
+        let output = rx.await.map_err(io::Error::other)??;
         Ok(output)
     }
 }

sdk/core/azure_core/src/process/standard.rs

@@ -0,0 +1,139 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+//! Asynchronous task execution utilities.
+//!
+//! This module provides a mechanism to spawn tasks asynchronously and wait for their completion.
+//!
+//! It abstracts away the underlying implementation details, allowing for different task execution strategies based on the target architecture and features enabled.
+//!
+//!
+//! Example usage:
+//!
+//! ```
+//! use azure_core::task::{new_task_spawner, TaskSpawner};
+//! use futures::FutureExt;
+//!
+//! #[tokio::main]
+//! async fn main() {
+//!     let spawner = new_task_spawner();
+//!     let handle = spawner.spawn(async {
+//!         // Simulate some work
+//!         std::thread::sleep(std::time::Duration::from_secs(1));
+//!     }.boxed());
+//!
+//!     handle.await.expect("Task should complete successfully");
+//!
+//!     println!("Task completed");
+//! }
+//! ```
+//!
+//!
+use std::{fmt::Debug, future::Future, pin::Pin, sync::Arc};
+
+mod standard_spawn;
+
+#[cfg(feature = "tokio")]
+mod tokio_spawn;
+
+#[cfg(test)]
+mod tests;
+
+#[cfg(not(target_arch = "wasm32"))]
+pub(crate) type TaskFuture = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
+
+// WASM32 does not support `Send` futures, so we use a non-Send future type.
+#[cfg(target_arch = "wasm32")]
+pub(crate) type TaskFuture = Pin<Box<dyn Future<Output = ()> + 'static>>;
+
+/// A `SpawnedTask` is a future that represents a running task.
+/// It can be awaited to block until the task has completed.
+#[cfg(not(target_arch = "wasm32"))]
+pub type SpawnedTask = Pin<
+    Box<
+        dyn Future<Output = std::result::Result<(), Box<dyn std::error::Error + Send>>>
+            + Send
+            + 'static,
+    >,
+>;
+
+#[cfg(target_arch = "wasm32")]
+pub type SpawnedTask =
+    Pin<Box<dyn Future<Output = std::result::Result<(), Box<dyn std::error::Error>>> + 'static>>;
+
+/// An async command runner.
+///
+pub trait TaskSpawner: Send + Sync + Debug {
+    /// Spawn a task that executes a given future and returns the output.
+    ///
+    /// # Arguments
+    ///
+    /// * `f` - A future representing the task to be spawned. This future cannot capture any variables
+    ///   from its environment by reference, as it will be executed in a different thread or context.
+    ///
+    /// # Returns
+    /// A future which can be awaited to block until the task has completed.
+    ///
+    /// # Example
+    /// ```
+    /// use azure_core::task::{new_task_spawner, TaskSpawner};
+    /// use futures::FutureExt;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///   let spawner = new_task_spawner();
+    ///   let future = spawner.spawn(async {
+    ///     // Simulate some work
+    ///     std::thread::sleep(std::time::Duration::from_secs(1));
+    ///   }.boxed());
+    ///   future.await.expect("Task should complete successfully");
+    /// }
+    /// ```
+    ///
+    /// # Note
+    ///
+    /// This trait intentionally does not use the *`async_trait`* macro because when the
+    /// `async_trait` attribute is applied to a trait  implementation, the rewritten
+    /// method cannot directly return a future, instead they wrap the return value
+    /// in a future, and we want the `spawn` method to directly return a future
+    /// that can be awaited.
+    ///
+    fn spawn(&self, f: TaskFuture) -> SpawnedTask;
+}
+
+/// Creates a new [`TaskSpawner`] to enable running tasks asynchronously.
+///
+///
+/// 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.
+///
+/// # Returns
+///  A new instance of a [`TaskSpawner`] which can be used to spawn background tasks.
+///
+/// # Example
+///
+/// ```
+/// use azure_core::task::{new_task_spawner, TaskSpawner};
+/// use futures::FutureExt;
+///
+/// #[tokio::main]
+/// async fn main() {
+///   let spawner = new_task_spawner();
+///   let handle = spawner.spawn(async {
+///     // Simulate some work
+///     std::thread::sleep(std::time::Duration::from_secs(1));
+///   }.boxed());
+/// }
+/// ```
+///
+pub fn new_task_spawner() -> Arc<dyn TaskSpawner> {
+    #[cfg(not(feature = "tokio"))]
+    {
+        Arc::new(standard_spawn::StdSpawner)
+    }
+    #[cfg(feature = "tokio")]
+    {
+        Arc::new(tokio_spawn::TokioSpawner) as Arc<dyn TaskSpawner>
+    }
+}

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

@@ -0,0 +1,146 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+use super::{SpawnedTask, TaskFuture, TaskSpawner};
+#[cfg(not(target_arch = "wasm32"))]
+use futures::{executor::LocalPool, task::SpawnExt};
+#[cfg(not(target_arch = "wasm32"))]
+use std::{
+    future,
+    future::Future,
+    pin::Pin,
+    sync::{Arc, Mutex},
+    task::Waker,
+    task::{Context, Poll},
+    thread,
+};
+#[cfg(not(target_arch = "wasm32"))]
+use tracing::debug;
+
+/// A future that completes when a thread join handle completes.
+#[cfg(not(target_arch = "wasm32"))]
+struct ThreadJoinFuture {
+    join_state: Arc<Mutex<ThreadJoinState>>,
+}
+
+#[cfg(not(target_arch = "wasm32"))]
+#[derive(Default)]
+struct ThreadJoinState {
+    join_handle:
+        Option<thread::JoinHandle<std::result::Result<(), Box<dyn std::error::Error + Send>>>>,
+    waker: Option<Waker>,
+    thread_finished: bool,
+}
+
+#[cfg(not(target_arch = "wasm32"))]
+impl Future for ThreadJoinFuture {
+    type Output = std::result::Result<(), Box<dyn std::error::Error + Send>>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let mut join_state = self.join_state.lock().map_err(|e| {
+            debug!("Failed to lock join state: {}", e);
+            Box::new(crate::Error::message(
+                crate::error::ErrorKind::Other,
+                format!("Thread panicked: {:?}", e),
+            )) as Box<dyn std::error::Error + Send>
+        })?;
+
+        // Join handle is present, so we can check if the thread has finished
+        // and take the handle if it has.
+        // This is safe because we are holding the lock on the join state.
+        // We can safely take the handle and join it without blocking.
+        // This allows us to retrieve the terminal state of the thread.
+        if join_state.thread_finished {
+            // Thread is finished, so we can safely take the handle
+            let Some(join_handle) = join_state.join_handle.take() else {
+                // The join handle was already removed from the state, we know we're done.
+                return Poll::Ready(Ok(()));
+            };
+
+            // Since we know the thread is finished, we can safely take the handle
+            // and join it. This allows us to retrieve the terminal state of the thread.
+            //
+            // Technically this might block (because the `thread_finished` flag
+            // is set before the thread *actually* finishes), but it should be negligible.
+            match join_handle.join() {
+                Ok(_) => Poll::Ready(Ok(())),
+                Err(e) => Poll::Ready(Err(Box::new(crate::Error::message(
+                    crate::error::ErrorKind::Other,
+                    format!("Thread panicked: {:?}", e),
+                )) as Box<dyn std::error::Error + Send>)),
+            }
+        } else {
+            // Thread is still running, so we need to register the waker
+            // for when it completes.
+            join_state.waker = Some(cx.waker().clone());
+            Poll::Pending
+        }
+    }
+}
+
+/// A [`TaskSpawner`] using [`std::thread::spawn`].
+#[derive(Debug)]
+pub struct StdSpawner;
+
+impl TaskSpawner for StdSpawner {
+    #[cfg_attr(target_arch = "wasm32", allow(unused_variables))]
+    fn spawn(&self, f: TaskFuture) -> SpawnedTask {
+        #[cfg(target_arch = "wasm32")]
+        {
+            panic!("std::thread::spawn is not supported on wasm32")
+        }
+        #[cfg(not(target_arch = "wasm32"))]
+        {
+            let join_state = Arc::new(Mutex::new(ThreadJoinState::default()));
+            {
+                let Ok(mut js) = join_state.lock() else {
+                    return Box::pin(future::ready(Err(Box::new(crate::Error::message(
+                        crate::error::ErrorKind::Other,
+                        "Thread panicked.",
+                    ))
+                        as Box<dyn std::error::Error + Send>)));
+                };
+
+                // Clone the join state so it can be moved into the thread
+                // and used to notify the waker when the thread finishes.
+                let join_state_clone = join_state.clone();
+
+                js.join_handle = Some(thread::spawn(move || {
+                    // Create a local executor
+                    let mut local_pool = LocalPool::new();
+                    let spawner = local_pool.spawner();
+
+                    // Spawn the future on the local executor
+                    let Ok(future_handle) = spawner.spawn_with_handle(f) else {
+                        return Err(Box::new(crate::Error::message(
+                            crate::error::ErrorKind::Other,
+                            "Failed to spawn future.",
+                        ))
+                            as Box<dyn std::error::Error + Send>);
+                    };
+                    // Drive the executor until the future completes
+                    local_pool.run_until(future_handle);
+
+                    let Ok(mut join_state) = join_state_clone.lock() else {
+                        return Err(Box::new(crate::Error::message(
+                            crate::error::ErrorKind::Other,
+                            "Failed to lock join state",
+                        ))
+                            as Box<dyn std::error::Error + Send>);
+                    };
+
+                    // The thread has finished, so we can take the waker
+                    // and notify it.
+                    join_state.thread_finished = true;
+                    if let Some(waker) = join_state.waker.take() {
+                        waker.wake();
+                    }
+                    Ok(())
+                }));
+            }
+            // Create a future that will complete when the thread joins
+            let join_future = ThreadJoinFuture { join_state };
+            Box::pin(join_future)
+        }
+    }
+}