1use crate::{
7 db::{
8 CompareProofAndAdvanceError, DbSession, ExhaustiveReadError, ReadSetRevisionProof,
9 ResumableJobAdvance, ResumableJobAdvanceReceipt, ResumableJobAdvanceRequest,
10 ResumableJobError, ResumableJobId, ResumableJobRecord, ResumableJobState,
11 ResumableJobStatus,
12 executor::budget::{ExecutionBudgetExceeded, HardExecutionContext},
13 integrity::with_resumable_progress_store,
14 },
15 traits::CanisterKind,
16};
17use icydb_diagnostic_code::{
18 DiagnosticExecutionBudgetResource, DiagnosticExecutionBudgetScope, DiagnosticExecutionLane,
19};
20
21const RESUMABLE_JOB_START_SHAPE: u64 = 0x7265_7375_6d65_0101;
22const RESUMABLE_JOB_LOAD_SHAPE: u64 = 0x7265_7375_6d65_0102;
23const RESUMABLE_JOB_ADVANCE_SHAPE: u64 = 0x7265_7375_6d65_0103;
24const RESUMABLE_JOB_ACKNOWLEDGE_SHAPE: u64 = 0x7265_7375_6d65_0104;
25
26impl<C: CanisterKind> DbSession<C> {
27 pub fn start_resumable_job(
30 &self,
31 job_id: ResumableJobId,
32 proof: ReadSetRevisionProof,
33 initial_application_state: Vec<u8>,
34 ) -> Result<ResumableJobState, ResumableJobError> {
35 self.charge_resumable_operation(
36 DiagnosticExecutionLane::Mutation,
37 RESUMABLE_JOB_START_SHAPE,
38 )?;
39 self.verify_durable_read_set_revision_proof(&proof)
40 .map_err(map_exhaustive_error)?;
41 let record = ResumableJobRecord::new(job_id, proof, initial_application_state)?;
42 with_resumable_progress_store::<C, _>(|store| store.insert_resumable(&record))?;
43 Ok(record.state().clone())
44 }
45
46 pub fn resumable_job_state(
48 &self,
49 job_id: ResumableJobId,
50 ) -> Result<ResumableJobState, ResumableJobError> {
51 self.charge_resumable_operation(
52 DiagnosticExecutionLane::TrustedRead,
53 RESUMABLE_JOB_LOAD_SHAPE,
54 )?;
55 with_resumable_progress_store::<C, _>(|store| store.load_resumable(job_id))
56 .map(|record| record.state().clone())
57 }
58
59 pub fn acknowledge_resumable_job(
66 &self,
67 job_id: ResumableJobId,
68 expected_sequence: u64,
69 ) -> Result<(), ResumableJobError> {
70 self.charge_resumable_operation(
71 DiagnosticExecutionLane::Mutation,
72 RESUMABLE_JOB_ACKNOWLEDGE_SHAPE,
73 )?;
74 with_resumable_progress_store::<C, _>(|store| {
75 let record = match store.load_resumable(job_id) {
76 Ok(record) => record,
77 Err(ResumableJobError::NotFound) => return Ok(()),
78 Err(error) => return Err(error),
79 };
80 if record.state().sequence != expected_sequence {
81 return Err(ResumableJobError::StaleSequence {
82 expected: expected_sequence,
83 actual: record.state().sequence,
84 });
85 }
86 let terminal = record.state().sequence > 0
87 && matches!(
88 record.state().status,
89 ResumableJobStatus::Completed | ResumableJobStatus::Invalidated
90 );
91 if !terminal {
92 return Err(ResumableJobError::NotTerminal);
93 }
94 store.remove_resumable(job_id)
95 })
96 }
97
98 pub fn compare_proof_and_advance<E>(
106 &self,
107 request: &ResumableJobAdvanceRequest,
108 operation: impl FnOnce(&ResumableJobState) -> Result<ResumableJobAdvance, E>,
109 ) -> Result<ResumableJobAdvanceReceipt, CompareProofAndAdvanceError<E>> {
110 self.charge_resumable_operation(
111 DiagnosticExecutionLane::Mutation,
112 RESUMABLE_JOB_ADVANCE_SHAPE,
113 )?;
114 request.idempotency_key.validate()?;
115 if request.job_id.to_bytes() == [0; 32] {
116 return Err(ResumableJobError::InvalidJobId.into());
117 }
118
119 let record =
120 with_resumable_progress_store::<C, _>(|store| store.load_resumable(request.job_id))?;
121 if let Some(receipt) = exact_replay(&record, request) {
122 return Ok(receipt.clone());
123 }
124 ensure_request_can_advance(&record, request)?;
125
126 if let Err(error) = self.verify_read_set_revision_proof(&record.state().proof) {
127 return Self::persist_or_report_pre_page_invalidation(record, request, error);
128 }
129
130 let advance = operation(record.state()).map_err(CompareProofAndAdvanceError::Operation)?;
131 advance.validate()?;
132
133 if let Err(error) = self.verify_read_set_revision_proof(&record.state().proof) {
134 return Self::persist_or_report_post_page_invalidation(record, request, error);
135 }
136
137 let (candidate, receipt) = record.apply_advance(request, advance)?;
138 Self::replace_resumable_job_if_current(request, &candidate)?;
139 Ok(receipt)
140 }
141
142 fn persist_or_report_pre_page_invalidation<E>(
143 record: ResumableJobRecord,
144 request: &ResumableJobAdvanceRequest,
145 error: ExhaustiveReadError,
146 ) -> Result<ResumableJobAdvanceReceipt, CompareProofAndAdvanceError<E>> {
147 Self::persist_or_report_invalidation(record, request, error)
148 }
149
150 fn persist_or_report_post_page_invalidation<E>(
151 record: ResumableJobRecord,
152 request: &ResumableJobAdvanceRequest,
153 error: ExhaustiveReadError,
154 ) -> Result<ResumableJobAdvanceReceipt, CompareProofAndAdvanceError<E>> {
155 Self::persist_or_report_invalidation(record, request, error)
156 }
157
158 fn persist_or_report_invalidation<E>(
159 record: ResumableJobRecord,
160 request: &ResumableJobAdvanceRequest,
161 error: ExhaustiveReadError,
162 ) -> Result<ResumableJobAdvanceReceipt, CompareProofAndAdvanceError<E>> {
163 match error {
164 ExhaustiveReadError::Revision(error) if error.is_source_change() => {
165 let (invalidated, receipt) = record.invalidate(request)?;
166 Self::replace_resumable_job_if_current(request, &invalidated)?;
167 Ok(receipt)
168 }
169 other => Err(map_exhaustive_error(other).into()),
170 }
171 }
172
173 fn replace_resumable_job_if_current<E>(
174 request: &ResumableJobAdvanceRequest,
175 candidate: &ResumableJobRecord,
176 ) -> Result<(), CompareProofAndAdvanceError<E>> {
177 with_resumable_progress_store::<C, _>(|store| {
178 let current = store.load_resumable(request.job_id)?;
179 if exact_replay(¤t, request).is_some() {
180 return Err(ResumableJobError::StaleSequence {
181 expected: request.expected_sequence,
182 actual: current.state().sequence,
183 });
184 }
185 ensure_request_can_advance(¤t, request)?;
186 store.replace_resumable(candidate)
187 })?;
188 Ok(())
189 }
190
191 fn charge_resumable_operation(
192 &self,
193 lane: DiagnosticExecutionLane,
194 shape: u64,
195 ) -> Result<(), ResumableJobError> {
196 self.db
197 .request_execution_scope()
198 .charge(
199 HardExecutionContext::new(DiagnosticExecutionBudgetScope::Execution, lane, shape),
200 DiagnosticExecutionBudgetResource::QueryExecutions,
201 1,
202 )
203 .map_err(resumable_execution_budget_error)
204 }
205}
206
207fn exact_replay<'a>(
208 record: &'a ResumableJobRecord,
209 request: &ResumableJobAdvanceRequest,
210) -> Option<&'a ResumableJobAdvanceReceipt> {
211 record.last_receipt().filter(|receipt| {
212 receipt.request_sequence == request.expected_sequence
213 && receipt.idempotency_key() == &request.idempotency_key
214 })
215}
216
217fn ensure_request_can_advance(
218 record: &ResumableJobRecord,
219 request: &ResumableJobAdvanceRequest,
220) -> Result<(), ResumableJobError> {
221 if record.state().status == ResumableJobStatus::Invalidated {
222 return Err(ResumableJobError::Invalidated);
223 }
224 if record.state().status == ResumableJobStatus::Completed {
225 return Err(ResumableJobError::Completed);
226 }
227 if record.state().sequence != request.expected_sequence {
228 return Err(ResumableJobError::StaleSequence {
229 expected: request.expected_sequence,
230 actual: record.state().sequence,
231 });
232 }
233 Ok(())
234}
235
236fn map_exhaustive_error(error: ExhaustiveReadError) -> ResumableJobError {
237 match error {
238 ExhaustiveReadError::Revision(error) => ResumableJobError::SourceProof(error),
239 ExhaustiveReadError::Query(_) => ResumableJobError::Internal,
240 }
241}
242
243const fn resumable_execution_budget_error(error: ExecutionBudgetExceeded) -> ResumableJobError {
244 ResumableJobError::ExecutionBudgetExceeded {
245 resource: error.resource().raw(),
246 limit: error.limit(),
247 observed: error.observed(),
248 scope: error.scope().raw(),
249 lane: error.lane().raw(),
250 normalized_shape_fingerprint_prefix: error.normalized_shape_fingerprint_prefix(),
251 }
252}