1#[cfg(test)]
7mod tests;
8
9use crate::{
10 db::{
11 Db,
12 codec::deserialize_persisted_payload,
13 commit::{
14 CommitMarker, CommitRowOp, begin_commit_with_migration_state,
15 clear_migration_state_bytes, finish_commit, load_migration_state_bytes,
16 },
17 },
18 error::InternalError,
19 serialize::serialize,
20 traits::CanisterKind,
21};
22use serde::{Deserialize, Serialize};
23
24const MAX_MIGRATION_STATE_BYTES: usize = 64 * 1024;
25
26#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
35pub struct MigrationCursor {
36 next_step: usize,
37}
38
39impl MigrationCursor {
40 #[must_use]
42 pub const fn start() -> Self {
43 Self { next_step: 0 }
44 }
45
46 #[must_use]
48 pub const fn next_step(self) -> usize {
49 self.next_step
50 }
51
52 const fn from_step(step_index: usize) -> Self {
53 Self {
54 next_step: step_index,
55 }
56 }
57
58 const fn advance(self) -> Self {
60 Self {
61 next_step: self.next_step.saturating_add(1),
62 }
63 }
64}
65
66#[derive(Clone, Debug, Deserialize, Serialize)]
77#[serde(deny_unknown_fields)]
78struct PersistedMigrationState {
79 migration_id: String,
80 migration_version: u64,
81 step_index: u64,
82 last_applied_row_key: Option<Vec<u8>>,
83}
84
85#[derive(Clone, Debug)]
94pub struct MigrationRowOp {
95 pub entity_path: String,
97 pub key: Vec<u8>,
99 pub before: Option<Vec<u8>>,
101 pub after: Option<Vec<u8>>,
103 pub schema_fingerprint: [u8; 16],
105}
106
107impl MigrationRowOp {
108 #[must_use]
110 pub fn new(
111 entity_path: impl Into<String>,
112 key: Vec<u8>,
113 before: Option<Vec<u8>>,
114 after: Option<Vec<u8>>,
115 schema_fingerprint: [u8; 16],
116 ) -> Self {
117 Self {
118 entity_path: entity_path.into(),
119 key,
120 before,
121 after,
122 schema_fingerprint,
123 }
124 }
125}
126
127impl TryFrom<MigrationRowOp> for CommitRowOp {
128 type Error = InternalError;
129
130 fn try_from(op: MigrationRowOp) -> Result<Self, Self::Error> {
131 Self::try_new_bytes(
132 op.entity_path,
133 op.key.as_slice(),
134 op.before,
135 op.after,
136 op.schema_fingerprint,
137 )
138 }
139}
140
141#[derive(Clone, Debug)]
150pub struct MigrationStep {
151 name: String,
152 row_ops: Vec<CommitRowOp>,
153}
154
155impl MigrationStep {
156 pub fn from_row_ops(
158 name: impl Into<String>,
159 row_ops: Vec<MigrationRowOp>,
160 ) -> Result<Self, InternalError> {
161 let commit_row_ops = row_ops
162 .into_iter()
163 .map(CommitRowOp::try_from)
164 .collect::<Result<Vec<_>, _>>()?;
165 Self::new(name, commit_row_ops)
166 }
167
168 pub(in crate::db) fn new(
170 name: impl Into<String>,
171 row_ops: Vec<CommitRowOp>,
172 ) -> Result<Self, InternalError> {
173 let name = name.into();
174 validate_non_empty_label(name.as_str(), "migration step name")?;
175
176 if row_ops.is_empty() {
177 return Err(InternalError::migration_step_row_ops_required(&name));
178 }
179
180 Ok(Self { name, row_ops })
181 }
182
183 #[must_use]
185 pub const fn name(&self) -> &str {
186 self.name.as_str()
187 }
188
189 #[must_use]
191 pub const fn row_op_count(&self) -> usize {
192 self.row_ops.len()
193 }
194}
195
196#[derive(Clone, Debug)]
206pub struct MigrationPlan {
207 id: String,
208 version: u64,
209 steps: Vec<MigrationStep>,
210}
211
212impl MigrationPlan {
213 pub fn new(
215 id: impl Into<String>,
216 version: u64,
217 steps: Vec<MigrationStep>,
218 ) -> Result<Self, InternalError> {
219 let id = id.into();
220 validate_non_empty_label(id.as_str(), "migration plan id")?;
221 if version == 0 {
222 return Err(InternalError::migration_plan_version_required(&id));
223 }
224
225 if steps.is_empty() {
226 return Err(InternalError::migration_plan_steps_required(&id));
227 }
228
229 Ok(Self { id, version, steps })
230 }
231
232 #[must_use]
234 pub const fn id(&self) -> &str {
235 self.id.as_str()
236 }
237
238 #[must_use]
240 pub const fn version(&self) -> u64 {
241 self.version
242 }
243
244 #[must_use]
246 pub const fn len(&self) -> usize {
247 self.steps.len()
248 }
249
250 #[must_use]
252 pub const fn is_empty(&self) -> bool {
253 self.steps.is_empty()
254 }
255
256 fn step_at(&self, index: usize) -> Result<&MigrationStep, InternalError> {
257 self.steps.get(index).ok_or_else(|| {
258 InternalError::migration_cursor_out_of_bounds(
259 self.id(),
260 self.version(),
261 index,
262 self.len(),
263 )
264 })
265 }
266}
267
268#[derive(Clone, Copy, Debug, Eq, PartialEq)]
275pub enum MigrationRunState {
276 Complete,
278 NeedsResume,
280}
281
282#[derive(Clone, Copy, Debug, Eq, PartialEq)]
291pub struct MigrationRunOutcome {
292 cursor: MigrationCursor,
293 applied_steps: usize,
294 applied_row_ops: usize,
295 state: MigrationRunState,
296}
297
298impl MigrationRunOutcome {
299 const fn new(
300 cursor: MigrationCursor,
301 applied_steps: usize,
302 applied_row_ops: usize,
303 state: MigrationRunState,
304 ) -> Self {
305 Self {
306 cursor,
307 applied_steps,
308 applied_row_ops,
309 state,
310 }
311 }
312
313 #[must_use]
315 pub const fn cursor(self) -> MigrationCursor {
316 self.cursor
317 }
318
319 #[must_use]
321 pub const fn applied_steps(self) -> usize {
322 self.applied_steps
323 }
324
325 #[must_use]
327 pub const fn applied_row_ops(self) -> usize {
328 self.applied_row_ops
329 }
330
331 #[must_use]
333 pub const fn state(self) -> MigrationRunState {
334 self.state
335 }
336}
337
338pub(in crate::db) fn execute_migration_plan<C: CanisterKind>(
347 db: &Db<C>,
348 plan: &MigrationPlan,
349 max_steps: usize,
350) -> Result<MigrationRunOutcome, InternalError> {
351 if max_steps == 0 {
353 return Err(InternalError::migration_execution_requires_max_steps(
354 plan.id(),
355 ));
356 }
357
358 db.ensure_recovered_state()?;
360
361 let mut next_cursor = load_durable_cursor_for_plan(plan)?;
363
364 let mut applied_steps = 0usize;
366 let mut applied_row_ops = 0usize;
367 while applied_steps < max_steps {
368 if next_cursor.next_step() >= plan.len() {
369 break;
370 }
371
372 let step_index = next_cursor.next_step();
373 let step = plan.step_at(step_index)?;
374 let next_cursor_after_step = next_cursor.advance();
375 let next_state_bytes =
376 encode_durable_cursor_state(plan, next_cursor_after_step, step.row_ops.last())?;
377 execute_migration_step(db, plan, step_index, step, next_state_bytes)?;
378
379 applied_steps = applied_steps.saturating_add(1);
380 applied_row_ops = applied_row_ops.saturating_add(step.row_op_count());
381 next_cursor = next_cursor_after_step;
382 }
383
384 let state = if next_cursor.next_step() == plan.len() {
385 clear_migration_state_bytes()?;
386 MigrationRunState::Complete
387 } else {
388 MigrationRunState::NeedsResume
389 };
390
391 Ok(MigrationRunOutcome::new(
392 next_cursor,
393 applied_steps,
394 applied_row_ops,
395 state,
396 ))
397}
398
399fn load_durable_cursor_for_plan(plan: &MigrationPlan) -> Result<MigrationCursor, InternalError> {
400 let Some(bytes) = load_migration_state_bytes()? else {
401 return Ok(MigrationCursor::start());
402 };
403 let state = decode_persisted_migration_state(&bytes)?;
404 if state.migration_id != plan.id() || state.migration_version != plan.version() {
405 return Err(InternalError::migration_in_progress_conflict(
406 plan.id(),
407 plan.version(),
408 &state.migration_id,
409 state.migration_version,
410 ));
411 }
412
413 let step_index = usize::try_from(state.step_index).map_err(|_| {
414 InternalError::migration_persisted_step_index_invalid_usize(
415 plan.id(),
416 plan.version(),
417 state.step_index,
418 )
419 })?;
420 if step_index > plan.len() {
421 return Err(InternalError::migration_persisted_step_index_out_of_bounds(
422 plan.id(),
423 plan.version(),
424 step_index,
425 plan.len(),
426 ));
427 }
428
429 if step_index == plan.len() {
430 clear_migration_state_bytes()?;
431 }
432
433 Ok(MigrationCursor::from_step(step_index))
434}
435
436fn encode_durable_cursor_state(
437 plan: &MigrationPlan,
438 cursor: MigrationCursor,
439 last_applied_row_op: Option<&CommitRowOp>,
440) -> Result<Vec<u8>, InternalError> {
441 let step_index = u64::try_from(cursor.next_step()).map_err(|_| {
442 InternalError::migration_next_step_index_u64_required(plan.id(), plan.version())
443 })?;
444 let state = PersistedMigrationState {
445 migration_id: plan.id().to_string(),
446 migration_version: plan.version(),
447 step_index,
448 last_applied_row_key: last_applied_row_op.map(|row_op| row_op.key.as_bytes().to_vec()),
449 };
450
451 encode_persisted_migration_state(&state)
452}
453
454fn decode_persisted_migration_state(
455 bytes: &[u8],
456) -> Result<PersistedMigrationState, InternalError> {
457 deserialize_persisted_payload::<PersistedMigrationState>(
458 bytes,
459 MAX_MIGRATION_STATE_BYTES,
460 "migration state",
461 )
462}
463
464fn encode_persisted_migration_state(
465 state: &PersistedMigrationState,
466) -> Result<Vec<u8>, InternalError> {
467 serialize(state).map_err(InternalError::migration_state_serialize_failed)
468}
469
470fn execute_migration_step<C: CanisterKind>(
471 db: &Db<C>,
472 plan: &MigrationPlan,
473 step_index: usize,
474 step: &MigrationStep,
475 next_state_bytes: Vec<u8>,
476) -> Result<(), InternalError> {
477 let marker = CommitMarker::new(step.row_ops.clone())
479 .map_err(|err| annotate_step_error(plan, step_index, step.name(), err))?;
480 let commit = begin_commit_with_migration_state(marker, next_state_bytes)
481 .map_err(|err| annotate_step_error(plan, step_index, step.name(), err))?;
482
483 finish_commit(commit, |_| apply_marker_row_ops(db, &step.row_ops))
485 .map_err(|err| annotate_step_error(plan, step_index, step.name(), err))?;
486
487 Ok(())
488}
489
490fn apply_marker_row_ops<C: CanisterKind>(
491 db: &Db<C>,
492 row_ops: &[CommitRowOp],
493) -> Result<(), InternalError> {
494 let mut prepared = Vec::with_capacity(row_ops.len());
496 for row_op in row_ops {
497 prepared.push(db.prepare_row_commit_op(row_op)?);
498 }
499
500 for prepared_op in prepared {
502 prepared_op.apply();
503 }
504
505 Ok(())
506}
507
508fn annotate_step_error(
509 plan: &MigrationPlan,
510 step_index: usize,
511 step_name: &str,
512 err: InternalError,
513) -> InternalError {
514 let source_message = err.message().to_string();
515
516 err.with_message(format!(
517 "migration '{}' step {} ('{}') failed: {}",
518 plan.id(),
519 step_index,
520 step_name,
521 source_message,
522 ))
523}
524
525fn validate_non_empty_label(value: &str, label: &str) -> Result<(), InternalError> {
526 if value.trim().is_empty() {
527 return Err(InternalError::migration_label_empty(label));
528 }
529
530 Ok(())
531}