1use crate::{
7 db::{
8 DbSession, MutationJobError, MutationJobId, MutationJobState, ProgressJobInventory,
9 executor::budget::{ExecutionBudgetExceeded, HardExecutionContext},
10 integrity::with_mutation_progress_store,
11 },
12 traits::CanisterKind,
13};
14use icydb_diagnostic_code::{
15 DiagnosticExecutionBudgetResource, DiagnosticExecutionBudgetScope, DiagnosticExecutionLane,
16};
17
18#[cfg(feature = "sql")]
19use crate::db::{
20 MutationJobAdvanceReceipt, MutationJobAdvanceRequest, MutationJobPhase,
21 integrity::InsertMutationJobResult,
22 mutation_job::{CanonicalMutationIntent, MutationJobRecord},
23 session::sql::validate_current_initial_mutation_job_continuation,
24};
25
26#[cfg(feature = "sql")]
27const MUTATION_JOB_START_SHAPE: u64 = 0x6d75_7461_7465_0100;
28const MUTATION_JOB_LOAD_SHAPE: u64 = 0x6d75_7461_7465_0101;
29const MUTATION_JOB_ACKNOWLEDGE_SHAPE: u64 = 0x6d75_7461_7465_0102;
30#[cfg(feature = "sql")]
31const MUTATION_JOB_ADVANCE_SHAPE: u64 = 0x6d75_7461_7465_0103;
32#[cfg(feature = "sql")]
33const MUTATION_JOB_CANCEL_UNADVANCED_SHAPE: u64 = 0x6d75_7461_7465_0104;
34const PROGRESS_JOB_INVENTORY_SHAPE: u64 = 0x6d75_7461_7465_0105;
35
36impl<C: CanisterKind> DbSession<C> {
37 #[cfg(feature = "sql")]
45 pub fn start_trusted_sql_mutation_job(
46 &self,
47 job_id: MutationJobId,
48 sql: &str,
49 ) -> Result<MutationJobState, MutationJobError> {
50 job_id.validate()?;
51 self.charge_mutation_job_operation(
52 DiagnosticExecutionLane::Mutation,
53 MUTATION_JOB_START_SHAPE,
54 )?;
55 let prepared = self.prepare_mutation_job_start(job_id, sql)?;
56 let submitted_intent = CanonicalMutationIntent::decode(&prepared.canonical_intent)?;
57 let submitted = MutationJobRecord::new(
58 job_id,
59 prepared.canonical_intent,
60 prepared.engine_continuation,
61 )?;
62 with_mutation_progress_store::<C, _>(|store| match store.insert_mutation(&submitted)? {
63 InsertMutationJobResult::Inserted => Ok(submitted.state().clone()),
64 InsertMutationJobResult::Occupied(retained) => {
65 resolve_occupied_mutation_job_start(&retained, &submitted_intent)
66 }
67 })
68 }
69
70 pub fn mutation_job_state(
72 &self,
73 job_id: MutationJobId,
74 ) -> Result<MutationJobState, MutationJobError> {
75 self.charge_mutation_job_operation(
76 DiagnosticExecutionLane::TrustedRead,
77 MUTATION_JOB_LOAD_SHAPE,
78 )?;
79 with_mutation_progress_store::<C, _>(|store| store.load_mutation(job_id))
80 .map(|record| record.state().clone())
81 }
82
83 #[cfg(feature = "sql")]
88 pub fn advance_trusted_mutation_job(
89 &self,
90 request: &MutationJobAdvanceRequest,
91 ) -> Result<MutationJobAdvanceReceipt, MutationJobError> {
92 let retained =
93 with_mutation_progress_store::<C, _>(|store| store.load_mutation(request.job_id))?;
94 if let Some(receipt) = retained.exact_replay(request)? {
95 return Ok(receipt.clone());
96 }
97 self.charge_mutation_job_operation(
98 DiagnosticExecutionLane::Mutation,
99 MUTATION_JOB_ADVANCE_SHAPE,
100 )?;
101 retained.ensure_can_advance(request)?;
102 match retained.state().phase {
103 MutationJobPhase::Forward => self.advance_mutation_job_forward(&retained, request),
104 MutationJobPhase::Verify => self.advance_mutation_job_verify(&retained, request),
105 }
106 }
107
108 pub fn acknowledge_mutation_job(
113 &self,
114 job_id: MutationJobId,
115 expected_terminal_sequence: u64,
116 ) -> Result<(), MutationJobError> {
117 self.charge_mutation_job_operation(
118 DiagnosticExecutionLane::Mutation,
119 MUTATION_JOB_ACKNOWLEDGE_SHAPE,
120 )?;
121 with_mutation_progress_store::<C, _>(|store| {
122 store.acknowledge_mutation(job_id, expected_terminal_sequence)
123 })
124 }
125
126 #[cfg(feature = "sql")]
132 pub fn cancel_unadvanced_mutation_job(
133 &self,
134 job_id: MutationJobId,
135 expected_sequence: u64,
136 ) -> Result<(), MutationJobError> {
137 job_id.validate()?;
138 self.charge_mutation_job_operation(
139 DiagnosticExecutionLane::Mutation,
140 MUTATION_JOB_CANCEL_UNADVANCED_SHAPE,
141 )?;
142 with_mutation_progress_store::<C, _>(|store| {
143 store.cancel_unadvanced_mutation(
144 job_id,
145 expected_sequence,
146 validate_current_initial_mutation_job_continuation,
147 )
148 })
149 }
150
151 pub fn progress_job_inventory(&self) -> Result<ProgressJobInventory, MutationJobError> {
156 self.charge_mutation_job_operation(
157 DiagnosticExecutionLane::TrustedRead,
158 PROGRESS_JOB_INVENTORY_SHAPE,
159 )?;
160 with_mutation_progress_store::<C, _>(|store| store.inventory())
161 }
162
163 fn charge_mutation_job_operation(
164 &self,
165 lane: DiagnosticExecutionLane,
166 shape: u64,
167 ) -> Result<(), MutationJobError> {
168 self.db
169 .request_execution_scope()
170 .charge(
171 HardExecutionContext::new(DiagnosticExecutionBudgetScope::Execution, lane, shape),
172 DiagnosticExecutionBudgetResource::QueryExecutions,
173 1,
174 )
175 .map_err(mutation_job_execution_budget_error)
176 }
177}
178
179#[cfg(feature = "sql")]
180fn resolve_occupied_mutation_job_start(
181 retained: &MutationJobRecord,
182 submitted_intent: &CanonicalMutationIntent,
183) -> Result<MutationJobState, MutationJobError> {
184 let retained_intent = CanonicalMutationIntent::decode(retained.canonical_intent())?;
185 if retained_intent.same_start_request(submitted_intent) {
186 return Ok(retained.state().clone());
187 }
188 if !retained_intent.same_authority(submitted_intent) {
189 return Err(MutationJobError::AuthorityMismatch);
190 }
191 Err(MutationJobError::IdentityConflict)
192}
193
194const fn mutation_job_execution_budget_error(error: ExecutionBudgetExceeded) -> MutationJobError {
195 MutationJobError::ExecutionBudgetExceeded {
196 resource: error.resource().raw(),
197 limit: error.limit(),
198 observed: error.observed(),
199 scope: error.scope().raw(),
200 lane: error.lane().raw(),
201 normalized_shape_fingerprint_prefix: error.normalized_shape_fingerprint_prefix(),
202 }
203}
204
205#[cfg(test)]
206mod tests {
207 use super::*;
208 use crate::{
209 db::{
210 MutationJobAdvanceRequest, MutationJobIdempotencyKey, MutationJobPhase,
211 MutationJobRestartReason, MutationJobStatus, RequestExecutionRoot, StoreRegistry,
212 mutation_job::{MutationJobRecord, MutationJobTransition},
213 },
214 traits::Path,
215 };
216 #[cfg(feature = "sql")]
217 use crate::{
218 db::{
219 data::{AcceptedFixedUpdatePatch, FieldSlot},
220 executor::budget::{HardExecutionBudget, HardExecutionFailureHeadroom},
221 query::plan::expr::{BinaryOp, Expr, FieldId},
222 },
223 types::Timestamp,
224 value::Value,
225 };
226
227 struct TestCanister;
228
229 impl Path for TestCanister {
230 const PATH: &'static str = "db::session::mutation_job::tests::Canister";
231 }
232
233 impl CanisterKind for TestCanister {
234 const COMMIT_MEMORY_ID: u8 = 244;
235 const COMMIT_STABLE_KEY: &'static str = "icydb.test.mutation_job.commit.v1";
236 const STARTUP_MEMORY_ID: u8 = 246;
237 const STARTUP_STABLE_KEY: &'static str = "icydb.test.mutation_job.startup.control.v1";
238 const INTEGRITY_PROGRESS_MEMORY_ID: u8 = 245;
239 const INTEGRITY_PROGRESS_STABLE_KEY: &'static str = "icydb.test.mutation_job.progress.v1";
240 }
241
242 thread_local! {
243 static STORE_REGISTRY: StoreRegistry = StoreRegistry::new();
244 }
245
246 fn job_id() -> MutationJobId {
247 MutationJobId::try_from_bytes([19; 32]).expect("nonzero mutation job id should admit")
248 }
249
250 fn session() -> DbSession<TestCanister> {
251 let root = RequestExecutionRoot::__new_runtime_root();
252 DbSession::new(&STORE_REGISTRY, &root)
253 }
254
255 #[cfg(feature = "sql")]
256 fn exhausted_session() -> DbSession<TestCanister> {
257 let budget =
258 HardExecutionBudget::uniform_for_tests(0, HardExecutionFailureHeadroom::new(500, 256));
259 let root = RequestExecutionRoot::new_for_tests(budget);
260 DbSession::new(&STORE_REGISTRY, &root)
261 }
262
263 #[test]
264 fn session_load_and_terminal_acknowledgement_preserve_the_store_contract() {
265 let initial = MutationJobRecord::new(job_id(), vec![1, 2], vec![3])
266 .expect("bounded initial record should admit");
267 with_mutation_progress_store::<TestCanister, _>(|store| {
268 store.insert_mutation(&initial).map(|_| ())
269 })
270 .expect("initial record should insert");
271
272 let session = session();
273 assert_eq!(
274 session.mutation_job_state(job_id()),
275 Ok(initial.state().clone())
276 );
277 assert_eq!(
278 session.acknowledge_mutation_job(job_id(), 0),
279 Err(MutationJobError::Active),
280 );
281
282 let request = MutationJobAdvanceRequest::new(
283 job_id(),
284 0,
285 MutationJobIdempotencyKey::new("authority-drift")
286 .expect("bounded idempotency key should admit"),
287 );
288 let (terminal, _) = initial
289 .apply_transition(
290 &request,
291 MutationJobTransition::new(
292 MutationJobStatus::RestartRequired(
293 MutationJobRestartReason::ManagedTimestampRegression,
294 ),
295 MutationJobPhase::Forward,
296 Vec::new(),
297 0,
298 0,
299 0,
300 ),
301 )
302 .expect("terminal restart receipt should admit");
303 with_mutation_progress_store::<TestCanister, _>(|store| store.replace_mutation(&terminal))
304 .expect("terminal record should replace active state");
305
306 assert_eq!(
307 session.acknowledge_mutation_job(job_id(), 0),
308 Err(MutationJobError::StaleSequence {
309 expected: 0,
310 actual: 1,
311 }),
312 );
313 assert_eq!(session.acknowledge_mutation_job(job_id(), 1), Ok(()));
314 assert_eq!(session.acknowledge_mutation_job(job_id(), 1), Ok(()));
315 assert_eq!(
316 session.mutation_job_state(job_id()),
317 Err(MutationJobError::NotFound),
318 );
319 }
320
321 #[cfg(feature = "sql")]
322 fn canonical_intent(
323 authority: u8,
324 scope_value: u64,
325 timestamp: i64,
326 ) -> CanonicalMutationIntent {
327 let scope = Expr::Binary {
328 op: BinaryOp::Eq,
329 left: Box::new(Expr::Field(FieldId::new("collection_id"))),
330 right: Box::new(Expr::Literal(Value::Nat64(scope_value))),
331 };
332 let patch = AcceptedFixedUpdatePatch::from_canonical_fields(vec![(
333 FieldSlot::from_validated_index(1),
334 vec![3, 4, 5],
335 )])
336 .expect("fixed patch should admit");
337 CanonicalMutationIntent::new(
338 [authority; 16],
339 [authority; 32],
340 "journaled".to_string(),
341 "schema::Token".to_string(),
342 7,
343 11,
344 1,
345 [authority; 16],
346 &scope,
347 &patch,
348 Timestamp::from_millis(timestamp),
349 17,
350 )
351 .expect("canonical intent should admit")
352 }
353
354 #[cfg(feature = "sql")]
355 #[test]
356 fn occupied_start_distinguishes_replay_authority_drift_and_identity_conflict() {
357 let retained_intent = canonical_intent(1, 7, 100);
358 let retained = MutationJobRecord::new(
359 job_id(),
360 retained_intent.encode().expect("intent should encode"),
361 vec![9],
362 )
363 .expect("record should admit");
364
365 assert_eq!(
366 resolve_occupied_mutation_job_start(&retained, &canonical_intent(1, 7, 200)),
367 Ok(retained.state().clone()),
368 );
369 assert_eq!(
370 resolve_occupied_mutation_job_start(&retained, &canonical_intent(2, 7, 200)),
371 Err(MutationJobError::AuthorityMismatch),
372 );
373 assert_eq!(
374 resolve_occupied_mutation_job_start(&retained, &canonical_intent(1, 8, 200)),
375 Err(MutationJobError::IdentityConflict),
376 );
377 }
378
379 #[cfg(feature = "sql")]
380 #[test]
381 fn aggregate_budget_exhaustion_does_not_advance_durable_state() {
382 let initial = MutationJobRecord::new(job_id(), vec![1, 2], vec![3])
383 .expect("bounded initial record should admit");
384 with_mutation_progress_store::<TestCanister, _>(|store| {
385 store.insert_mutation(&initial).map(|_| ())
386 })
387 .expect("initial record should insert");
388 let request = MutationJobAdvanceRequest::new(
389 job_id(),
390 0,
391 MutationJobIdempotencyKey::new("budget-exhausted")
392 .expect("bounded idempotency key should admit"),
393 );
394
395 assert!(matches!(
396 exhausted_session().advance_trusted_mutation_job(&request),
397 Err(MutationJobError::ExecutionBudgetExceeded {
398 limit: 0,
399 observed: 1,
400 ..
401 })
402 ));
403 assert_eq!(
404 session().mutation_job_state(job_id()),
405 Ok(initial.state().clone()),
406 );
407 }
408
409 #[cfg(feature = "sql")]
410 #[test]
411 fn exact_replay_precedes_advance_budget_accounting() {
412 let initial = MutationJobRecord::new(job_id(), vec![1, 2], vec![3])
413 .expect("bounded initial record should admit");
414 let request = MutationJobAdvanceRequest::new(
415 job_id(),
416 0,
417 MutationJobIdempotencyKey::new("lost-response")
418 .expect("bounded idempotency key should admit"),
419 );
420 let (advanced, receipt) = initial
421 .apply_transition(
422 &request,
423 MutationJobTransition::new(
424 MutationJobStatus::Active,
425 MutationJobPhase::Forward,
426 vec![4],
427 1,
428 1,
429 0,
430 ),
431 )
432 .expect("bounded successor should admit");
433 with_mutation_progress_store::<TestCanister, _>(|store| {
434 store.insert_mutation(&advanced).map(|_| ())
435 })
436 .expect("advanced record should insert");
437
438 assert_eq!(
439 exhausted_session().advance_trusted_mutation_job(&request),
440 Ok(receipt),
441 );
442 assert_eq!(
443 session().mutation_job_state(job_id()),
444 Ok(advanced.state().clone()),
445 );
446 }
447}