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