Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: fence pending queue controls

Browse Source
This commit is contained in:
Hare
2026-09-05 23:16:19 +09:00
parent bb56283063
commit 5b0a6691f8
11 changed files with 347 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+54 -32
View File
@@ -232,8 +232,9 @@ enum PendingRun {
fn prepare_pending_run<St: Store + Clone>(
pending_submissions: &crate::worker::PendingSubmissionHandle<St>,
notify_buffer: &NotifyBuffer,
fence: Option<(u64, &str)>,
) -> Result<Option<PendingRun>, crate::worker::PendingSubmissionError> {
Ok(match pending_submissions.prepare_next_activation()? {
Ok(match pending_submissions.prepare_next_activation(fence)? {
Some(crate::worker::PendingActivation::Submission(submission)) => {
Some(PendingRun::Submit(submission))
}
@@ -1420,7 +1421,7 @@ async fn controller_loop<C, St>(
}
if !shutdown && may_drain_pending && new_status == WorkerStatus::Idle {
match prepare_pending_run(&pending_submissions, &notify_buffer) {
match prepare_pending_run(&pending_submissions, &notify_buffer, None) {
Ok(Some(next)) => {
pending = Some(next);
new_status = WorkerStatus::Running;
@@ -1505,17 +1506,18 @@ async fn controller_loop<C, St>(
if auto_run {
match pending_submissions.accept_notification(notification_request_id, message)
{
Ok(true) => match prepare_pending_run(&pending_submissions, &notify_buffer)
{
Ok(Some(next)) => pending = Some(next),
Ok(None) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: error.to_string(),
});
Ok(true) => {
match prepare_pending_run(&pending_submissions, &notify_buffer, None) {
Ok(Some(next)) => pending = Some(next),
Ok(None) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: error.to_string(),
});
}
}
},
}
Ok(false) => {}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
@@ -1534,8 +1536,24 @@ async fn controller_loop<C, St>(
pending: pending_submissions.snapshot(),
});
}
Method::CancelPendingSubmission { submission_id } => {
match pending_submissions.cancel(&submission_id) {
Method::CancelPendingSubmission {
submission_id,
expected_revision,
} => match pending_submissions.cancel(&submission_id, expected_revision) {
Ok(pending_snapshot) => {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_snapshot,
});
}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
}
},
Method::ClearPendingSubmissions { expected_revision } => {
match pending_submissions.clear(expected_revision) {
Ok(pending_snapshot) => {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_snapshot,
@@ -1549,21 +1567,22 @@ async fn controller_loop<C, St>(
}
}
}
Method::ClearPendingSubmissions => match pending_submissions.clear() {
Ok(pending_snapshot) => {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged {
pending: pending_snapshot,
});
}
Err(error) => {
Method::ContinuePending {
expected_revision,
expected_head_id,
} => {
if shared_state.get_status() != WorkerStatus::Idle {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: error.to_string(),
code: ErrorCode::InvalidRequest,
message: "ContinuePending requires an idle Worker; Resume or Cancel a paused run first".into(),
});
continue;
}
},
Method::ContinuePending => {
match prepare_pending_run(&pending_submissions, &notify_buffer) {
match prepare_pending_run(
&pending_submissions,
&notify_buffer,
Some((expected_revision, &expected_head_id)),
) {
Ok(Some(next)) => pending = Some(next),
Ok(None) => {
let _ = working_event_tx.send(Event::Error {
@@ -2077,7 +2096,7 @@ where
}
}
}
Some(Method::Resume | Method::ContinuePending) => {
Some(Method::Resume | Method::ContinuePending { .. }) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn".into(),
@@ -2088,8 +2107,11 @@ where
pending: pending_submissions.snapshot(),
});
}
Some(Method::CancelPendingSubmission { submission_id }) => {
match pending_submissions.cancel(&submission_id) {
Some(Method::CancelPendingSubmission {
submission_id,
expected_revision,
}) => {
match pending_submissions.cancel(&submission_id, expected_revision) {
Ok(pending) => {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged { pending });
}
@@ -2101,14 +2123,14 @@ where
}
}
}
Some(Method::ClearPendingSubmissions) => {
match pending_submissions.clear() {
Some(Method::ClearPendingSubmissions { expected_revision }) => {
match pending_submissions.clear(expected_revision) {
Ok(pending) => {
let _ = working_event_tx.send(Event::PendingSubmissionsChanged { pending });
}
Err(error) => {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
}
crates/worker/src/worker.rs
+67 -5
View File
@@ -130,9 +130,20 @@ pub(crate) struct PendingActivationState {
impl PendingActivationState {
pub(crate) fn snapshot(&self) -> protocol::PendingSubmissionsSnapshot {
let head_id = match (self.pending.front(), self.pending_notifications.front()) {
(Some(submission), Some(notification))
if notification.activation_sequence < submission.activation_sequence =>
{
Some(notification.notification_request_id.clone())
}
(Some(submission), _) => Some(submission.submission_id.clone()),
(None, Some(notification)) => Some(notification.notification_request_id.clone()),
(None, None) => None,
};
protocol::PendingSubmissionsSnapshot {
revision: self.revision,
notification_count: u32::try_from(self.pending_notifications.len()).unwrap_or(u32::MAX),
head_id,
submissions: self
.pending
.iter()
@@ -1125,6 +1136,13 @@ pub(crate) enum PendingSubmissionError {
ByteLimit,
#[error("pending submission artifact references exceed {MAX_PENDING_ARTIFACT_REFS}")]
ArtifactLimit,
#[error("pending queue revision conflict: expected {expected}, current {current}")]
RevisionConflict { expected: u64, current: u64 },
#[error("pending queue head conflict: expected {expected}, current {current:?}")]
HeadConflict {
expected: String,
current: Option<String>,
},
#[error("pending submission not found: {0}")]
NotFound(String),
#[error("pending submission state persistence failed: {0}")]
@@ -1141,6 +1159,29 @@ impl<St> PendingSubmissionHandle<St>
where
St: Store + Clone,
{
fn validate_fence(
state: &PendingActivationState,
expected_revision: u64,
expected_head_id: Option<&str>,
) -> Result<(), PendingSubmissionError> {
if state.revision != expected_revision {
return Err(PendingSubmissionError::RevisionConflict {
expected: expected_revision,
current: state.revision,
});
}
if let Some(expected) = expected_head_id {
let current = state.snapshot().head_id;
if current.as_deref() != Some(expected) {
return Err(PendingSubmissionError::HeadConflict {
expected: expected.to_owned(),
current,
});
}
}
Ok(())
}
fn persist_locked(&self, state: &PendingActivationState) -> Result<(), PendingSubmissionError> {
self.writer.append_entry_locked(LogEntry::Extension {
ts: segment_log::now_millis(),
@@ -1354,6 +1395,7 @@ where
pub(crate) fn prepare_next_activation(
&self,
fence: Option<(u64, &str)>,
) -> Result<Option<PendingActivation>, PendingSubmissionError> {
let _append_guard = self
.writer
@@ -1365,6 +1407,9 @@ where
.state
.lock()
.expect("pending activation state poisoned");
if let Some((expected_revision, expected_head_id)) = fence {
Self::validate_fence(&state, expected_revision, Some(expected_head_id))?;
}
if state.activating.is_some() || state.activating_notification.is_some() {
return Ok(None);
}
@@ -1495,6 +1540,7 @@ where
pub(crate) fn cancel(
&self,
submission_id: &str,
expected_revision: u64,
) -> Result<protocol::PendingSubmissionsSnapshot, PendingSubmissionError> {
let _append_guard = self
.writer
@@ -1506,6 +1552,7 @@ where
.state
.lock()
.expect("pending activation state poisoned");
Self::validate_fence(&state, expected_revision, None)?;
let original = state.clone();
let Some(index) = state
.pending
@@ -1525,6 +1572,7 @@ where
pub(crate) fn clear(
&self,
expected_revision: u64,
) -> Result<protocol::PendingSubmissionsSnapshot, PendingSubmissionError> {
let _append_guard = self
.writer
@@ -1536,6 +1584,7 @@ where
.state
.lock()
.expect("pending activation state poisoned");
Self::validate_fence(&state, expected_revision, None)?;
let original = state.clone();
state.pending.clear();
state.pending_notifications.clear();
@@ -9189,10 +9238,23 @@ mod build_summary_prompt_tests {
assert_eq!(restored.pending.len(), 1);
assert_eq!(restored.pending[0].submission_id, accepted.submission_id);
let snapshot = handle.cancel(&accepted.submission_id).unwrap();
let fence = handle.snapshot();
assert!(matches!(
handle.cancel(&accepted.submission_id, fence.revision.saturating_sub(1)),
Err(PendingSubmissionError::RevisionConflict { .. })
));
assert!(matches!(
handle.prepare_next_activation(Some((fence.revision, "wrong-head"))),
Err(PendingSubmissionError::HeadConflict { .. })
));
assert_eq!(handle.snapshot(), fence);
let snapshot = handle
.cancel(&accepted.submission_id, handle.snapshot().revision)
.unwrap();
assert!(snapshot.submissions.is_empty());
assert!(matches!(
handle.cancel(&accepted.submission_id),
handle.cancel(&accepted.submission_id, handle.snapshot().revision),
Err(PendingSubmissionError::NotFound(_))
));
@@ -9210,7 +9272,7 @@ mod build_summary_prompt_tests {
Err(PendingSubmissionError::CountLimit)
));
assert_eq!(handle.snapshot().submissions.len(), MAX_PENDING_SUBMISSIONS);
let cleared = handle.clear().unwrap();
let cleared = handle.clear(handle.snapshot().revision).unwrap();
assert!(cleared.submissions.is_empty());
assert_eq!(cleared.notification_count, 0);
}
@@ -9237,7 +9299,7 @@ mod build_summary_prompt_tests {
.accept("request-1".into(), vec![Segment::text("submit")], false)
.unwrap();
let first = handle.prepare_next_activation().unwrap().unwrap();
let first = handle.prepare_next_activation(None).unwrap().unwrap();
assert!(matches!(
first,
PendingActivation::Notification(PendingNotification { ref message, .. })
@@ -9249,7 +9311,7 @@ mod build_summary_prompt_tests {
assert!(committed_state.pending_notifications.is_empty());
assert!(committed_state.activating_notification.is_none());
handle.finish_notification_activation("notification-1");
let second = handle.prepare_next_activation().unwrap().unwrap();
let second = handle.prepare_next_activation(None).unwrap().unwrap();
assert!(matches!(second, PendingActivation::Submission(_)));
}
crates/worker/tests/controller_test.rs
+27 -4
View File
@@ -1322,7 +1322,7 @@ async fn submit_while_running_is_durably_queued() {
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
let mut accepted = None;
let mut pending_count = None;
let mut pending_snapshot = None;
while tokio::time::Instant::now() < deadline {
match tokio::time::timeout(std::time::Duration::from_millis(100), rx.recv()).await {
Ok(Ok(Event::SubmissionAccepted {
@@ -1333,21 +1333,44 @@ async fn submit_while_running_is_durably_queued() {
Ok(Ok(Event::PendingSubmissionsChanged { pending }))
if pending.submissions.len() == 1 =>
{
pending_count = Some(1)
pending_snapshot = Some(pending)
}
Ok(Ok(Event::Error { code, message })) if code == worker::ErrorCode::AlreadyRunning => {
panic!("Submit was busy-rejected: {message}")
}
_ => {}
}
if accepted.is_some() && pending_count.is_some() {
if accepted.is_some() && pending_snapshot.is_some() {
break;
}
}
assert_eq!(accepted, Some(protocol::SubmissionDisposition::Queued));
assert_eq!(pending_count, Some(1));
let pending_snapshot = pending_snapshot.expect("pending snapshot");
assert_eq!(pending_snapshot.submissions.len(), 1);
handle.send(Method::Pause).await.unwrap();
wait_for_status(&handle, WorkerStatus::Paused).await;
handle
.send(Method::ContinuePending {
expected_revision: pending_snapshot.revision,
expected_head_id: pending_snapshot.head_id.expect("pending head"),
})
.await
.unwrap();
let rejection = tokio::time::timeout(std::time::Duration::from_secs(1), async {
loop {
if let Ok(Event::Error { code, message }) = rx.recv().await
&& code == worker::ErrorCode::InvalidRequest
&& message.contains("requires an idle Worker")
{
break message;
}
}
})
.await
.expect("paused ContinuePending rejection");
assert!(rejection.contains("Resume or Cancel"));
assert_eq!(handle.shared_state.get_status(), WorkerStatus::Paused);
}
#[tokio::test]