Use the permit in send_retry

apache · May 27, 2024 · ee03292 · ee03292
1 parent 8b3bbf0
commit ee03292
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/object_store/src/azure/credential.rs b/object_store/src/azure/credential.rs
@@ -942,7 +942,6 @@ mod tests {
     use super::*;
     use crate::azure::MicrosoftAzureBuilder;
     use crate::client::mock_server::MockServer;
-    use crate::RetryConfig;
     use crate::{ObjectStore, Path};
 
     #[tokio::test]

diff --git a/object_store/src/client/retry.rs b/object_store/src/client/retry.rs
@@ -401,7 +401,13 @@ impl RetryExt for reqwest::RequestBuilder {
 
     fn send_retry(self, ctx: &RequestContext) -> BoxFuture<'static, Result<Response>> {
         let request = self.retryable(&ctx.config);
-        Box::pin(async move { request.send().await })
+        let semaphore = Arc::clone(&ctx.semaphore);
+        Box::pin(async move {
+            let permit = semaphore.acquire_owned().await.unwrap();
+            let response = request.send().await;
+            drop(permit);
+            response
+        })
     }
 }