mz_persist_client/internal/
restore.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! See documentation on [[restore_blob]].
11
12use crate::ShardId;
13use crate::internal::encoding::UntypedState;
14use crate::internal::paths::BlobKey;
15use crate::internal::state::{BatchPart, RunPart, State};
16use crate::internal::state_diff::{StateDiff, StateFieldValDiff};
17use crate::internal::state_versions::StateVersions;
18use crate::metrics::Metrics;
19use anyhow::anyhow;
20use mz_persist::location::Blob;
21use tracing::info;
22
23/// Attempt to restore all the blobs referenced by the current state in consensus.
24/// Returns a list of blobs that were not possible to restore.
25pub(crate) async fn restore_blob(
26    versions: &StateVersions,
27    blob: &dyn Blob,
28    build_version: &semver::Version,
29    shard_id: ShardId,
30    metrics: &Metrics,
31) -> anyhow::Result<Vec<BlobKey>> {
32    let diffs = versions.fetch_all_live_diffs(&shard_id).await;
33    let Some(first_live_seqno) = diffs.0.first().map(|d| d.seqno) else {
34        info!("No diffs for shard {shard_id}.");
35        return Ok(vec![]);
36    };
37
38    fn after<A>(diff: StateFieldValDiff<A>) -> Option<A> {
39        match diff {
40            StateFieldValDiff::Insert(a) => Some(a),
41            StateFieldValDiff::Update(_, a) => Some(a),
42            StateFieldValDiff::Delete(_) => None,
43        }
44    }
45
46    let mut not_restored = vec![];
47    let mut check_restored = |key: &BlobKey, result: Result<(), _>| {
48        if result.is_err() {
49            not_restored.push(key.clone());
50        }
51    };
52
53    for diff in diffs.0 {
54        let mut diff: StateDiff<u64> = StateDiff::decode(build_version, diff.data);
55        let mut part_queue = vec![];
56
57        for rollup in std::mem::take(&mut diff.rollups) {
58            // We never actually reference rollups from before the first live diff.
59            if rollup.key < first_live_seqno {
60                continue;
61            }
62            let Some(value) = after(rollup.val) else {
63                continue;
64            };
65            let key = value.key.complete(&shard_id);
66            let rollup_result = blob.restore(&key).await;
67            let rollup_restored = rollup_result.is_ok();
68            check_restored(&key, rollup_result);
69
70            // Elsewhere, we restore any state referenced in live diffs... but we also
71            // need to restore everything referenced in that first rollup.
72            // If restoring the rollup failed, let's
73            // keep going to try and recover the rest of the blobs before bailing out.
74            if rollup.key != first_live_seqno || !rollup_restored {
75                continue;
76            }
77            let rollup_bytes = blob
78                .get(&key)
79                .await?
80                .ok_or_else(|| anyhow!("fetching just-restored rollup"))?;
81            let rollup_state: State<u64> =
82                UntypedState::decode(build_version, rollup_bytes).check_ts_codec(&shard_id)?;
83            for (seqno, rollup) in &rollup_state.collections.rollups {
84                // We never actually reference rollups from before the first live diff.
85                if *seqno < first_live_seqno {
86                    continue;
87                }
88                let key = rollup.key.complete(&shard_id);
89                check_restored(&key, blob.restore(&key).await);
90            }
91            for batch in rollup_state.collections.trace.batches() {
92                part_queue.extend(batch.parts.iter().cloned());
93            }
94        }
95        for diff in diff.referenced_batches() {
96            if let Some(after) = after(diff) {
97                part_queue.extend(after.parts.iter().cloned())
98            }
99        }
100        while let Some(part) = part_queue.pop() {
101            match part {
102                RunPart::Single(BatchPart::Inline { .. }) => {}
103                RunPart::Single(BatchPart::Hollow(part)) => {
104                    let key = part.key.complete(&shard_id);
105                    check_restored(&key, blob.restore(&key).await);
106                }
107                RunPart::Many(runs) => {
108                    let key = runs.key.complete(&shard_id);
109                    check_restored(&key, blob.restore(&key).await);
110                    let runs = runs
111                        .get(shard_id, blob, metrics)
112                        .await
113                        .ok_or_else(|| anyhow!("fetching just-restored run"))?;
114                    part_queue.extend(runs.parts);
115                }
116            }
117        }
118    }
119    Ok(not_restored)
120}
mz_persist_client/internal/restore.rs

mz_persist_client/internal/
restore.rs
