mirror of
https://github.com/openai/codex.git
synced 2026-09-08 15:50:34 +00:00
refactor: drop data locks before async work
This commit is contained in:
@@ -13,6 +13,7 @@ use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use tokio::process::Command;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct BearerTokenRefresher {
|
||||
@@ -34,8 +35,26 @@ 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) = self.state.cached_token.lock().await.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(),
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
let _refresh_guard = self
|
||||
.state
|
||||
.refresh_lock
|
||||
.acquire()
|
||||
.await
|
||||
.map_err(|_| io::Error::other("external bearer token refresh semaphore closed"))?;
|
||||
{
|
||||
let 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,
|
||||
@@ -47,14 +66,14 @@ impl ExternalAuth for BearerTokenRefresher {
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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 access_token = run_provider_auth_command(&self.state.config).await?;
|
||||
let mut cached = self.state.cached_token.lock().await;
|
||||
*cached = Some(CachedExternalBearerToken {
|
||||
access_token: access_token.clone(),
|
||||
fetched_at: Instant::now(),
|
||||
});
|
||||
Ok(Some(ExternalAuthTokens::access_token_only(access_token)))
|
||||
}
|
||||
|
||||
@@ -62,6 +81,12 @@ impl ExternalAuth for BearerTokenRefresher {
|
||||
&self,
|
||||
_context: ExternalAuthRefreshContext,
|
||||
) -> io::Result<ExternalAuthTokens> {
|
||||
let _refresh_guard = self
|
||||
.state
|
||||
.refresh_lock
|
||||
.acquire()
|
||||
.await
|
||||
.map_err(|_| io::Error::other("external bearer token refresh semaphore closed"))?;
|
||||
let access_token = run_provider_auth_command(&self.state.config).await?;
|
||||
let mut cached = self.state.cached_token.lock().await;
|
||||
*cached = Some(CachedExternalBearerToken {
|
||||
@@ -82,6 +107,7 @@ impl fmt::Debug for BearerTokenRefresher {
|
||||
struct ExternalBearerAuthState {
|
||||
config: ModelProviderAuthInfo,
|
||||
cached_token: Mutex<Option<CachedExternalBearerToken>>,
|
||||
refresh_lock: Semaphore,
|
||||
}
|
||||
|
||||
impl ExternalBearerAuthState {
|
||||
@@ -89,6 +115,7 @@ impl ExternalBearerAuthState {
|
||||
Self {
|
||||
config,
|
||||
cached_token: Mutex::new(None),
|
||||
refresh_lock: Semaphore::new(1),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -420,32 +420,27 @@ 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()
|
||||
total, host, reason, decision, source, protocol, port, buffered
|
||||
);
|
||||
debug!("{violation_line}");
|
||||
drop(guard);
|
||||
|
||||
if let Some(observer) = blocked_request_observer {
|
||||
observer.on_blocked_request(blocked_for_observer).await;
|
||||
|
||||
Reference in New Issue
Block a user