refactor: drop data locks before async work

This commit is contained in:
Michael Bolin
2026-04-17 15:21:04 -07:00
parent 29dd204720
commit 4820ebd984
2 changed files with 47 additions and 25 deletions

View File

@@ -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),
}
}
}

View File

@@ -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;