refactor: drop data locks before async work

codex-rs/login/src/auth/external_bearer.rs

@@ -13,6 +13,8 @@ use std::sync::Arc;
 use std::time::Instant;
 use tokio::process::Command;
 use tokio::sync::Mutex;
+use tokio::sync::Semaphore;
+use tokio::sync::SemaphorePermit;
 
 #[derive(Clone)]
 pub(crate) struct BearerTokenRefresher {
@@ -34,34 +36,44 @@ impl ExternalAuth for BearerTokenRefresher {
     }
 
     async fn resolve(&self) -> io::Result<Option<ExternalAuthTokens>> {
-        let access_token = {
-            let mut cached = self.state.cached_token.lock().await;
-            if let Some(cached_token) = cached.as_ref() {
-                let should_use_cached_token = match self.state.config.refresh_interval() {
-                    Some(refresh_interval) => cached_token.fetched_at.elapsed() < refresh_interval,
-                    None => true,
-                };
-                if should_use_cached_token {
-                    return Ok(Some(ExternalAuthTokens::access_token_only(
-                        cached_token.access_token.clone(),
-                    )));
-                }
-            }
+        if let Some(tokens) = self.fresh_cached_tokens().await {
+            return Ok(Some(tokens));
+        }
 
-            let access_token = run_provider_auth_command(&self.state.config).await?;
-            *cached = Some(CachedExternalBearerToken {
-                access_token: access_token.clone(),
-                fetched_at: Instant::now(),
-            });
-            access_token
+        let refresh_permit = self.acquire_refresh_permit().await?;
+        let tokens = if let Some(tokens) = self.fresh_cached_tokens().await {
+            tokens
+        } else {
+            self.refresh_tokens().await?
         };
-        Ok(Some(ExternalAuthTokens::access_token_only(access_token)))
+        drop(refresh_permit);
+        Ok(Some(tokens))
     }
 
     async fn refresh(
         &self,
         _context: ExternalAuthRefreshContext,
     ) -> io::Result<ExternalAuthTokens> {
+        let refresh_permit = self.acquire_refresh_permit().await?;
+        let tokens = self.refresh_tokens().await?;
+        drop(refresh_permit);
+        Ok(tokens)
+    }
+}
+
+impl BearerTokenRefresher {
+    async fn fresh_cached_tokens(&self) -> Option<ExternalAuthTokens> {
+        let cached = self.state.cached_token.lock().await;
+        let cached_token = cached.as_ref()?;
+        let should_use_cached_token = match self.state.config.refresh_interval() {
+            Some(refresh_interval) => cached_token.fetched_at.elapsed() < refresh_interval,
+            None => true,
+        };
+        should_use_cached_token
+            .then(|| ExternalAuthTokens::access_token_only(cached_token.access_token.clone()))
+    }
+
+    async fn refresh_tokens(&self) -> io::Result<ExternalAuthTokens> {
         let access_token = run_provider_auth_command(&self.state.config).await?;
         let mut cached = self.state.cached_token.lock().await;
         *cached = Some(CachedExternalBearerToken {
@@ -70,6 +82,14 @@ impl ExternalAuth for BearerTokenRefresher {
         });
         Ok(ExternalAuthTokens::access_token_only(access_token))
     }
+
+    async fn acquire_refresh_permit(&self) -> io::Result<SemaphorePermit<'_>> {
+        self.state
+            .refresh_lock
+            .acquire()
+            .await
+            .map_err(|_| io::Error::other("external bearer token refresh semaphore closed"))
+    }
 }
 
 impl fmt::Debug for BearerTokenRefresher {
@@ -82,6 +102,7 @@ impl fmt::Debug for BearerTokenRefresher {
 struct ExternalBearerAuthState {
     config: ModelProviderAuthInfo,
     cached_token: Mutex<Option<CachedExternalBearerToken>>,
+    refresh_lock: Semaphore,
 }
 
 impl ExternalBearerAuthState {
@@ -89,6 +110,7 @@ impl ExternalBearerAuthState {
         Self {
             config,
             cached_token: Mutex::new(None),
+            refresh_lock: Semaphore::new(/*permits*/ 1),
         }
     }
 }

codex-rs/network-proxy/src/runtime.rs

@@ -420,32 +420,29 @@ impl NetworkProxyState {
         let blocked_for_observer = entry.clone();
         let blocked_request_observer = self.blocked_request_observer.read().await.clone();
         let violation_line = blocked_request_violation_log_line(&entry);
-        let mut guard = self.state.write().await;
         let host = entry.host.clone();
         let reason = entry.reason.clone();
         let decision = entry.decision.clone();
         let source = entry.source.clone();
         let protocol = entry.protocol.clone();
         let port = entry.port;
-        guard.blocked.push_back(entry);
-        guard.blocked_total = guard.blocked_total.saturating_add(1);
-        let total = guard.blocked_total;
-        while guard.blocked.len() > MAX_BLOCKED_EVENTS {
-            guard.blocked.pop_front();
-        }
+        let (total, buffered) = {
+            let mut guard = self.state.write().await;
+            guard.blocked.push_back(entry);
+            guard.blocked_total = guard.blocked_total.saturating_add(1);
+            let total = guard.blocked_total;
+            while guard.blocked.len() > MAX_BLOCKED_EVENTS {
+                guard.blocked.pop_front();
+            }
+            (total, guard.blocked.len())
+        };
         debug!(
-            "recorded blocked request telemetry (total={}, host={}, reason={}, decision={:?}, source={:?}, protocol={}, port={:?}, buffered={})",
-            total,
-            host,
-            reason,
-            decision,
-            source,
-            protocol,
-            port,
-            guard.blocked.len()
+            "recorded blocked request telemetry (\
+             total={total}, host={host}, reason={reason}, \
+             decision={decision:?}, source={source:?}, \
+             protocol={protocol}, port={port:?}, buffered={buffered})"
         );
         debug!("{violation_line}");
-        drop(guard);
 
         if let Some(observer) = blocked_request_observer {
             observer.on_blocked_request(blocked_for_observer).await;