Yield in blob container client; read environment variables from env file (Azure#3245)

Author: LarryOsterman

sdk/core/azure_core/README.md

@@ -546,6 +546,9 @@ impl AsyncRuntime for CustomRuntime {
     fn sleep(&self, duration: Duration) -> TaskFuture {
       unimplemented!("Custom sleep not implemented");
     }
+    fn yield_now(&self) -> TaskFuture {
+        unimplemented!("Custom yield not implemented");
+    }
   }
 
   set_async_runtime(Arc::new(CustomRuntime)).expect("Failed to set async runtime");

sdk/core/typespec_client_core/src/async_runtime/mod.rs

@@ -115,6 +115,9 @@ pub trait AsyncRuntime: Send + Sync {
     /// # Returns
     /// A future that resolves after the specified duration has elapsed.
     fn sleep(&self, duration: Duration) -> TaskFuture;
+
+    /// Yield the current task back to the runtime scheduler.
+    fn yield_now(&self) -> TaskFuture;
 }
 
 static ASYNC_RUNTIME_IMPLEMENTATION: OnceLock<Arc<dyn AsyncRuntime>> = OnceLock::new();
@@ -176,6 +179,9 @@ pub fn get_async_runtime() -> Arc<dyn AsyncRuntime> {
 ///    fn sleep(&self, duration: typespec_client_core::time::Duration) -> TaskFuture {
 ///      unimplemented!("Custom sleep not implemented");
 ///    }
+///    fn yield_now(&self) -> TaskFuture {
+///      unimplemented!("Custom yield not implemented");
+///    }
 ///  }
 ///
 ///  set_async_runtime(Arc::new(CustomRuntime)).expect("Failed to set async runtime");

sdk/core/typespec_client_core/src/async_runtime/standard_runtime.rs

@@ -161,6 +161,18 @@ impl AsyncRuntime for StdRuntime {
             duration,
         })
     }
+
+    fn yield_now(&self) -> TaskFuture {
+        #[cfg(target_arch = "wasm32")]
+        {
+            panic!("yield_now is not supported on wasm32")
+        }
+        #[cfg(not(target_arch = "wasm32"))]
+        {
+            std::thread::yield_now();
+            Box::pin(future::ready(()))
+        }
+    }
 }
 
 #[cfg(not(target_arch = "wasm32"))]

sdk/core/typespec_client_core/src/async_runtime/tests.rs

@@ -260,6 +260,10 @@ impl AsyncRuntime for TestRuntime {
     fn sleep(&self, _duration: Duration) -> TaskFuture {
         unimplemented!("TestRuntime does not support sleeping");
     }
+
+    fn yield_now(&self) -> TaskFuture {
+        unimplemented!("TestRuntime does not support yielding");
+    }
 }
 
 // This test is ignored because by default, cargo test runs all tests in parallel, but

sdk/core/typespec_client_core/src/async_runtime/tokio_runtime.rs

@@ -25,4 +25,10 @@ impl AsyncRuntime for TokioRuntime {
                 .expect("Failed to convert duration to tokio format"),
         ))
     }
+
+    fn yield_now(&self) -> TaskFuture {
+        Box::pin(async {
+            tokio::task::yield_now().await;
+        })
+    }
 }

sdk/storage/azure_storage_blob/perf/list_blob_test.rs

@@ -13,7 +13,7 @@ use futures::{FutureExt, TryStreamExt};
 
 pub struct ListBlobTest {
     count: u32,
-    endpoint: String,
+    endpoint: Option<String>,
     client: OnceLock<BlobContainerClient>,
 }
 
@@ -27,19 +27,10 @@ impl ListBlobTest {
             println!("Parsed count: {}", count);
 
             let endpoint: Option<&String> = runner.try_get_test_arg("endpoint")?;
-            let endpoint = match endpoint {
-                Some(e) => e.clone(),
-                None => format!(
-                    "https://{}.blob.core.windows.net",
-                    std::env::var("AZURE_STORAGE_ACCOUNT_NAME")
-                        .expect("AZURE_STORAGE_ACCOUNT_NAME is not set")
-                ),
-            };
-            println!("Using endpoint: {}", endpoint);
 
             Ok(Box::new(ListBlobTest {
                 count,
-                endpoint,
+                endpoint: endpoint.cloned(),
                 client: OnceLock::new(),
             }) as Box<dyn PerfTest>)
         }
@@ -83,7 +74,15 @@ impl PerfTest for ListBlobTest {
         let recording = context.recording();
         let credential = recording.credential();
         let container_name = format!("perf-container-{}", uuid::Uuid::new_v4());
-        let client = BlobContainerClient::new(&self.endpoint, container_name, credential, None)?;
+        let endpoint = match &self.endpoint {
+            Some(e) => e.clone(),
+            None => format!(
+                "https://{}.blob.core.windows.net",
+                recording.var("AZURE_STORAGE_ACCOUNT_NAME", None)
+            ),
+        };
+        println!("Using endpoint: {}", endpoint);
+        let client = BlobContainerClient::new(&endpoint, container_name, credential, None)?;
         self.client.set(client).map_err(|_| {
             azure_core::Error::with_message(ErrorKind::Other, "Failed to set client")
         })?;