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

refactor: make worker state snapshots authoritative

Browse Source
This commit is contained in:
Hare
2026-09-17 02:02:51 +09:00
parent d6bef5d1c7
commit 10ebac142e
31 changed files with 473 additions and 1132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workspace-server/src/hosts.rs
-1
View File
@@ -5836,7 +5836,6 @@ mod tests {
self.backend_id(),
),
worker_state: protocol::WorkerStateSnapshot {
execution_generation: request.run_generation,
..protocol::WorkerStatus::Idle.into()
},
working_directory: request
crates/workspace-server/src/latest_schema.sql
+1 -1
View File
@@ -775,7 +775,7 @@ CREATE TABLE "worker_registry" (
CREATE TABLE worker_removal_operations (
operation_id TEXT PRIMARY KEY, plan_id TEXT NOT NULL UNIQUE, input_fingerprint TEXT NOT NULL,
workspace_id TEXT NOT NULL, runtime_id TEXT NOT NULL, worker_id TEXT NOT NULL,
worker_revision TEXT NOT NULL, run_generation INTEGER NOT NULL CHECK(run_generation>=0),
worker_revision TEXT NOT NULL,
policy_id TEXT NOT NULL, policy_revision INTEGER NOT NULL,
session_disposition TEXT NOT NULL, metadata_disposition TEXT NOT NULL,
archive_retention_kind TEXT NOT NULL, archive_retention_seconds INTEGER,
crates/workspace-server/src/retention.rs
+19 -27
View File
@@ -91,7 +91,6 @@ pub struct WorkerRemovalPlan {
pub workspace_id: String,
pub worker: RuntimeWorkerRef,
pub worker_revision: String,
pub run_generation: u64,
pub policy_id: String,
pub policy_revision: u64,
pub session_disposition: SessionDisposition,
@@ -219,7 +218,7 @@ impl SqliteWorkspaceStore {
let plan_id=stable("wrp",&fp); let operation_id=stable("wro",&fp);
let archive_id=(policy.session_disposition==SessionDisposition::Archive).then(||stable("wra",&fp));
let state=if blockers.is_empty(){WorkerRemovalPlanState::Planned}else{WorkerRemovalPlanState::Blocked};
tx.execute("INSERT OR IGNORE INTO worker_removal_operations(operation_id,plan_id,input_fingerprint,workspace_id,runtime_id,worker_id,worker_revision,run_generation,policy_id,policy_revision,session_disposition,metadata_disposition,archive_retention_kind,archive_retention_seconds,diagnostics_disposition,diagnostics_retention_seconds,archive_id,blockers_json,state,reason,created_at,updated_at) VALUES(?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13,?14,?15,?16,?17,?18,?19,?20,?21,?21)",params![operation_id,plan_id,fp,req.workspace_id,req.worker.runtime_id,req.worker.worker_id,worker.updated_at,inv.run_generation,policy.policy_id,policy.revision,sess(policy.session_disposition),meta(policy.metadata_disposition),archive_kind(policy.archive_retention),archive_seconds(policy.archive_retention),diag(policy.diagnostics_disposition),policy.diagnostics_retention_seconds,archive_id,serde_json::to_string(&blockers).map_err(|e|StoreError::InvalidInput(e.to_string()))?,state_s(state),req.reason,now])?;
tx.execute("INSERT OR IGNORE INTO worker_removal_operations(operation_id,plan_id,input_fingerprint,workspace_id,runtime_id,worker_id,worker_revision,policy_id,policy_revision,session_disposition,metadata_disposition,archive_retention_kind,archive_retention_seconds,diagnostics_disposition,diagnostics_retention_seconds,archive_id,blockers_json,state,reason,created_at,updated_at) VALUES(?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13,?14,?15,?16,?17,?18,?19,?20,?20)",params![operation_id,plan_id,fp,req.workspace_id,req.worker.runtime_id,req.worker.worker_id,worker.updated_at,policy.policy_id,policy.revision,sess(policy.session_disposition),meta(policy.metadata_disposition),archive_kind(policy.archive_retention),archive_seconds(policy.archive_retention),diag(policy.diagnostics_disposition),policy.diagnostics_retention_seconds,archive_id,serde_json::to_string(&blockers).map_err(|e|StoreError::InvalidInput(e.to_string()))?,state_s(state),req.reason,now])?;
let plan=load_plan(&tx,&plan_id)?.ok_or_else(||StoreError::InvalidInput("plan missing".into()))?;
if plan.input_fingerprint!=fp{return Err(StoreError::InvalidInput(format!("fingerprint:{}",plan.operation_id)));}
tx.commit()?; Ok(plan)
@@ -304,7 +303,6 @@ impl SqliteWorkspaceStore {
source_runtime_id: plan.worker.runtime_id.clone(),
worker_id: worker_id,
expected_worker_revision: plan.worker_revision.clone(),
expected_run_generation: plan.run_generation,
source_created_at: worker.created_at,
removed_at,
effective_profile: worker.profile,
@@ -376,7 +374,6 @@ impl SqliteWorkspaceStore {
source_runtime_id: plan.worker.runtime_id.clone(),
worker_id: worker_id,
expected_worker_revision: plan.worker_revision.clone(),
expected_run_generation: plan.run_generation,
source_created_at: worker
.as_ref()
.map(|worker| worker.created_at.clone())
@@ -682,20 +679,20 @@ fn load_plan_op(c: &Connection, id: &str) -> crate::Result<Option<WorkerRemovalP
fn load_plan_q(c: &Connection, key: &str, id: &str) -> crate::Result<Option<WorkerRemovalPlan>> {
let query = format!(
"SELECT plan_id,operation_id,input_fingerprint,workspace_id,runtime_id,worker_id,
worker_revision,run_generation,policy_id,policy_revision,session_disposition,
worker_revision,policy_id,policy_revision,session_disposition,
metadata_disposition,archive_retention_kind,archive_retention_seconds,
diagnostics_disposition,diagnostics_retention_seconds,archive_id,blockers_json,
state,reason,created_at,updated_at,failure_category
FROM worker_removal_operations WHERE {key}=?1"
);
c.query_row(&query, params![id], |row| {
let session: String = row.get(10)?;
let metadata: String = row.get(11)?;
let archive_kind: String = row.get(12)?;
let archive_seconds: Option<i64> = row.get(13)?;
let diagnostics: String = row.get(14)?;
let blockers: String = row.get(17)?;
let state: String = row.get(18)?;
let session: String = row.get(9)?;
let metadata: String = row.get(10)?;
let archive_kind: String = row.get(11)?;
let archive_seconds: Option<i64> = row.get(12)?;
let diagnostics: String = row.get(13)?;
let blockers: String = row.get(16)?;
let state: String = row.get(17)?;
Ok(WorkerRemovalPlan {
plan_id: row.get(0)?,
operation_id: row.get(1)?,
@@ -706,27 +703,26 @@ fn load_plan_q(c: &Connection, key: &str, id: &str) -> crate::Result<Option<Work
worker_id: row.get(5)?,
},
worker_revision: row.get(6)?,
run_generation: row.get::<_, i64>(7)? as u64,
policy_id: row.get(8)?,
policy_revision: row.get::<_, i64>(9)? as u64,
policy_id: row.get(7)?,
policy_revision: row.get::<_, i64>(8)? as u64,
session_disposition: parse_s(&session)?,
metadata_disposition: parse_m(&metadata)?,
archive_retention: parse_archive(&archive_kind, archive_seconds)?,
diagnostics_disposition: parse_d(&diagnostics)?,
diagnostics_retention_seconds: row.get::<_, Option<i64>>(15)?.map(|v| v as u64),
archive_id: row.get(16)?,
diagnostics_retention_seconds: row.get::<_, Option<i64>>(14)?.map(|v| v as u64),
archive_id: row.get(15)?,
blockers: serde_json::from_str(&blockers).map_err(|error| {
rusqlite::Error::FromSqlConversionFailure(
17,
16,
rusqlite::types::Type::Text,
Box::new(error),
)
})?,
state: parse_state(&state)?,
reason: row.get(19)?,
created_at: row.get(20)?,
updated_at: row.get(21)?,
failure_category: row.get(22)?,
reason: row.get(18)?,
created_at: row.get(19)?,
updated_at: row.get(20)?,
failure_category: row.get(21)?,
})
})
.optional()
@@ -755,7 +751,6 @@ fn fingerprint(
r.worker.runtime_id,
r.worker.worker_id,
worker_revision,
i.run_generation,
i.session_id,
i.segment_ids,
p.policy_id,
@@ -992,7 +987,6 @@ mod tests {
workspace_id: "w".into(),
runtime_id: "r".into(),
worker_id: worker_id(),
run_generation: 2,
session_id: Some("s".into()),
segment_ids: vec!["a".into()],
session_bytes: 1,
@@ -1139,13 +1133,12 @@ mod tests {
}
#[test]
fn prepared_execution_is_derived_from_pinned_plan_generation() {
fn prepared_execution_is_derived_from_pinned_plan() {
let s = setup();
let plan = s.plan_worker_removal(&req(), &inv()).unwrap();
let prepared = s
.prepare_worker_removal_execution("w", &plan.plan_id, &plan.input_fingerprint)
.unwrap();
assert_eq!(prepared.runtime_request.expected_run_generation, 2);
assert_eq!(
prepared.runtime_request.session_disposition,
SessionDisposition::Archive
@@ -1264,7 +1257,6 @@ mod tests {
workspace_id: "w".into(),
runtime_id: "r".into(),
worker_id: WorkerId::from_legacy_u64(2),
run_generation: 1,
session_id: Some("orphan-session".into()),
segment_ids: vec![],
session_bytes: 10,
crates/workspace-server/src/runtime_subscription_tests.rs
-3
View File
@@ -20,7 +20,6 @@ impl WorkerExecutionBackend for TestExecutionBackend {
WorkerExecutionSpawnResult::connected(
WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
protocol::WorkerStateSnapshot {
execution_generation: request.run_generation,
..protocol::WorkerStatus::Idle.into()
},
None,
@@ -182,7 +181,6 @@ async fn equal_downstream_selectors_share_one_upstream_subscription() {
.worker_state
.clone()
.expect("connected test Worker must expose its initial state");
running.revision += 1;
running.state = protocol::WorkerState::Busy(protocol::WorkerBusyState::Run(
protocol::WorkerRunState::Running,
));
@@ -348,7 +346,6 @@ async fn embedded_runtime_uses_in_process_subscription_source() {
.worker_state
.clone()
.expect("connected test Worker must expose its initial state");
running.revision += 1;
running.state = protocol::WorkerState::Busy(protocol::WorkerBusyState::Run(
protocol::WorkerRunState::Running,
));
crates/workspace-server/src/server.rs
+5 -9
View File
@@ -21861,7 +21861,6 @@ mod tests {
self.backend_id(),
),
worker_state: protocol::WorkerStateSnapshot {
execution_generation: request.run_generation,
..protocol::WorkerStatus::Idle.into()
},
working_directory,
@@ -21901,10 +21900,11 @@ mod tests {
context_window: 0,
context_tokens: 0,
},
state: protocol::WorkerStateSnapshot::initial(1),
state: protocol::WorkerStateSnapshot::initial(),
in_flight: protocol::InFlightSnapshot {
blocks: Vec::new(),
commands: Vec::new(),
compaction: None,
},
internal_workers: Vec::new(),
})
@@ -21965,12 +21965,12 @@ mod tests {
uuid::Uuid::now_v7().to_string(),
protocol::SubmissionDisposition::Started,
)
.with_worker_state(protocol::WorkerStateSnapshot::initial(1))
.with_worker_state(protocol::WorkerStateSnapshot::initial())
} else {
worker_runtime::execution::WorkerExecutionResult::accepted(
worker_runtime::execution::WorkerExecutionOperation::Input,
)
.with_worker_state(protocol::WorkerStateSnapshot::initial(1))
.with_worker_state(protocol::WorkerStateSnapshot::initial())
}
}
}
@@ -31548,11 +31548,7 @@ mod tests {
protocol::subscription::SubscriptionWorkerProtocolMethod {
subscription_id: second_protocol_subscription_id,
method: protocol::Method::Resume {
command: protocol::WorkerCommandEnvelope {
command_id: 1,
expected_execution_generation: 1,
expected_worker_state_revision: 0,
},
command: protocol::WorkerCommandEnvelope { command_id: 1 },
},
},
),
crates/workspace-server/src/store.rs
+77 -4
View File
@@ -18,7 +18,7 @@ use crate::workspace_deletion::WorkspaceDeletionStore;
use crate::{Error, Result};
const OLDEST_SCHEMA_VERSION: i64 = 50;
const LATEST_SCHEMA_VERSION: i64 = 60;
const LATEST_SCHEMA_VERSION: i64 = 61;
const SCHEMA_BASELINE_NAME: &str = "workspace schema baseline";
const WORKSPACE_RUNTIME_BINDINGS_MIGRATION_NAME: &str = "workspace runtime bindings";
const RUNTIME_BINDING_AUDIT_MIGRATION_NAME: &str = "workspace Runtime binding revision and audit";
@@ -35,6 +35,7 @@ const REMOVE_WORKDIR_CACHE_GENERATION_MIGRATION_NAME: &str =
const WORKDIR_CREDENTIAL_CANDIDATE_SNAPSHOT_MIGRATION_NAME: &str =
"Workdir create credential candidate snapshots";
const RUNTIME_REMOVAL_OPERATION_MIGRATION_NAME: &str = "guarded Runtime removal operations";
const REMOVE_WORKER_RUN_GENERATION_MIGRATION_NAME: &str = "remove obsolete Worker run generation";
const MIGRATIONS: &[Migration] = &[
Migration {
@@ -87,6 +88,11 @@ const MIGRATIONS: &[Migration] = &[
name: RUNTIME_REMOVAL_OPERATION_MIGRATION_NAME,
apply: migrate_runtime_removal_operations_v59_to_v60,
},
Migration {
version: 61,
name: REMOVE_WORKER_RUN_GENERATION_MIGRATION_NAME,
apply: migrate_worker_run_generation_v60_to_v61,
},
];
#[derive(Clone, Copy)]
@@ -9767,11 +9773,22 @@ fn migrate_runtime_removal_operations_v59_to_v60(conn: &Connection) -> Result<()
BEGIN SELECT RAISE(ABORT, 'runtime_removal_in_progress'); END;
"#,
)?;
tx.execute(
"INSERT INTO __yoi_schema_migrations (version, name) VALUES (?1, ?2)",
params![60, RUNTIME_REMOVAL_OPERATION_MIGRATION_NAME],
)?;
tx.commit()?;
Ok(())
}
fn migrate_worker_run_generation_v60_to_v61(conn: &Connection) -> Result<()> {
let tx = rusqlite::Transaction::new_unchecked(conn, TransactionBehavior::Exclusive)?;
tx.execute_batch("ALTER TABLE worker_removal_operations DROP COLUMN run_generation;")?;
tx.execute(
"INSERT INTO __yoi_schema_migrations (version, name) VALUES (?1, ?2)",
params![
LATEST_SCHEMA_VERSION,
RUNTIME_REMOVAL_OPERATION_MIGRATION_NAME
REMOVE_WORKER_RUN_GENERATION_MIGRATION_NAME
],
)?;
tx.commit()?;
@@ -10779,6 +10796,8 @@ mod tests {
DROP TRIGGER workdir_removal_insert_blocked_by_runtime_removal; \
DROP TRIGGER workdir_removal_update_blocked_by_runtime_removal; \
DROP TABLE runtime_removal_operations; \
ALTER TABLE worker_removal_operations \
ADD COLUMN run_generation INTEGER NOT NULL DEFAULT 0; \
DELETE FROM __yoi_schema_migrations; \
INSERT INTO __yoi_schema_migrations(version, name) \
VALUES (59, 'workspace schema baseline');",
@@ -10818,6 +10837,50 @@ mod tests {
.unwrap();
}
#[test]
fn schema_v60_upgrade_removes_worker_run_generation_column() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().join("server.db");
drop(SqliteWorkspaceStore::open(&path).unwrap());
{
let conn = Connection::open(&path).unwrap();
conn.execute_batch(
"ALTER TABLE worker_removal_operations \
ADD COLUMN run_generation INTEGER NOT NULL DEFAULT 0; \
DELETE FROM __yoi_schema_migrations; \
INSERT INTO __yoi_schema_migrations(version,name) \
VALUES (60,'workspace schema baseline');",
)
.unwrap();
}
drop(SqliteWorkspaceStore::open(&path).unwrap());
let conn = Connection::open(&path).unwrap();
assert_eq!(
current_schema_version(&conn).unwrap(),
LATEST_SCHEMA_VERSION
);
let columns = conn
.prepare("PRAGMA table_info(worker_removal_operations)")
.unwrap()
.query_map([], |row| row.get::<_, String>(1))
.unwrap()
.collect::<rusqlite::Result<Vec<_>>>()
.unwrap();
assert!(!columns.iter().any(|column| column == "run_generation"));
assert_eq!(
conn.query_row("PRAGMA quick_check", [], |row| row.get::<_, String>(0))
.unwrap(),
"ok"
);
let foreign_key_failures: i64 = conn
.query_row("SELECT COUNT(*) FROM pragma_foreign_key_check", [], |row| {
row.get(0)
})
.unwrap();
assert_eq!(foreign_key_failures, 0);
}
#[test]
fn schema_v59_runtime_removal_migration_rolls_back_partial_ddl_and_marker() {
let temp = tempfile::tempdir().unwrap();
@@ -10948,6 +11011,8 @@ mod tests {
DROP TRIGGER workdir_removal_insert_blocked_by_runtime_removal;
DROP TRIGGER workdir_removal_update_blocked_by_runtime_removal;
DROP TABLE runtime_removal_operations;
ALTER TABLE worker_removal_operations
ADD COLUMN run_generation INTEGER NOT NULL DEFAULT 0;
DROP INDEX workspace_signing_identity_audit_workspace_idx;
DROP TABLE workspace_signing_identity_audit;
DROP TABLE workspace_signing_identity_provisioning_operations;
@@ -11097,6 +11162,10 @@ mod tests {
version: 60,
name: RUNTIME_REMOVAL_OPERATION_MIGRATION_NAME.to_string(),
},
WorkspaceSchemaMigrationStep {
version: 61,
name: REMOVE_WORKER_RUN_GENERATION_MIGRATION_NAME.to_string(),
},
]
);
@@ -11143,6 +11212,10 @@ mod tests {
60,
RUNTIME_REMOVAL_OPERATION_MIGRATION_NAME.to_string(),
),
(
61,
REMOVE_WORKER_RUN_GENERATION_MIGRATION_NAME.to_string(),
),
]
);
assert!(!table_exists(conn, "trusted_runtime_records")?);
@@ -11213,7 +11286,7 @@ mod tests {
.iter()
.map(|migration| migration.version)
.collect::<Vec<_>>(),
vec![52, 53, 54, 55, 56, 57, 58, 59, 60]
vec![52, 53, 54, 55, 56, 57, 58, 59, 60, 61]
);
SqliteWorkspaceStore::migrate_database(&path).unwrap();
let conn = Connection::open(&path).unwrap();
@@ -11221,7 +11294,7 @@ mod tests {
current_schema_version(&conn).unwrap(),
LATEST_SCHEMA_VERSION
);
assert_eq!(workspace_schema_migration_history(&conn).unwrap().len(), 11);
assert_eq!(workspace_schema_migration_history(&conn).unwrap().len(), 12);
}
#[test]