fix: retain Workdir close confirmation across retries
This commit is contained in:
@@ -556,6 +556,7 @@ mod client {
|
||||
capabilities: WorkdirSessionCapabilities,
|
||||
environment: crate::WorkdirSessionEnvironment,
|
||||
closed: AtomicBool,
|
||||
close_lock: tokio::sync::Mutex<()>,
|
||||
}
|
||||
|
||||
impl RemoteWorkdirSession {
|
||||
@@ -616,6 +617,7 @@ mod client {
|
||||
capabilities: opened.capabilities,
|
||||
environment: opened.environment,
|
||||
closed: AtomicBool::new(false),
|
||||
close_lock: tokio::sync::Mutex::new(()),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1014,7 +1016,11 @@ mod client {
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), WorkdirError> {
|
||||
if self.closed.swap(true, Ordering::AcqRel) {
|
||||
// In-flight is not closed. Dropping this guard on error or cancellation
|
||||
// leaves confirmation unset, so the next caller can retry Runtime's
|
||||
// idempotent DELETE even when the previous response was lost.
|
||||
let _close = self.close_lock.lock().await;
|
||||
if self.closed.load(Ordering::Acquire) {
|
||||
return Ok(());
|
||||
}
|
||||
let url = endpoint(
|
||||
@@ -1033,6 +1039,7 @@ mod client {
|
||||
.await
|
||||
.map_err(http_unavailable)?;
|
||||
if response.status() == StatusCode::NO_CONTENT || response.status().is_success() {
|
||||
self.closed.store(true, Ordering::Release);
|
||||
Ok(())
|
||||
} else {
|
||||
Err(decode_error(response).await)
|
||||
@@ -1211,7 +1218,7 @@ mod client {
|
||||
|
||||
// Seed an already-open session to exercise the public operation and
|
||||
// close boundaries independently of the rejected open request.
|
||||
let remote = RemoteWorkdirSession {
|
||||
let mut remote = RemoteWorkdirSession {
|
||||
client,
|
||||
base_url,
|
||||
authorization: authorization.clone(),
|
||||
@@ -1222,6 +1229,7 @@ mod client {
|
||||
WorkdirSessionCapabilities::EMPTY,
|
||||
),
|
||||
closed: AtomicBool::new(false),
|
||||
close_lock: tokio::sync::Mutex::new(()),
|
||||
};
|
||||
let stat = StatRequest {
|
||||
path: crate::WorkdirPath::new("file").unwrap(),
|
||||
@@ -1263,6 +1271,15 @@ mod client {
|
||||
listener.accept().unwrap_err().kind(),
|
||||
std::io::ErrorKind::WouldBlock
|
||||
);
|
||||
// An invalid endpoint cannot be produced by a successful open, but
|
||||
// must not turn a pre-dispatch error into a confirmed close either.
|
||||
remote.base_url = Url::parse("data:invalid-base").unwrap();
|
||||
for _ in 0..2 {
|
||||
assert!(matches!(
|
||||
remote.close().await,
|
||||
Err(WorkdirError::InvalidArgument(_))
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,287 @@
|
||||
#![cfg(feature = "http-client")]
|
||||
//! Remote close protocol tests with a loopback HTTP boundary. Requests and responses
|
||||
//! are explicitly gated; Runtime lease integration belongs to workspace-server.
|
||||
use std::{
|
||||
future::Future,
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
},
|
||||
task::Poll,
|
||||
time::Duration,
|
||||
};
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::{TcpListener, TcpStream},
|
||||
sync::mpsc,
|
||||
};
|
||||
use workdir::{
|
||||
Workdir, WorkdirError, WorkdirSession,
|
||||
http::{
|
||||
OpenWorkdirSessionRequest, OpenWorkdirSessionResponse, RemoteWorkdirSession,
|
||||
WorkdirHttpAuthorization, WorkdirHttpAuthorizationFuture, WorkdirSessionId,
|
||||
},
|
||||
};
|
||||
|
||||
struct Provider {
|
||||
url: reqwest::Url,
|
||||
requests: mpsc::UnboundedReceiver<TcpStream>,
|
||||
task: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
impl Drop for Provider {
|
||||
fn drop(&mut self) {
|
||||
self.task.abort();
|
||||
}
|
||||
}
|
||||
impl Provider {
|
||||
async fn new() -> Self {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let url = format!("http://{}", listener.local_addr().unwrap())
|
||||
.parse()
|
||||
.unwrap();
|
||||
let (tx, requests) = mpsc::unbounded_channel();
|
||||
let task = tokio::spawn(async move {
|
||||
loop {
|
||||
let (mut stream, _) = listener.accept().await.unwrap();
|
||||
let mut bytes = Vec::new();
|
||||
while !bytes.ends_with(b"\r\n\r\n") {
|
||||
bytes.push(stream.read_u8().await.unwrap());
|
||||
assert!(bytes.len() < 16384);
|
||||
}
|
||||
let headers = String::from_utf8(bytes).unwrap();
|
||||
let length = headers
|
||||
.lines()
|
||||
.find_map(|line| {
|
||||
let (name, value) = line.split_once(':')?;
|
||||
name.eq_ignore_ascii_case("content-length")
|
||||
.then(|| value.trim().parse::<usize>().unwrap())
|
||||
})
|
||||
.unwrap_or(0);
|
||||
assert!(length < 16384);
|
||||
stream.read_exact(&mut vec![0; length]).await.unwrap();
|
||||
assert!(
|
||||
headers.starts_with("POST /v1/working-directories/w/sessions ")
|
||||
|| headers.starts_with("DELETE /v1/workdir-sessions/s ")
|
||||
);
|
||||
if tx.send(stream).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
Self {
|
||||
url,
|
||||
requests,
|
||||
task,
|
||||
}
|
||||
}
|
||||
async fn request_while<F: Future>(&mut self, future: std::pin::Pin<&mut F>) -> TcpStream {
|
||||
tokio::time::timeout(Duration::from_secs(5), async {
|
||||
tokio::select! {
|
||||
request = self.requests.recv() => request.expect("provider stopped"),
|
||||
_ = future => panic!("close/open completed without the expected HTTP request"),
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("request timed out")
|
||||
}
|
||||
async fn open(&mut self, auth: Arc<dyn WorkdirHttpAuthorization>) -> RemoteWorkdirSession {
|
||||
let future = RemoteWorkdirSession::open_with_authorization(
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.timeout(Duration::from_secs(5))
|
||||
.build()
|
||||
.unwrap(),
|
||||
self.url.clone(),
|
||||
auth,
|
||||
Workdir::new("w").id().clone(),
|
||||
OpenWorkdirSessionRequest::default(),
|
||||
);
|
||||
tokio::pin!(future);
|
||||
let stream = self.request_while(future.as_mut()).await;
|
||||
let body = serde_json::to_string(&OpenWorkdirSessionResponse {
|
||||
session_id: WorkdirSessionId::new("s").unwrap(),
|
||||
workdir_id: Workdir::new("w").id().clone(),
|
||||
capabilities: workdir::WorkdirSessionCapabilities::ALL,
|
||||
environment: workdir::WorkdirSessionEnvironment::root(
|
||||
workdir::WorkdirSessionCapabilities::ALL,
|
||||
),
|
||||
})
|
||||
.unwrap();
|
||||
reply(stream, "200 OK", &body).await;
|
||||
future.await.unwrap()
|
||||
}
|
||||
}
|
||||
async fn reply(mut stream: TcpStream, status: &str, body: &str) {
|
||||
stream.write_all(format!("HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
|
||||
}
|
||||
#[derive(Debug, Default)]
|
||||
struct Authorization {
|
||||
reject_delete: AtomicBool,
|
||||
}
|
||||
impl WorkdirHttpAuthorization for Authorization {
|
||||
fn bearer_token(
|
||||
&self,
|
||||
method: String,
|
||||
_: String,
|
||||
_: Vec<u8>,
|
||||
) -> WorkdirHttpAuthorizationFuture {
|
||||
let reject = method == "DELETE" && self.reject_delete.swap(false, Ordering::SeqCst);
|
||||
Box::pin(async move {
|
||||
if reject {
|
||||
Err(WorkdirError::Unavailable(
|
||||
"fixture authorization rejected".into(),
|
||||
))
|
||||
} else {
|
||||
Ok("fixture-token".into())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn close_authorization_failure_retries_same_session_and_confirmed_close_skips_http() {
|
||||
let mut provider = Provider::new().await;
|
||||
let auth = Arc::new(Authorization::default());
|
||||
let remote = provider.open(auth.clone()).await;
|
||||
auth.reject_delete.store(true, Ordering::SeqCst);
|
||||
assert!(remote.close().await.is_err());
|
||||
assert!(provider.requests.try_recv().is_err());
|
||||
let retry = remote.close();
|
||||
tokio::pin!(retry);
|
||||
let request = provider.request_while(retry.as_mut()).await;
|
||||
reply(request, "204 No Content", "").await;
|
||||
retry.await.unwrap();
|
||||
remote.close().await.unwrap();
|
||||
assert!(provider.requests.try_recv().is_err());
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct GatedAuthorization {
|
||||
entered: mpsc::UnboundedSender<()>,
|
||||
release: std::sync::Mutex<Option<tokio::sync::oneshot::Receiver<()>>>,
|
||||
}
|
||||
impl WorkdirHttpAuthorization for GatedAuthorization {
|
||||
fn bearer_token(
|
||||
&self,
|
||||
method: String,
|
||||
_: String,
|
||||
_: Vec<u8>,
|
||||
) -> WorkdirHttpAuthorizationFuture {
|
||||
let gate = if method == "DELETE" {
|
||||
self.release.lock().unwrap().take()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let entered = self.entered.clone();
|
||||
Box::pin(async move {
|
||||
if let Some(gate) = gate {
|
||||
entered.send(()).unwrap();
|
||||
gate.await.unwrap();
|
||||
}
|
||||
Ok("fixture-token".into())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn close_cancel_during_authorization_allows_an_actual_retry() {
|
||||
let mut provider = Provider::new().await;
|
||||
let (entered, mut events) = mpsc::unbounded_channel();
|
||||
let (_release, gate) = tokio::sync::oneshot::channel();
|
||||
let remote = provider
|
||||
.open(Arc::new(GatedAuthorization {
|
||||
entered,
|
||||
release: std::sync::Mutex::new(Some(gate)),
|
||||
}))
|
||||
.await;
|
||||
let mut first = Box::pin(remote.close());
|
||||
tokio::time::timeout(Duration::from_secs(5), async {
|
||||
tokio::select! {
|
||||
_ = events.recv() => (),
|
||||
_ = &mut first => panic!("close completed while authorization was gated"),
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(first);
|
||||
assert!(provider.requests.try_recv().is_err());
|
||||
let retry = remote.close();
|
||||
tokio::pin!(retry);
|
||||
let request = provider.request_while(retry.as_mut()).await;
|
||||
reply(request, "204 No Content", "").await;
|
||||
retry.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn close_http_rejection_or_lost_response_requires_another_delete() {
|
||||
for response in [
|
||||
Some((
|
||||
"503 Service Unavailable",
|
||||
r#"{"code":"unavailable","message":"fixture"}"#,
|
||||
)),
|
||||
None,
|
||||
] {
|
||||
let mut provider = Provider::new().await;
|
||||
let remote = provider.open(Arc::new(Authorization::default())).await;
|
||||
let first = remote.close();
|
||||
tokio::pin!(first);
|
||||
let request = provider.request_while(first.as_mut()).await;
|
||||
if let Some((status, body)) = response {
|
||||
reply(request, status, body).await;
|
||||
} else {
|
||||
drop(request);
|
||||
} // Runtime may have deleted; no response confirms it.
|
||||
assert!(first.await.is_err());
|
||||
let retry = remote.close();
|
||||
tokio::pin!(retry);
|
||||
let request = provider.request_while(retry.as_mut()).await;
|
||||
reply(request, "204 No Content", "").await;
|
||||
retry.await.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn close_cancel_releases_serialization_and_waiter_must_confirm_delete() {
|
||||
let mut provider = Provider::new().await;
|
||||
let remote = provider.open(Arc::new(Authorization::default())).await;
|
||||
let mut first = Box::pin(remote.close());
|
||||
let request = provider.request_while(first.as_mut()).await;
|
||||
let mut waiter = Box::pin(remote.close());
|
||||
assert!(
|
||||
std::future::poll_fn(|cx| Poll::Ready(waiter.as_mut().poll(cx)))
|
||||
.await
|
||||
.is_pending()
|
||||
);
|
||||
drop(first);
|
||||
drop(request);
|
||||
let request = provider.request_while(waiter.as_mut()).await;
|
||||
reply(request, "204 No Content", "").await;
|
||||
waiter.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn concurrent_close_waits_for_confirmation_and_retries_after_failure() {
|
||||
for success in [false, true] {
|
||||
let mut provider = Provider::new().await;
|
||||
let remote = provider.open(Arc::new(Authorization::default())).await;
|
||||
let mut first = Box::pin(remote.close());
|
||||
let request = provider.request_while(first.as_mut()).await;
|
||||
let mut waiter = Box::pin(remote.close());
|
||||
assert!(
|
||||
std::future::poll_fn(|cx| Poll::Ready(waiter.as_mut().poll(cx)))
|
||||
.await
|
||||
.is_pending()
|
||||
);
|
||||
if success {
|
||||
reply(request, "204 No Content", "").await;
|
||||
first.await.unwrap();
|
||||
} else {
|
||||
drop(request);
|
||||
assert!(first.await.is_err());
|
||||
let request = provider.request_while(waiter.as_mut()).await;
|
||||
reply(request, "204 No Content", "").await;
|
||||
}
|
||||
waiter.await.unwrap();
|
||||
assert!(provider.requests.try_recv().is_err());
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,8 @@
|
||||
//! filesystem/Git inspection and unlink; conflicting target use is rejected promptly.
|
||||
|
||||
mod runtime_management_api;
|
||||
#[cfg(test)]
|
||||
mod workdir_close_tests;
|
||||
|
||||
use crate::auth::{RuntimeAuthContext, new_token_id, unix_now_seconds};
|
||||
use crate::catalog::{
|
||||
@@ -373,9 +375,12 @@ pub struct WorkspaceRuntimeHttpAuth {
|
||||
pub runtime_id: String,
|
||||
}
|
||||
|
||||
type RuntimeHttpWorkdirCloseResult = Result<(), RuntimeHttpWorkdirError>;
|
||||
|
||||
struct RuntimeHttpWorkdirSession {
|
||||
owner: RuntimeWorkspaceScope,
|
||||
session: WorkdirSessionHandle,
|
||||
close_attempt: Option<tokio::sync::watch::Receiver<Option<RuntimeHttpWorkdirCloseResult>>>,
|
||||
}
|
||||
|
||||
/// `GET /v1/runtime` response.
|
||||
@@ -878,7 +883,11 @@ async fn open_workdir_session(
|
||||
.map_err(|_| RuntimeHttpWorkdirError::internal())?
|
||||
.insert(
|
||||
session_id.as_str().to_string(),
|
||||
RuntimeHttpWorkdirSession { owner, session },
|
||||
RuntimeHttpWorkdirSession {
|
||||
owner,
|
||||
session,
|
||||
close_attempt: None,
|
||||
},
|
||||
);
|
||||
Ok(Json(response))
|
||||
}
|
||||
@@ -898,7 +907,7 @@ async fn run_workdir_session_operation(
|
||||
.map_err(|_| RuntimeHttpWorkdirError::internal())?;
|
||||
let record = sessions
|
||||
.get(&session_id)
|
||||
.filter(|record| record.owner == owner)
|
||||
.filter(|record| record.owner == owner && record.close_attempt.is_none())
|
||||
.ok_or_else(RuntimeHttpWorkdirError::not_found)?;
|
||||
record.session.clone()
|
||||
};
|
||||
@@ -934,24 +943,55 @@ async fn close_workdir_session(
|
||||
auth: Option<Extension<RuntimeAuthContext>>,
|
||||
) -> Result<StatusCode, RuntimeHttpWorkdirError> {
|
||||
let owner = required_workdir_owner(auth)?;
|
||||
let session = {
|
||||
let mut completion = {
|
||||
let mut sessions = state
|
||||
.workdir_sessions
|
||||
.lock()
|
||||
.map_err(|_| RuntimeHttpWorkdirError::internal())?;
|
||||
if sessions
|
||||
.get(&session_id)
|
||||
.is_some_and(|record| record.owner == owner)
|
||||
{
|
||||
sessions.remove(&session_id).map(|record| record.session)
|
||||
let Some(record) = sessions
|
||||
.get_mut(&session_id)
|
||||
.filter(|record| record.owner == owner)
|
||||
else {
|
||||
return Ok(StatusCode::NO_CONTENT);
|
||||
};
|
||||
if let Some(completion) = &record.close_attempt {
|
||||
completion.clone()
|
||||
} else {
|
||||
None
|
||||
let (result_tx, completion) = tokio::sync::watch::channel(None);
|
||||
record.close_attempt = Some(completion.clone());
|
||||
let session = record.session.clone();
|
||||
let registry = state.workdir_sessions.clone();
|
||||
// The task, not an HTTP waiter, owns provider teardown. In particular,
|
||||
// Local close marks itself closed before awaiting command teardown;
|
||||
// cancelling that future would make a later provider close a false ack.
|
||||
tokio::spawn(async move {
|
||||
let result = session.close().await.map_err(RuntimeHttpWorkdirError::from);
|
||||
let result = match registry.lock() {
|
||||
Ok(mut sessions) => {
|
||||
if result.is_ok() {
|
||||
sessions.remove(&session_id);
|
||||
} else if let Some(record) = sessions.get_mut(&session_id) {
|
||||
record.close_attempt = None;
|
||||
}
|
||||
result
|
||||
}
|
||||
Err(_) => Err(RuntimeHttpWorkdirError::internal()),
|
||||
};
|
||||
result_tx.send_replace(Some(result));
|
||||
});
|
||||
completion
|
||||
}
|
||||
};
|
||||
if let Some(session) = session {
|
||||
session.close().await?;
|
||||
loop {
|
||||
if let Some(result) = completion.borrow_and_update().clone() {
|
||||
result?;
|
||||
return Ok(StatusCode::NO_CONTENT);
|
||||
}
|
||||
completion
|
||||
.changed()
|
||||
.await
|
||||
.map_err(|_| RuntimeHttpWorkdirError::internal())?;
|
||||
}
|
||||
Ok(StatusCode::NO_CONTENT)
|
||||
}
|
||||
|
||||
fn required_workdir_owner(
|
||||
@@ -2149,7 +2189,7 @@ fn required_runtime_permission(method: &Method, path: &str) -> Option<&'static s
|
||||
None
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
#[derive(Clone, Debug)]
|
||||
struct RuntimeHttpWorkdirError {
|
||||
status: StatusCode,
|
||||
payload: WorkdirTransportError,
|
||||
@@ -2890,6 +2930,7 @@ mod tests {
|
||||
RuntimeHttpWorkdirSession {
|
||||
owner: RuntimeWorkspaceScope::new("workspace-a", "https://backend.test"),
|
||||
session,
|
||||
close_attempt: None,
|
||||
},
|
||||
)]))),
|
||||
ssh_keyscan_program: Arc::new(PathBuf::from("ssh-keyscan")),
|
||||
@@ -3194,6 +3235,7 @@ mod tests {
|
||||
RuntimeHttpWorkdirSession {
|
||||
owner: owner.clone(),
|
||||
session: session.clone(),
|
||||
close_attempt: None,
|
||||
},
|
||||
)]))),
|
||||
ssh_keyscan_program: Arc::new(PathBuf::from("ssh-keyscan")),
|
||||
@@ -3314,6 +3356,41 @@ mod tests {
|
||||
.expect_err("cross-workspace session access must fail");
|
||||
assert_eq!(error.status, StatusCode::NOT_FOUND);
|
||||
|
||||
for unauthorized in [
|
||||
None,
|
||||
Some(Extension(RuntimeAuthContext {
|
||||
workspace_id: "workspace-b".into(),
|
||||
..auth.clone()
|
||||
})),
|
||||
Some(Extension(RuntimeAuthContext {
|
||||
server_id: "server-b".into(),
|
||||
..auth.clone()
|
||||
})),
|
||||
] {
|
||||
let missing_auth = unauthorized.is_none();
|
||||
let result =
|
||||
close_workdir_session(State(state.clone()), Path("session-1".into()), unauthorized)
|
||||
.await;
|
||||
if missing_auth {
|
||||
assert_eq!(result.unwrap_err().status, StatusCode::FORBIDDEN);
|
||||
} else {
|
||||
assert_eq!(result.unwrap(), StatusCode::NO_CONTENT);
|
||||
}
|
||||
assert!(
|
||||
state
|
||||
.workdir_sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.contains_key("session-1")
|
||||
);
|
||||
session
|
||||
.stat(StatRequest {
|
||||
path: WorkdirPath::new("hello.txt").unwrap(),
|
||||
})
|
||||
.await
|
||||
.expect("unauthorized close must not reach provider");
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
close_workdir_session(
|
||||
State(state.clone()),
|
||||
|
||||
@@ -99,6 +99,7 @@ impl Harness {
|
||||
RuntimeHttpWorkdirSession {
|
||||
owner: RuntimeWorkspaceScope::new("workspace-a", "https://backend.test"),
|
||||
session: session.clone(),
|
||||
close_attempt: None,
|
||||
},
|
||||
)]))),
|
||||
ssh_keyscan_program: Arc::new(PathBuf::from("ssh-keyscan")),
|
||||
|
||||
@@ -0,0 +1,415 @@
|
||||
//! Runtime HTTP close integration: production local-token authentication, session registry,
|
||||
//! Local command teardown, and Runtime active-use guard. The execution-backend mapping is
|
||||
//! a fixture, not a materializer; a provider wrapper gates a genuinely pending Local close.
|
||||
//! Run with `cargo test -p worker-runtime --lib http_server::workdir_close_tests`.
|
||||
|
||||
use super::*;
|
||||
use crate::RuntimeOptions;
|
||||
use crate::execution::{
|
||||
WorkerExecutionBackend, WorkerExecutionResult, WorkerExecutionSpawnRequest,
|
||||
WorkerExecutionSpawnResult,
|
||||
};
|
||||
use crate::working_directory::WorkingDirectoryDiagnostic;
|
||||
use axum::body::to_bytes;
|
||||
use std::future::Future;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::task::Poll;
|
||||
use tokio::sync::{Notify, Semaphore};
|
||||
use tower::ServiceExt;
|
||||
use workdir::{LocalWorkdirSession, Workdir, WorkdirError, WorkdirSession};
|
||||
|
||||
#[derive(Debug)]
|
||||
struct CloseProvider {
|
||||
local: LocalWorkdirSession,
|
||||
calls: AtomicUsize,
|
||||
cancelled: AtomicUsize,
|
||||
pause_first: bool,
|
||||
fail_first: bool,
|
||||
entered: Notify,
|
||||
release: Semaphore,
|
||||
}
|
||||
|
||||
struct AttemptGuard<'a> {
|
||||
cancelled: &'a AtomicUsize,
|
||||
finished: bool,
|
||||
}
|
||||
|
||||
impl Drop for AttemptGuard<'_> {
|
||||
fn drop(&mut self) {
|
||||
if !self.finished {
|
||||
self.cancelled.fetch_add(1, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Only close scheduling/error injection is substituted. All session operations and
|
||||
// actual teardown remain Local; in particular its early closed.swap is exercised.
|
||||
#[async_trait::async_trait]
|
||||
impl WorkdirSession for CloseProvider {
|
||||
fn workdir(&self) -> &Workdir {
|
||||
self.local.workdir()
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> workdir::WorkdirSessionCapabilities {
|
||||
self.local.capabilities()
|
||||
}
|
||||
|
||||
async fn stat(
|
||||
&self,
|
||||
request: workdir::StatRequest,
|
||||
) -> Result<workdir::StatResult, WorkdirError> {
|
||||
self.local.stat(request).await
|
||||
}
|
||||
async fn read(
|
||||
&self,
|
||||
request: workdir::ReadRequest,
|
||||
) -> Result<workdir::ReadResult, WorkdirError> {
|
||||
self.local.read(request).await
|
||||
}
|
||||
async fn write(
|
||||
&self,
|
||||
request: workdir::WriteRequest,
|
||||
) -> Result<workdir::WriteResult, WorkdirError> {
|
||||
self.local.write(request).await
|
||||
}
|
||||
async fn edit(
|
||||
&self,
|
||||
request: workdir::EditRequest,
|
||||
) -> Result<workdir::EditResult, WorkdirError> {
|
||||
self.local.edit(request).await
|
||||
}
|
||||
async fn list(
|
||||
&self,
|
||||
request: workdir::ListRequest,
|
||||
) -> Result<workdir::ListResult, WorkdirError> {
|
||||
self.local.list(request).await
|
||||
}
|
||||
async fn glob(
|
||||
&self,
|
||||
request: workdir::GlobRequest,
|
||||
) -> Result<workdir::GlobResult, WorkdirError> {
|
||||
self.local.glob(request).await
|
||||
}
|
||||
async fn grep(
|
||||
&self,
|
||||
request: workdir::GrepRequest,
|
||||
) -> Result<workdir::GrepResult, WorkdirError> {
|
||||
self.local.grep(request).await
|
||||
}
|
||||
async fn start_command(
|
||||
&self,
|
||||
request: workdir::CommandRequest,
|
||||
) -> Result<workdir::CommandHandle, WorkdirError> {
|
||||
self.local.start_command(request).await
|
||||
}
|
||||
async fn command_status(
|
||||
&self,
|
||||
request: workdir::CommandHandle,
|
||||
) -> Result<workdir::CommandStatus, WorkdirError> {
|
||||
self.local.command_status(request).await
|
||||
}
|
||||
async fn command_output(
|
||||
&self,
|
||||
request: workdir::CommandOutputRequest,
|
||||
) -> Result<workdir::CommandOutput, WorkdirError> {
|
||||
self.local.command_output(request).await
|
||||
}
|
||||
async fn cancel_command(&self, request: workdir::CommandHandle) -> Result<(), WorkdirError> {
|
||||
self.local.cancel_command(request).await
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), WorkdirError> {
|
||||
let attempt = self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
if attempt == 0 && self.fail_first {
|
||||
return Err(WorkdirError::Unavailable("injected close failure".into()));
|
||||
}
|
||||
let mut guard = AttemptGuard {
|
||||
cancelled: &self.cancelled,
|
||||
finished: false,
|
||||
};
|
||||
if attempt == 0 && self.pause_first {
|
||||
self.local
|
||||
.start_command(workdir::CommandRequest {
|
||||
command: "printf close-probe".into(),
|
||||
timeout_secs: 5,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
tool_call_id: None,
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
let mut closing = Box::pin(self.local.close());
|
||||
if attempt == 0 && self.pause_first {
|
||||
// On the current-thread executor the command task cannot run before
|
||||
// this poll. Local close sets closed, drains/cancels the command,
|
||||
// then is genuinely Pending on command.task.await (reviewer's probe).
|
||||
assert!(
|
||||
std::future::poll_fn(|cx| Poll::Ready(closing.as_mut().poll(cx)))
|
||||
.await
|
||||
.is_pending()
|
||||
);
|
||||
self.entered.notify_one();
|
||||
// Hold the already-polled provider future, not a fake release flag.
|
||||
// Cancelling this future would strand Local's unconfirmed teardown.
|
||||
let _permit = self.release.acquire().await.unwrap();
|
||||
}
|
||||
let result = closing.await;
|
||||
guard.finished = true;
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
struct Backend {
|
||||
session: Arc<CloseProvider>,
|
||||
}
|
||||
|
||||
impl WorkerExecutionBackend for Backend {
|
||||
fn backend_id(&self) -> &str {
|
||||
"close-confirmation-fixture"
|
||||
}
|
||||
|
||||
fn spawn_worker(&self, _: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
|
||||
panic!("not used")
|
||||
}
|
||||
|
||||
fn dispatch_input(&self, _: &WorkerRef, _: WorkerInput) -> WorkerExecutionResult {
|
||||
panic!("not used")
|
||||
}
|
||||
|
||||
fn working_directory(
|
||||
&self,
|
||||
_: &str,
|
||||
) -> Result<WorkingDirectoryStatus, WorkingDirectoryDiagnostic> {
|
||||
Ok(serde_json::from_value(serde_json::json!({
|
||||
"summary": {
|
||||
"working_directory_id": "w", "repository_id": "r",
|
||||
"materializer_kind": "runtime_git_clone", "status": "active"
|
||||
}
|
||||
}))
|
||||
.unwrap())
|
||||
}
|
||||
|
||||
fn open_workdir_session(
|
||||
&self,
|
||||
_: &str,
|
||||
) -> Result<WorkdirSessionHandle, WorkingDirectoryDiagnostic> {
|
||||
Ok(self.session.clone())
|
||||
}
|
||||
|
||||
fn cleanup_working_directory(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<WorkingDirectoryStatus, WorkingDirectoryDiagnostic> {
|
||||
self.working_directory(id)
|
||||
}
|
||||
}
|
||||
|
||||
struct Fixture {
|
||||
_temp: tempfile::TempDir,
|
||||
provider: Arc<CloseProvider>,
|
||||
runtime: Runtime,
|
||||
router: Router,
|
||||
path: String,
|
||||
}
|
||||
|
||||
impl Fixture {
|
||||
async fn new(pause_first: bool, fail_first: bool) -> Self {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
std::fs::write(temp.path().join("hello.txt"), "hello").unwrap();
|
||||
let provider = Arc::new(CloseProvider {
|
||||
local: LocalWorkdirSession::materialized_bound(
|
||||
Workdir::new("w"),
|
||||
temp.path().to_owned(),
|
||||
temp.path().to_owned(),
|
||||
manifest::SharedScope::new(manifest::Scope::writable(temp.path()).unwrap()),
|
||||
workdir::WorkdirSessionCapabilities::ALL,
|
||||
),
|
||||
calls: AtomicUsize::new(0),
|
||||
cancelled: AtomicUsize::new(0),
|
||||
pause_first,
|
||||
fail_first,
|
||||
entered: Notify::new(),
|
||||
release: Semaphore::new(0),
|
||||
});
|
||||
let runtime = Runtime::with_execution_backend(
|
||||
RuntimeOptions::default(),
|
||||
Arc::new(Backend {
|
||||
session: provider.clone(),
|
||||
}),
|
||||
)
|
||||
.unwrap();
|
||||
let router = runtime_http_router(runtime.clone(), "probe-token".into());
|
||||
let opened = router
|
||||
.clone()
|
||||
.oneshot(request("POST", "/v1/working-directories/w/sessions", "{}"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(opened.status(), StatusCode::OK);
|
||||
let opened: OpenWorkdirSessionResponse =
|
||||
serde_json::from_slice(&to_bytes(opened.into_body(), 65536).await.unwrap()).unwrap();
|
||||
let path = format!("/v1/workdir-sessions/{}", opened.session_id.as_str());
|
||||
Self {
|
||||
_temp: temp,
|
||||
provider,
|
||||
runtime,
|
||||
router,
|
||||
path,
|
||||
}
|
||||
}
|
||||
|
||||
fn close(&self) -> impl Future<Output = Response> + use<> {
|
||||
let router = self.router.clone();
|
||||
let request = request("DELETE", &self.path, "");
|
||||
async move { router.oneshot(request).await.unwrap() }
|
||||
}
|
||||
|
||||
fn assert_busy(&self) {
|
||||
let error = self
|
||||
.runtime
|
||||
.cleanup_working_directory("w")
|
||||
.expect_err("unconfirmed close must retain the real Runtime use guard");
|
||||
assert!(
|
||||
matches!(error, RuntimeError::WorkingDirectory(ref d)
|
||||
if d.code == "working_directory_cleanup_resource_busy"),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
async fn assert_operations_unavailable(&self) {
|
||||
let response = self
|
||||
.router
|
||||
.clone()
|
||||
.oneshot(request(
|
||||
"POST",
|
||||
&format!("{}/operations", self.path),
|
||||
r#"{"operation":{"operation":"stat","request":{"path":"hello.txt"}}}"#,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::NOT_FOUND);
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Fixture {
|
||||
fn drop(&mut self) {
|
||||
// Let a detached close finish even if a test assertion unwinds.
|
||||
self.provider.release.add_permits(1);
|
||||
}
|
||||
}
|
||||
|
||||
fn request(method: &str, path: &str, body: &str) -> Request<Body> {
|
||||
Request::builder()
|
||||
.method(method)
|
||||
.uri(path)
|
||||
.header("authorization", "Bearer probe-token")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(body.to_owned()))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn bounded<T>(future: impl Future<Output = T>) -> T {
|
||||
tokio::time::timeout(std::time::Duration::from_secs(10), future)
|
||||
.await
|
||||
.expect("close probe must not hang")
|
||||
}
|
||||
|
||||
async fn assert_pending<T>(mut future: std::pin::Pin<&mut impl Future<Output = T>>) {
|
||||
assert!(
|
||||
std::future::poll_fn(|cx| Poll::Ready(future.as_mut().poll(cx)))
|
||||
.await
|
||||
.is_pending(),
|
||||
"close acknowledged before teardown confirmation"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn duplicate_delete_waits_for_actual_local_teardown_and_runtime_guard_release() {
|
||||
let fixture = Fixture::new(true, false).await;
|
||||
fixture.assert_busy();
|
||||
let mut first = Box::pin(fixture.close());
|
||||
assert_pending(first.as_mut()).await;
|
||||
bounded(fixture.provider.entered.notified()).await;
|
||||
let mut duplicate = Box::pin(fixture.close());
|
||||
assert_pending(duplicate.as_mut()).await;
|
||||
fixture.assert_busy();
|
||||
fixture.assert_operations_unavailable().await;
|
||||
assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 1);
|
||||
|
||||
fixture.provider.release.add_permits(1);
|
||||
assert_eq!(bounded(first).await.status(), StatusCode::NO_CONTENT);
|
||||
assert_eq!(bounded(duplicate).await.status(), StatusCode::NO_CONTENT);
|
||||
fixture
|
||||
.runtime
|
||||
.cleanup_working_directory("w")
|
||||
.expect("confirmed close releases use guard");
|
||||
assert_eq!(
|
||||
bounded(fixture.close()).await.status(),
|
||||
StatusCode::NO_CONTENT
|
||||
);
|
||||
assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn cancelled_http_waiter_does_not_cancel_provider_or_let_retries_ack_early() {
|
||||
let fixture = Fixture::new(true, false).await;
|
||||
let mut first = Box::pin(fixture.close());
|
||||
assert_pending(first.as_mut()).await;
|
||||
bounded(fixture.provider.entered.notified()).await;
|
||||
drop(first);
|
||||
let mut retry = Box::pin(fixture.close());
|
||||
let mut duplicate = Box::pin(fixture.close());
|
||||
assert_pending(retry.as_mut()).await;
|
||||
assert_pending(duplicate.as_mut()).await;
|
||||
fixture.assert_busy();
|
||||
assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(fixture.provider.cancelled.load(Ordering::SeqCst), 0);
|
||||
|
||||
fixture.provider.release.add_permits(1);
|
||||
assert_eq!(bounded(retry).await.status(), StatusCode::NO_CONTENT);
|
||||
assert_eq!(bounded(duplicate).await.status(), StatusCode::NO_CONTENT);
|
||||
assert_eq!(fixture.provider.cancelled.load(Ordering::SeqCst), 0);
|
||||
fixture
|
||||
.runtime
|
||||
.cleanup_working_directory("w")
|
||||
.expect("task completes real provider close");
|
||||
assert_eq!(
|
||||
bounded(fixture.close()).await.status(),
|
||||
StatusCode::NO_CONTENT
|
||||
);
|
||||
assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn provider_close_error_retains_session_and_runtime_guard_for_retry() {
|
||||
let fixture = Fixture::new(false, true).await;
|
||||
let mut first = Box::pin(fixture.close());
|
||||
let mut duplicate = Box::pin(fixture.close());
|
||||
assert_pending(first.as_mut()).await;
|
||||
assert_pending(duplicate.as_mut()).await;
|
||||
let response = bounded(first).await;
|
||||
assert_eq!(
|
||||
bounded(duplicate).await.status(),
|
||||
StatusCode::SERVICE_UNAVAILABLE
|
||||
);
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let error: WorkdirTransportError =
|
||||
serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap();
|
||||
assert_eq!(error.code, WorkdirTransportErrorCode::Unavailable);
|
||||
fixture.assert_busy();
|
||||
assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 1);
|
||||
|
||||
assert_eq!(
|
||||
bounded(fixture.close()).await.status(),
|
||||
StatusCode::NO_CONTENT
|
||||
);
|
||||
assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 2);
|
||||
fixture
|
||||
.runtime
|
||||
.cleanup_working_directory("w")
|
||||
.expect("retry releases the retained Runtime guard");
|
||||
assert_eq!(
|
||||
bounded(fixture.close()).await.status(),
|
||||
StatusCode::NO_CONTENT
|
||||
);
|
||||
assert_eq!(fixture.provider.calls.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
@@ -704,28 +704,52 @@ impl WorkdirSessionRegistry {
|
||||
}
|
||||
}
|
||||
|
||||
// Extraction is not release confirmation. Retain every unconfirmed entry across
|
||||
// the await, including entries not yet attempted when the close future is dropped.
|
||||
// Callers keep their Worker session fence until this restoration has completed.
|
||||
struct UnconfirmedWorkdirSessionCloses<'a> {
|
||||
registry: &'a Arc<Mutex<WorkdirSessionRegistry>>,
|
||||
registered: Vec<RegisteredWorkdirSession>,
|
||||
}
|
||||
|
||||
impl Drop for UnconfirmedWorkdirSessionCloses<'_> {
|
||||
fn drop(&mut self) {
|
||||
if self.registered.is_empty() {
|
||||
return;
|
||||
}
|
||||
let mut sessions = self
|
||||
.registry
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
for item in self.registered.drain(..) {
|
||||
sessions.restore(item);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn close_registered_workdir_sessions(
|
||||
registry: &Arc<Mutex<WorkdirSessionRegistry>>,
|
||||
registered: Vec<RegisteredWorkdirSession>,
|
||||
) -> std::result::Result<(), String> {
|
||||
let mut failed = Vec::new();
|
||||
let mut unconfirmed = UnconfirmedWorkdirSessionCloses {
|
||||
registry,
|
||||
registered,
|
||||
};
|
||||
let mut errors = Vec::new();
|
||||
for item in registered {
|
||||
let session = match &item {
|
||||
let mut index = 0;
|
||||
while let Some(item) = unconfirmed.registered.get(index) {
|
||||
let session = match item {
|
||||
RegisteredWorkdirSession::Attachment { session, .. } => session,
|
||||
RegisteredWorkdirSession::Command { session, .. } => &session.source,
|
||||
};
|
||||
if let Err(error) = session.close().await {
|
||||
errors.push(error.to_string());
|
||||
failed.push(item);
|
||||
}
|
||||
}
|
||||
if !failed.is_empty() {
|
||||
let mut sessions = registry
|
||||
.lock()
|
||||
.map_err(|_| "Workdir session registry was poisoned".to_string())?;
|
||||
for item in failed {
|
||||
sessions.restore(item);
|
||||
match session.close().await {
|
||||
Ok(()) => {
|
||||
unconfirmed.registered.remove(index);
|
||||
}
|
||||
Err(error) => {
|
||||
errors.push(error.to_string());
|
||||
index += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
if errors.is_empty() {
|
||||
@@ -21407,17 +21431,19 @@ async fn close_current_worker_attachment_session_locked(
|
||||
.expect("Workdir session registry lock poisoned")
|
||||
.remove_attachment(worker);
|
||||
if let Some(session) = session {
|
||||
if let Err(error) = session.close().await {
|
||||
api.workdir_sessions
|
||||
.lock()
|
||||
.expect("Workdir session registry lock poisoned")
|
||||
.insert_attachment(worker.clone(), session);
|
||||
return Err(Error::RuntimeOperationFailed {
|
||||
runtime_id: worker.runtime_id.clone(),
|
||||
code: "workdir_session_close_failed".to_string(),
|
||||
message: error.to_string(),
|
||||
});
|
||||
}
|
||||
close_registered_workdir_sessions(
|
||||
&api.workdir_sessions,
|
||||
vec![RegisteredWorkdirSession::Attachment {
|
||||
worker: worker.clone(),
|
||||
session,
|
||||
}],
|
||||
)
|
||||
.await
|
||||
.map_err(|message| Error::RuntimeOperationFailed {
|
||||
runtime_id: worker.runtime_id.clone(),
|
||||
code: "workdir_session_close_failed".to_string(),
|
||||
message,
|
||||
})?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -39310,6 +39336,7 @@ mod tests {
|
||||
mod ticket_worker_launch_tests;
|
||||
mod value_profiles_tests;
|
||||
mod workdir_delegation_tests;
|
||||
mod workdir_session_close_tests;
|
||||
mod worker_drive_tests;
|
||||
mod worker_operations_tests;
|
||||
mod worker_protocol_tests;
|
||||
@@ -51253,6 +51280,10 @@ mod tests {
|
||||
let worker_root = root
|
||||
.join(".test-embedded-runtime-store/workers")
|
||||
.join(&worker.worker_id);
|
||||
seed_worker_session_at_root(&worker_root);
|
||||
}
|
||||
|
||||
fn seed_worker_session_at_root(worker_root: &Path) {
|
||||
use session_store::Store;
|
||||
let store = session_store::WorkerSessionStore::new(worker_root.join("session")).unwrap();
|
||||
let session = session_store::new_session_id();
|
||||
|
||||
@@ -13,6 +13,7 @@ use worker_runtime::execution::{
|
||||
#[derive(Default)]
|
||||
struct LaunchExecutionBackend {
|
||||
execution: DeterministicExecutionBackend,
|
||||
provider_close_gate: Mutex<Option<Arc<super::workdir_session_close_tests::ProviderCloseGate>>>,
|
||||
bundles: Mutex<std::collections::HashMap<String, worker_runtime::config_bundle::ConfigBundle>>,
|
||||
creates: Mutex<Vec<worker_runtime::catalog::CreateWorkerRequest>>,
|
||||
submits: Mutex<
|
||||
@@ -122,6 +123,35 @@ impl WorkerExecutionBackend for LaunchExecutionBackend {
|
||||
self.execution.working_directory(id)
|
||||
}
|
||||
|
||||
fn open_workdir_session(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> std::result::Result<
|
||||
WorkdirSessionHandle,
|
||||
worker_runtime::working_directory::WorkingDirectoryDiagnostic,
|
||||
> {
|
||||
let binding = self
|
||||
.execution
|
||||
.materializer
|
||||
.bind_working_directory(id, None)?;
|
||||
let scope = manifest::Scope::writable(binding.root()).unwrap();
|
||||
let session: WorkdirSessionHandle = Arc::new(
|
||||
workdir::LocalWorkdirSession::materialized_bound_with_environment(
|
||||
workdir::Workdir::new(id),
|
||||
binding.root().to_owned(),
|
||||
binding.root().to_owned(),
|
||||
manifest::SharedScope::new(scope),
|
||||
workdir::WorkdirSessionCapabilities::ALL,
|
||||
binding.command_environment(),
|
||||
binding.session_resources()?,
|
||||
),
|
||||
);
|
||||
Ok(match self.provider_close_gate.lock().unwrap().as_ref() {
|
||||
Some(gate) => gate.wrap(session),
|
||||
None => session,
|
||||
})
|
||||
}
|
||||
|
||||
fn create_working_directory(
|
||||
&self,
|
||||
request: &worker_runtime::catalog::WorkingDirectoryRequest,
|
||||
@@ -180,11 +210,15 @@ pub(super) struct LaunchFixture {
|
||||
backend: Arc<LaunchExecutionBackend>,
|
||||
runtime_id: String,
|
||||
runtime_server: Option<tokio::task::JoinHandle<()>>,
|
||||
runtime_router: Option<Router>,
|
||||
runtime_identity: Option<RuntimeIdentityMaterial>,
|
||||
}
|
||||
|
||||
impl Drop for LaunchFixture {
|
||||
fn drop(&mut self) {
|
||||
if let Some(gate) = self.backend.provider_close_gate.lock().unwrap().as_ref() {
|
||||
gate.release();
|
||||
}
|
||||
if let Some(server) = &self.runtime_server {
|
||||
server.abort();
|
||||
}
|
||||
@@ -192,13 +226,38 @@ impl Drop for LaunchFixture {
|
||||
}
|
||||
|
||||
impl LaunchFixture {
|
||||
pub(super) fn hold_provider_close(
|
||||
&self,
|
||||
gate: Arc<super::workdir_session_close_tests::ProviderCloseGate>,
|
||||
) {
|
||||
*self.backend.provider_close_gate.lock().unwrap() = Some(gate);
|
||||
}
|
||||
|
||||
pub(super) fn runtime_router(&self) -> Router {
|
||||
self.runtime_router.clone().unwrap()
|
||||
}
|
||||
|
||||
pub(super) fn publish_profile(&self, profile: &str) {
|
||||
self.backend.publish_profile(&self.api, profile);
|
||||
}
|
||||
|
||||
pub(super) fn seed_remote_session_for_cleanup(&self, worker: &RuntimeWorkerRef) {
|
||||
seed_worker_session_at_root(
|
||||
&self
|
||||
._temp
|
||||
.path()
|
||||
.join(".test-embedded-runtime-store/source-runtime/workers")
|
||||
.join(&worker.worker_id),
|
||||
);
|
||||
}
|
||||
|
||||
// Reuse the authenticated Runtime HTTP harness for Drive/source lifetime
|
||||
// assertions without duplicating its trust/bootstrap/listener lifecycle.
|
||||
pub(super) async fn with_source_workdir() -> (
|
||||
Self,
|
||||
worker_runtime::working_directory::WorkingDirectoryBinding,
|
||||
) {
|
||||
let fixture = Self::new(true).await;
|
||||
let fixture = Self::new_with_runtime_store(true, true).await;
|
||||
init_clean_git_workspace(fixture._temp.path());
|
||||
let (_, Json(created)) = create_workspace_working_directory(
|
||||
&fixture.api,
|
||||
@@ -224,6 +283,10 @@ impl LaunchFixture {
|
||||
}
|
||||
|
||||
async fn new(remote: bool) -> Self {
|
||||
Self::new_with_runtime_store(remote, false).await
|
||||
}
|
||||
|
||||
async fn new_with_runtime_store(remote: bool, persistent: bool) -> Self {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let config = test_server_config(temp.path());
|
||||
let store = SqliteWorkspaceStore::open_for_test(config.database_path.clone()).unwrap();
|
||||
@@ -247,6 +310,7 @@ impl LaunchFixture {
|
||||
backend,
|
||||
runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(),
|
||||
runtime_server: None,
|
||||
runtime_router: None,
|
||||
runtime_identity: None,
|
||||
};
|
||||
if remote {
|
||||
@@ -282,11 +346,26 @@ impl LaunchFixture {
|
||||
Arc::new(InMemoryWorkspaceClaimReplayProtection::default()),
|
||||
)
|
||||
.unwrap();
|
||||
let runtime = worker_runtime::Runtime::with_execution_backend(
|
||||
worker_runtime::RuntimeOptions::default(),
|
||||
fixture.backend.clone(),
|
||||
)
|
||||
let runtime = if persistent {
|
||||
worker_runtime::Runtime::with_fs_store_and_execution_backend(
|
||||
worker_runtime::fs_store::FsRuntimeStoreOptions {
|
||||
root: fixture
|
||||
._temp
|
||||
.path()
|
||||
.join(".test-embedded-runtime-store/source-runtime"),
|
||||
runtime_id: fixture.runtime_id.clone(),
|
||||
display_name: None,
|
||||
},
|
||||
fixture.backend.clone(),
|
||||
)
|
||||
} else {
|
||||
worker_runtime::Runtime::with_execution_backend(
|
||||
worker_runtime::RuntimeOptions::default(),
|
||||
fixture.backend.clone(),
|
||||
)
|
||||
}
|
||||
.unwrap();
|
||||
runtime.bind_runtime_identity(&fixture.runtime_id).unwrap();
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let endpoint = format!("http://{}", listener.local_addr().unwrap());
|
||||
let binding = WorkspaceRuntimeBinding {
|
||||
@@ -318,18 +397,17 @@ impl LaunchFixture {
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let runtime_id = fixture.runtime_id.clone();
|
||||
let router = worker_runtime::http_server::runtime_http_router_with_workspace_auth(
|
||||
runtime,
|
||||
None,
|
||||
worker_runtime::http_server::WorkspaceRuntimeHttpAuth {
|
||||
verifier,
|
||||
runtime_id,
|
||||
},
|
||||
);
|
||||
fixture.runtime_router = Some(router.clone());
|
||||
fixture.runtime_server = Some(tokio::spawn(async move {
|
||||
worker_runtime::http_server::serve_runtime_http_with_workspace_auth(
|
||||
runtime,
|
||||
listener,
|
||||
None,
|
||||
worker_runtime::http_server::WorkspaceRuntimeHttpAuth {
|
||||
verifier,
|
||||
runtime_id,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
axum::serve(listener, router).await.unwrap();
|
||||
}));
|
||||
let mut remote_config = RemoteRuntimeConfig {
|
||||
runtime_id: fixture.runtime_id.clone(),
|
||||
|
||||
@@ -0,0 +1,772 @@
|
||||
//! Backend registry/current-attachment/WorkerRemove integration with authenticated
|
||||
//! Runtime HTTP sessions and real Git materializer resource leases. Only Worker
|
||||
//! execution and the failing/held close transport are replaced. Ordinary library
|
||||
//! tests use LaunchFixture's isolated SQLite, materializations and loopback Runtime.
|
||||
use super::ticket_worker_launch_tests::LaunchFixture;
|
||||
use super::*;
|
||||
use std::{future::Future, task::Poll};
|
||||
use tokio::sync::{Notify, Semaphore};
|
||||
use tower::ServiceExt;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(super) struct ProviderCloseGate {
|
||||
entered: Notify,
|
||||
release: Semaphore,
|
||||
}
|
||||
|
||||
impl ProviderCloseGate {
|
||||
pub(super) fn release(&self) {
|
||||
self.release.add_permits(1);
|
||||
}
|
||||
|
||||
pub(super) fn wrap(self: &Arc<Self>, source: WorkdirSessionHandle) -> WorkdirSessionHandle {
|
||||
Arc::new(GatedProviderClose {
|
||||
source,
|
||||
gate: self.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct GatedProviderClose {
|
||||
source: WorkdirSessionHandle,
|
||||
gate: Arc<ProviderCloseGate>,
|
||||
}
|
||||
|
||||
// Forward the real provider operations; only suspend an already-pending close.
|
||||
macro_rules! gated_provider {
|
||||
($($name:ident($request:ident: $input:ty) -> $output:ty;)*) => {
|
||||
#[async_trait::async_trait]
|
||||
impl WorkdirSession for GatedProviderClose {
|
||||
fn workdir(&self) -> &workdir::Workdir { self.source.workdir() }
|
||||
fn capabilities(&self) -> workdir::WorkdirSessionCapabilities { self.source.capabilities() }
|
||||
fn environment(&self) -> workdir::WorkdirSessionEnvironment { self.source.environment() }
|
||||
$(async fn $name(&self, $request: $input) -> std::result::Result<$output, workdir::WorkdirError> {
|
||||
self.source.$name($request).await
|
||||
})*
|
||||
async fn close(&self) -> std::result::Result<(), workdir::WorkdirError> {
|
||||
let mut closing = Box::pin(self.source.close());
|
||||
match std::future::poll_fn(|cx| Poll::Ready(closing.as_mut().poll(cx))).await {
|
||||
Poll::Ready(result) => result,
|
||||
Poll::Pending => {
|
||||
self.gate.entered.notify_one();
|
||||
self.gate.release.acquire().await.unwrap().forget();
|
||||
closing.await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
gated_provider! {
|
||||
stat(request: workdir::StatRequest) -> workdir::StatResult;
|
||||
read(request: workdir::ReadRequest) -> workdir::ReadResult;
|
||||
write(request: workdir::WriteRequest) -> workdir::WriteResult;
|
||||
edit(request: workdir::EditRequest) -> workdir::EditResult;
|
||||
list(request: workdir::ListRequest) -> workdir::ListResult;
|
||||
glob(request: workdir::GlobRequest) -> workdir::GlobResult;
|
||||
grep(request: workdir::GrepRequest) -> workdir::GrepResult;
|
||||
start_command(request: workdir::CommandRequest) -> workdir::CommandHandle;
|
||||
command_status(handle: workdir::CommandHandle) -> workdir::CommandStatus;
|
||||
command_output(request: workdir::CommandOutputRequest) -> workdir::CommandOutput;
|
||||
cancel_command(handle: workdir::CommandHandle) -> ();
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum CloseFault {
|
||||
Reject,
|
||||
HoldThenReject,
|
||||
RejectAfterRuntimeClose,
|
||||
}
|
||||
|
||||
struct CloseTransport {
|
||||
upstream: String,
|
||||
client: reqwest::Client,
|
||||
fault: Mutex<Option<CloseFault>>,
|
||||
attempts: Mutex<Vec<String>>,
|
||||
runtime_closes: Mutex<Vec<(String, StatusCode)>>,
|
||||
direct_runtime: Mutex<Option<Router>>,
|
||||
pending_dispatches: Mutex<Vec<bool>>,
|
||||
dispatched: Notify,
|
||||
entered: Notify,
|
||||
release: Semaphore,
|
||||
}
|
||||
|
||||
impl CloseTransport {
|
||||
async fn forward(State(state): State<Arc<Self>>, request: Request<Body>) -> Response {
|
||||
let is_close = request.method() == Method::DELETE
|
||||
&& request.uri().path().starts_with("/v1/workdir-sessions/");
|
||||
let path = request.uri().path().to_owned();
|
||||
let fault = if is_close {
|
||||
state.attempts.lock().unwrap().push(path.clone());
|
||||
state.fault.lock().unwrap().take()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if matches!(fault, Some(CloseFault::Reject | CloseFault::HoldThenReject)) {
|
||||
if matches!(fault, Some(CloseFault::HoldThenReject)) {
|
||||
state.entered.notify_one();
|
||||
state.release.acquire().await.unwrap().forget();
|
||||
}
|
||||
// The actual Runtime has not seen this request or released its lease.
|
||||
return StatusCode::SERVICE_UNAVAILABLE.into_response();
|
||||
}
|
||||
let direct = state.direct_runtime.lock().unwrap().clone();
|
||||
if let Some(router) = direct.filter(|_| is_close) {
|
||||
// Poll the production authenticated handler without another socket
|
||||
// scheduling boundary, so a duplicate's Pending/204 is observable.
|
||||
let mut response = Box::pin(router.oneshot(request));
|
||||
let first_poll =
|
||||
std::future::poll_fn(|cx| Poll::Ready(response.as_mut().poll(cx))).await;
|
||||
state
|
||||
.pending_dispatches
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(first_poll.is_pending());
|
||||
state.dispatched.notify_one();
|
||||
let response = match first_poll {
|
||||
Poll::Ready(result) => result.unwrap(),
|
||||
Poll::Pending => response.await.unwrap(),
|
||||
};
|
||||
state
|
||||
.runtime_closes
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push((path, response.status()));
|
||||
return response;
|
||||
}
|
||||
let (parts, body) = request.into_parts();
|
||||
let body = to_bytes(body, 1024 * 1024).await.unwrap();
|
||||
let response = state
|
||||
.client
|
||||
.request(parts.method, format!("{}{}", state.upstream, parts.uri))
|
||||
.headers(parts.headers)
|
||||
.body(body)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
let status = response.status();
|
||||
let headers = response.headers().clone();
|
||||
let bytes = response.bytes().await.unwrap();
|
||||
if is_close {
|
||||
state.runtime_closes.lock().unwrap().push((path, status));
|
||||
}
|
||||
if matches!(fault, Some(CloseFault::RejectAfterRuntimeClose)) {
|
||||
assert_eq!(status, StatusCode::NO_CONTENT);
|
||||
// Lose only the successful acknowledgement, after Runtime has closed
|
||||
// the actual session and released its active-use/materializer leases.
|
||||
return StatusCode::BAD_GATEWAY.into_response();
|
||||
}
|
||||
let mut response = Response::new(Body::from(bytes));
|
||||
*response.status_mut() = status;
|
||||
*response.headers_mut() = headers;
|
||||
response
|
||||
}
|
||||
|
||||
fn fail_next(&self, fault: CloseFault) {
|
||||
assert!(self.fault.lock().unwrap().replace(fault).is_none());
|
||||
}
|
||||
|
||||
fn attempts(&self) -> Vec<String> {
|
||||
self.attempts.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
struct Fixture {
|
||||
source: LaunchFixture,
|
||||
transport: Arc<CloseTransport>,
|
||||
proxy: tokio::task::JoinHandle<()>,
|
||||
runtime_id: String,
|
||||
workdir_id: String,
|
||||
materialization_root: PathBuf,
|
||||
}
|
||||
|
||||
impl Drop for Fixture {
|
||||
fn drop(&mut self) {
|
||||
self.proxy.abort();
|
||||
}
|
||||
}
|
||||
|
||||
impl Fixture {
|
||||
async fn new() -> Self {
|
||||
let (source, binding) = LaunchFixture::with_source_workdir().await;
|
||||
let workdir_id = binding.working_directory.id.clone();
|
||||
let materialization_root = binding.working_directory_root().to_owned();
|
||||
// Do not accidentally retain a second materializer lease in the fixture.
|
||||
drop(binding);
|
||||
let record = source
|
||||
.api
|
||||
.store
|
||||
.get_workdir_registry(TEST_WORKSPACE_ID, &workdir_id)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let runtime_id = record.source.runtime_id().unwrap().to_owned();
|
||||
let binding = source
|
||||
.api
|
||||
.store
|
||||
.get_workspace_runtime_binding(TEST_WORKSPACE_ID, &runtime_id)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let transport = Arc::new(CloseTransport {
|
||||
upstream: binding.base_url.clone(),
|
||||
client: reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.timeout(std::time::Duration::from_secs(10))
|
||||
.build()
|
||||
.unwrap(),
|
||||
fault: Mutex::new(None),
|
||||
attempts: Mutex::new(Vec::new()),
|
||||
runtime_closes: Mutex::new(Vec::new()),
|
||||
direct_runtime: Mutex::new(None),
|
||||
pending_dispatches: Mutex::new(Vec::new()),
|
||||
dispatched: Notify::new(),
|
||||
entered: Notify::new(),
|
||||
release: Semaphore::new(0),
|
||||
});
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let endpoint = format!("http://{}", listener.local_addr().unwrap());
|
||||
let router = Router::new()
|
||||
.fallback(axum::routing::any(CloseTransport::forward))
|
||||
.with_state(transport.clone());
|
||||
let proxy = tokio::spawn(async move {
|
||||
axum::serve(listener, router).await.unwrap();
|
||||
});
|
||||
let backend_url = source.api.config.backend_base_url.clone().unwrap();
|
||||
let remote = RemoteWorkerRuntime::new(
|
||||
RemoteRuntimeConfig {
|
||||
runtime_id: runtime_id.clone(),
|
||||
workspace_id: Some(TEST_WORKSPACE_ID.into()),
|
||||
display_name: binding.display_name.clone(),
|
||||
base_url: endpoint,
|
||||
bearer_token: None,
|
||||
workspace_authorization: Some(crate::hosts::WorkspaceRuntimeAuthorization::new(
|
||||
source.api.store.clone(),
|
||||
source.api.signing_identities.clone(),
|
||||
backend_url.clone(),
|
||||
binding,
|
||||
)),
|
||||
strict_public_egress: false,
|
||||
cached_worker_creation_available: true,
|
||||
cached_os: "linux".into(),
|
||||
cached_arch: "x86_64".into(),
|
||||
cached_status: "active".into(),
|
||||
timeout: std::time::Duration::from_secs(10),
|
||||
},
|
||||
TEST_WORKSPACE_ID.into(),
|
||||
backend_url,
|
||||
)
|
||||
.unwrap();
|
||||
source.api.runtime.register_or_replace(remote);
|
||||
Self {
|
||||
source,
|
||||
transport,
|
||||
proxy,
|
||||
runtime_id,
|
||||
workdir_id,
|
||||
materialization_root,
|
||||
}
|
||||
}
|
||||
|
||||
fn api(&self) -> &WorkspaceApi {
|
||||
&self.source.api
|
||||
}
|
||||
|
||||
fn registry_owner(&self) -> RuntimeWorkerRef {
|
||||
RuntimeWorkerRef::new(&self.runtime_id, "registry-close-owner")
|
||||
}
|
||||
|
||||
async fn open(&self) -> WorkdirSessionHandle {
|
||||
self.api()
|
||||
.runtime
|
||||
.open_workdir_session(&self.runtime_id, &self.workdir_id, None)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn assert_restored(&self, worker: &RuntimeWorkerRef, expected: &WorkdirSessionHandle) {
|
||||
let registry = self.api().workdir_sessions.lock().unwrap();
|
||||
assert!(Arc::ptr_eq(
|
||||
registry
|
||||
.attachments
|
||||
.get(worker)
|
||||
.expect("unconfirmed session retained"),
|
||||
expected
|
||||
));
|
||||
}
|
||||
|
||||
async fn assert_live(&self, session: &WorkdirSessionHandle) {
|
||||
session
|
||||
.stat(workdir::StatRequest {
|
||||
path: workdir::WorkdirPath::new("README.md").unwrap(),
|
||||
})
|
||||
.await
|
||||
.expect("same Runtime session still serves filesystem operations");
|
||||
}
|
||||
|
||||
async fn remove_workdir(&self) -> WorkingDirectoryRemovalResponse {
|
||||
let api = self.api().clone();
|
||||
let workdir_id = self.workdir_id.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
execute_workdir_removal(
|
||||
&api,
|
||||
&workdir_id,
|
||||
"account:close-test",
|
||||
"remove normally after session close",
|
||||
)
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn assert_busy(&self) {
|
||||
let result = self.remove_workdir().await;
|
||||
assert_eq!(
|
||||
result.disposition,
|
||||
WorkingDirectoryRemovalDisposition::AttentionRequired,
|
||||
"{result:?}"
|
||||
);
|
||||
assert_eq!(
|
||||
result.failure_category.as_deref(),
|
||||
Some("resource_busy"),
|
||||
"{result:?}"
|
||||
);
|
||||
assert!(result.retryable);
|
||||
assert!(self.materialization_root.exists());
|
||||
}
|
||||
|
||||
async fn assert_removed(&self) {
|
||||
let result = self.remove_workdir().await;
|
||||
assert_eq!(
|
||||
result.disposition,
|
||||
WorkingDirectoryRemovalDisposition::Removed,
|
||||
"{result:?}"
|
||||
);
|
||||
assert!(!result.retryable);
|
||||
assert!(!self.materialization_root.exists());
|
||||
assert!(
|
||||
self.api()
|
||||
.store
|
||||
.get_workdir_registry(TEST_WORKSPACE_ID, &self.workdir_id)
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum ClosePath {
|
||||
Registry,
|
||||
Attachment,
|
||||
}
|
||||
|
||||
async fn close(api: &WorkspaceApi, worker: &RuntimeWorkerRef, path: ClosePath) -> bool {
|
||||
let lock = current_worker_session_lock(api, worker);
|
||||
let _guard = lock.lock().await;
|
||||
match path {
|
||||
ClosePath::Registry => close_worker_workdir_sessions(&api.workdir_sessions, worker)
|
||||
.await
|
||||
.is_ok(),
|
||||
ClosePath::Attachment => close_current_worker_attachment_session_locked(api, worker)
|
||||
.await
|
||||
.is_ok(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn failure_restoration(path: ClosePath) {
|
||||
let fixture = Fixture::new().await;
|
||||
let worker = fixture.registry_owner();
|
||||
let session = fixture.open().await;
|
||||
fixture
|
||||
.api()
|
||||
.workdir_sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert_attachment(worker.clone(), session.clone());
|
||||
fixture.assert_busy().await;
|
||||
fixture.transport.fail_next(CloseFault::Reject);
|
||||
assert!(!close(fixture.api(), &worker, path).await);
|
||||
fixture.assert_restored(&worker, &session);
|
||||
fixture.assert_live(&session).await;
|
||||
fixture.assert_busy().await;
|
||||
assert!(close(fixture.api(), &worker, path).await);
|
||||
assert!(
|
||||
!fixture
|
||||
.api()
|
||||
.workdir_sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.attachments
|
||||
.contains_key(&worker)
|
||||
);
|
||||
let attempts = fixture.transport.attempts();
|
||||
assert_eq!(attempts.len(), 2, "retry must actually contact Runtime");
|
||||
assert_eq!(
|
||||
attempts[0], attempts[1],
|
||||
"retry closes the SAME Runtime session"
|
||||
);
|
||||
fixture.assert_removed().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancelled_backend_close_retry_waits_for_actual_provider_teardown_before_confirmation() {
|
||||
let fixture = Fixture::new().await;
|
||||
let gate = Arc::new(ProviderCloseGate {
|
||||
entered: Notify::new(),
|
||||
release: Semaphore::new(0),
|
||||
});
|
||||
fixture.source.hold_provider_close(gate.clone());
|
||||
*fixture.transport.direct_runtime.lock().unwrap() = Some(fixture.source.runtime_router());
|
||||
let worker = fixture.registry_owner();
|
||||
let session = fixture.open().await;
|
||||
session
|
||||
.start_command(workdir::CommandRequest {
|
||||
command: "sleep 60".into(),
|
||||
timeout_secs: 60,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
tool_call_id: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
fixture
|
||||
.api()
|
||||
.workdir_sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert_attachment(worker.clone(), session.clone());
|
||||
fixture.assert_busy().await;
|
||||
let api = fixture.api().clone();
|
||||
let first_worker = worker.clone();
|
||||
let first = tokio::spawn(async move { close(&api, &first_worker, ClosePath::Registry).await });
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), gate.entered.notified())
|
||||
.await
|
||||
.expect("actual Local provider teardown reached Pending");
|
||||
fixture.transport.dispatched.notified().await;
|
||||
first.abort();
|
||||
assert!(first.await.unwrap_err().is_cancelled());
|
||||
fixture.assert_restored(&worker, &session);
|
||||
fixture.assert_busy().await;
|
||||
|
||||
let api = fixture.api().clone();
|
||||
let retry_worker = worker.clone();
|
||||
let retry = tokio::spawn(async move { close(&api, &retry_worker, ClosePath::Registry).await });
|
||||
tokio::time::timeout(
|
||||
std::time::Duration::from_secs(5),
|
||||
fixture.transport.dispatched.notified(),
|
||||
)
|
||||
.await
|
||||
.expect("retry reached production Runtime handler");
|
||||
assert_eq!(
|
||||
*fixture.transport.pending_dispatches.lock().unwrap(),
|
||||
vec![true, true],
|
||||
"duplicate DELETE must wait, not acknowledge an extracted record"
|
||||
);
|
||||
assert!(
|
||||
!retry.is_finished(),
|
||||
"Backend cannot confirm while provider close is held"
|
||||
);
|
||||
assert!(fixture.transport.runtime_closes.lock().unwrap().is_empty());
|
||||
fixture.assert_busy().await;
|
||||
gate.release.add_permits(1);
|
||||
assert!(
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), retry)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
);
|
||||
assert!(
|
||||
!fixture
|
||||
.api()
|
||||
.workdir_sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.attachments
|
||||
.contains_key(&worker)
|
||||
);
|
||||
let attempts = fixture.transport.attempts();
|
||||
assert_eq!(attempts.len(), 2);
|
||||
assert_eq!(attempts[0], attempts[1]);
|
||||
fixture.assert_removed().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn registry_close_failure_restores_same_remote_session_and_retry_releases_runtime_lease() {
|
||||
failure_restoration(ClosePath::Registry).await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn attachment_close_failure_restores_same_remote_session_and_retry_releases_runtime_lease() {
|
||||
failure_restoration(ClosePath::Attachment).await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn lost_close_response_restores_handle_and_retry_confirms_actual_runtime_idempotent_release()
|
||||
{
|
||||
let fixture = Fixture::new().await;
|
||||
let worker = fixture.registry_owner();
|
||||
let session = fixture.open().await;
|
||||
fixture
|
||||
.api()
|
||||
.workdir_sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert_attachment(worker.clone(), session.clone());
|
||||
fixture.assert_busy().await;
|
||||
fixture
|
||||
.transport
|
||||
.fail_next(CloseFault::RejectAfterRuntimeClose);
|
||||
|
||||
assert!(!close(fixture.api(), &worker, ClosePath::Registry).await);
|
||||
fixture.assert_restored(&worker, &session);
|
||||
let first_attempt = fixture.transport.attempts();
|
||||
assert_eq!(first_attempt.len(), 1);
|
||||
assert_eq!(
|
||||
*fixture.transport.runtime_closes.lock().unwrap(),
|
||||
vec![(first_attempt[0].clone(), StatusCode::NO_CONTENT)],
|
||||
"fault must follow the real authenticated Runtime's successful DELETE"
|
||||
);
|
||||
let absent = session
|
||||
.stat(workdir::StatRequest {
|
||||
path: workdir::WorkdirPath::new("README.md").unwrap(),
|
||||
})
|
||||
.await
|
||||
.expect_err("original session must already be absent on Runtime");
|
||||
assert!(
|
||||
matches!(
|
||||
absent.classification_source(),
|
||||
workdir::WorkdirError::NotFound(_)
|
||||
),
|
||||
"must observe Runtime absence, not a local closed flag: {absent:?}"
|
||||
);
|
||||
|
||||
assert!(close(fixture.api(), &worker, ClosePath::Registry).await);
|
||||
let attempts = fixture.transport.attempts();
|
||||
assert_eq!(
|
||||
attempts,
|
||||
vec![first_attempt[0].clone(), first_attempt[0].clone()],
|
||||
"retry must send DELETE to the same already-absent Runtime session"
|
||||
);
|
||||
assert_eq!(
|
||||
*fixture.transport.runtime_closes.lock().unwrap(),
|
||||
vec![
|
||||
(attempts[0].clone(), StatusCode::NO_CONTENT),
|
||||
(attempts[1].clone(), StatusCode::NO_CONTENT)
|
||||
],
|
||||
"Runtime itself must acknowledge the idempotent retry"
|
||||
);
|
||||
assert!(
|
||||
!fixture
|
||||
.api()
|
||||
.workdir_sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.attachments
|
||||
.contains_key(&worker)
|
||||
);
|
||||
// Runtime has released its leases, but only now has Backend confirmed it.
|
||||
fixture.assert_removed().await;
|
||||
}
|
||||
|
||||
async fn cancellation_restoration(path: ClosePath) {
|
||||
let fixture = Fixture::new().await;
|
||||
let worker = fixture.registry_owner();
|
||||
let session = fixture.open().await;
|
||||
fixture
|
||||
.api()
|
||||
.workdir_sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert_attachment(worker.clone(), session.clone());
|
||||
// A second real Runtime session proves the batch retains not-yet-attempted
|
||||
// command sources as well as the close currently awaiting a response.
|
||||
let command = if matches!(path, ClosePath::Registry) {
|
||||
let source = fixture.open().await;
|
||||
let handle = source
|
||||
.start_command(workdir::CommandRequest {
|
||||
command: "printf close-test".into(),
|
||||
timeout_secs: 5,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
tool_call_id: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let external = fixture
|
||||
.api()
|
||||
.workdir_sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.register_command(worker.clone(), "checkout".into(), source.clone(), handle);
|
||||
Some((external, source))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
fixture.assert_busy().await;
|
||||
fixture.transport.fail_next(CloseFault::HoldThenReject);
|
||||
let api = fixture.api().clone();
|
||||
let closing_worker = worker.clone();
|
||||
let closing = tokio::spawn(async move { close(&api, &closing_worker, path).await });
|
||||
tokio::time::timeout(
|
||||
std::time::Duration::from_secs(5),
|
||||
fixture.transport.entered.notified(),
|
||||
)
|
||||
.await
|
||||
.expect("close reached pre-Runtime hold");
|
||||
closing.abort();
|
||||
assert!(closing.await.unwrap_err().is_cancelled());
|
||||
fixture.assert_restored(&worker, &session);
|
||||
fixture.assert_live(&session).await;
|
||||
if let Some((external, source)) = &command {
|
||||
let restored = fixture
|
||||
.api()
|
||||
.workdir_sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.command(&worker, "checkout", external)
|
||||
.expect("unattempted command session retained");
|
||||
assert!(Arc::ptr_eq(&restored.source, source));
|
||||
}
|
||||
fixture.transport.release.add_permits(1);
|
||||
fixture.assert_busy().await;
|
||||
assert!(close(fixture.api(), &worker, path).await);
|
||||
let attempts = fixture.transport.attempts();
|
||||
assert_eq!(attempts.len(), if command.is_some() { 3 } else { 2 });
|
||||
assert_eq!(
|
||||
attempts[0], attempts[1],
|
||||
"retry closes cancelled session, not a replacement"
|
||||
);
|
||||
fixture.assert_removed().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn registry_close_cancellation_restores_inflight_and_unattempted_remote_sessions() {
|
||||
cancellation_restoration(ClosePath::Registry).await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn attachment_close_cancellation_restores_remote_session_for_retry() {
|
||||
cancellation_restoration(ClosePath::Attachment).await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn worker_removal_close_failure_retains_attachment_and_retry_releases_runtime_lease() {
|
||||
let fixture = Fixture::new().await;
|
||||
let api = fixture.api();
|
||||
fixture.source.publish_profile("builtin:companion");
|
||||
let Json(created) = create_workspace_worker(
|
||||
State(api.clone()),
|
||||
HeaderMap::new(),
|
||||
Json(CreateWorkspaceWorkerRequest {
|
||||
runtime_id: fixture.runtime_id.clone(),
|
||||
display_name: "session-close removal".into(),
|
||||
singleton_key: None,
|
||||
profile: Some("builtin:companion".into()),
|
||||
ticket_assignment: None,
|
||||
initial_submit: Vec::new(),
|
||||
workdir_attachments: Vec::new(),
|
||||
feature_connections: Default::default(),
|
||||
control_operation_id: None,
|
||||
execution_grant: None,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let worker = RuntimeWorkerRef::new(created.runtime_id, created.worker_id);
|
||||
api.runtime
|
||||
.stop_worker(
|
||||
&worker,
|
||||
WorkerLifecycleRequest {
|
||||
reason: Some("prepare removal".into()),
|
||||
ticket_assignment: None,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
let summary = api.runtime.worker(&worker).unwrap();
|
||||
assert_eq!(summary.state, "stopped");
|
||||
sync_worker_observation(api, &summary).unwrap();
|
||||
fixture.source.seed_remote_session_for_cleanup(&worker);
|
||||
let session = fixture.open().await;
|
||||
api.workdir_sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert_attachment(worker.clone(), session.clone());
|
||||
fixture.assert_busy().await;
|
||||
api.store
|
||||
.attach_worker_workdir(&WorkerWorkdirLinkRecord {
|
||||
connection_id: String::new(),
|
||||
workspace_id: TEST_WORKSPACE_ID.into(),
|
||||
worker: worker.clone(),
|
||||
workdir_id: fixture.workdir_id.clone(),
|
||||
alias: "checkout".into(),
|
||||
capabilities: workdir::WorkdirSessionCapabilities::ALL,
|
||||
linked_at: TEST_CREATED_AT.into(),
|
||||
unlinked_at: None,
|
||||
})
|
||||
.unwrap();
|
||||
fixture.transport.fail_next(CloseFault::Reject);
|
||||
let removal = WorkerRemovalService::new(api);
|
||||
let failed = removal
|
||||
.execute_target_removal(&api.runtime, &worker, "remove stopped Worker", None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
failed.status,
|
||||
StatusCode::SERVICE_UNAVAILABLE.as_u16(),
|
||||
"{}",
|
||||
failed.body
|
||||
);
|
||||
assert!(
|
||||
failed.body.contains("attachment_close_failed"),
|
||||
"{}",
|
||||
failed.body
|
||||
);
|
||||
fixture.assert_restored(&worker, &session);
|
||||
fixture.assert_live(&session).await;
|
||||
assert!(
|
||||
api.store
|
||||
.list_worker_workdir_links(TEST_WORKSPACE_ID, &worker)
|
||||
.unwrap()
|
||||
.iter()
|
||||
.any(|link| link.unlinked_at.is_none())
|
||||
);
|
||||
assert!(
|
||||
api.store
|
||||
.get_worker_registry(TEST_WORKSPACE_ID, &worker)
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
let operation = api
|
||||
.store
|
||||
.recover_worker_removal_execution(TEST_WORKSPACE_ID, &worker)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
operation.plan.state,
|
||||
crate::retention::WorkerRemovalPlanState::Failed
|
||||
);
|
||||
let retried = removal
|
||||
.execute_target_removal(&api.runtime, &worker, "remove stopped Worker", None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(retried.status, StatusCode::OK.as_u16(), "{}", retried.body);
|
||||
assert!(
|
||||
api.store
|
||||
.get_worker_registry(TEST_WORKSPACE_ID, &worker)
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
api.store
|
||||
.list_worker_workdir_links(TEST_WORKSPACE_ID, &worker)
|
||||
.unwrap()
|
||||
.iter()
|
||||
.all(|link| link.unlinked_at.is_some())
|
||||
);
|
||||
let attempts = fixture.transport.attempts();
|
||||
assert_eq!(attempts.len(), 2);
|
||||
assert_eq!(attempts[0], attempts[1]);
|
||||
fixture.assert_removed().await;
|
||||
}
|
||||
Reference in New Issue
Block a user