From 48497e258e8485402501bcc2773a5b8d1ba8154a Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 20 Apr 2026 14:58:01 -0700 Subject: [PATCH] refactor: drop data locks before async work --- codex-rs/login/src/auth/external_bearer.rs | 62 +++++++++++++++------- codex-rs/network-proxy/src/runtime.rs | 31 +++++------ 2 files changed, 56 insertions(+), 37 deletions(-) diff --git a/codex-rs/login/src/auth/external_bearer.rs b/codex-rs/login/src/auth/external_bearer.rs index c179f48619..e3f673a544 100644 --- a/codex-rs/login/src/auth/external_bearer.rs +++ b/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> { - 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 { + 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 { + 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 { 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> { + 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>, + refresh_lock: Semaphore, } impl ExternalBearerAuthState { @@ -89,6 +110,7 @@ impl ExternalBearerAuthState { Self { config, cached_token: Mutex::new(None), + refresh_lock: Semaphore::new(/*permits*/ 1), } } } diff --git a/codex-rs/network-proxy/src/runtime.rs b/codex-rs/network-proxy/src/runtime.rs index da090a69d4..daafeec737 100644 --- a/codex-rs/network-proxy/src/runtime.rs +++ b/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;