Skip to main content

reddb_wire/replication/
basebackup.rs

1use serde_json::Value as JsonValue;
2
3use super::util::{
4    get_bool_default, get_opt_string, get_opt_u64, get_string, get_u64, hex_decode, hex_encode,
5    object_from_slice, ReplicationPayloadError, Result,
6};
7
8pub const BASEBACKUP_MANIFEST_HEX_FIELD: &str = "basebackup_manifest_hex";
9pub const BASEBACKUP_CHUNK_ORDINAL_FIELD: &str = "basebackup_chunk_ordinal";
10pub const BASEBACKUP_CHUNK_HEX_FIELD: &str = "basebackup_chunk_hex";
11pub const BASEBACKUP_CHUNK_PAIR_FIELD: &str = "basebackup_chunk_ordinal/basebackup_chunk_hex";
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct BaseBackupRequest {
15    pub replica_id: Option<String>,
16    pub max_bytes: Option<u64>,
17    pub snapshot_offset: u64,
18    pub snapshot_token: Option<String>,
19}
20
21impl BaseBackupRequest {
22    pub fn encode_json(&self) -> Vec<u8> {
23        let mut obj = serde_json::Map::new();
24        if let Some(replica_id) = &self.replica_id {
25            obj.insert(
26                "replica_id".to_string(),
27                JsonValue::String(replica_id.clone()),
28            );
29        }
30        if let Some(max_bytes) = self.max_bytes {
31            obj.insert("max_bytes".to_string(), JsonValue::Number(max_bytes.into()));
32        }
33        obj.insert(
34            "snapshot_offset".to_string(),
35            JsonValue::Number(self.snapshot_offset.into()),
36        );
37        if let Some(token) = &self.snapshot_token {
38            obj.insert(
39                "snapshot_token".to_string(),
40                JsonValue::String(token.clone()),
41            );
42        }
43        serde_json::to_vec(&JsonValue::Object(obj)).unwrap_or_default()
44    }
45
46    pub fn decode_json(bytes: &[u8]) -> Result<Self> {
47        let obj = object_from_slice(bytes)?;
48        Ok(Self {
49            replica_id: get_opt_string(&obj, "replica_id"),
50            max_bytes: get_opt_u64(&obj, "max_bytes"),
51            snapshot_offset: get_opt_u64(&obj, "snapshot_offset").unwrap_or(0),
52            snapshot_token: get_opt_string(&obj, "snapshot_token"),
53        })
54    }
55}
56
57#[derive(Debug, Clone, PartialEq, Eq)]
58pub struct BaseBackupManifestChunk {
59    pub ordinal: u32,
60    pub snapshot_offset: u64,
61    pub bytes: u64,
62    pub checksum: u64,
63    pub relative_path: String,
64}
65
66impl BaseBackupManifestChunk {
67    fn to_json(&self) -> JsonValue {
68        let mut obj = serde_json::Map::new();
69        obj.insert(
70            "ordinal".to_string(),
71            JsonValue::Number(self.ordinal.into()),
72        );
73        obj.insert(
74            "snapshot_offset".to_string(),
75            JsonValue::Number(self.snapshot_offset.into()),
76        );
77        obj.insert("bytes".to_string(), JsonValue::Number(self.bytes.into()));
78        obj.insert(
79            "checksum".to_string(),
80            JsonValue::Number(self.checksum.into()),
81        );
82        obj.insert(
83            "relative_path".to_string(),
84            JsonValue::String(self.relative_path.clone()),
85        );
86        JsonValue::Object(obj)
87    }
88
89    fn from_json(value: &JsonValue) -> Result<Self> {
90        let obj = value
91            .as_object()
92            .ok_or(ReplicationPayloadError::InvalidField("basebackup_chunks"))?;
93        let ordinal = get_u64(obj, "ordinal")?;
94        Ok(Self {
95            ordinal: u32::try_from(ordinal)
96                .map_err(|_| ReplicationPayloadError::InvalidField("ordinal"))?,
97            snapshot_offset: get_u64(obj, "snapshot_offset")?,
98            bytes: get_u64(obj, "bytes")?,
99            checksum: get_u64(obj, "checksum")?,
100            relative_path: get_string(obj, "relative_path")?,
101        })
102    }
103}
104
105#[derive(Debug, Clone, PartialEq, Eq)]
106pub struct BaseBackupChunk {
107    pub snapshot_available: bool,
108    pub replica_id: String,
109    pub slot_restart_lsn: u64,
110    pub snapshot_lsn: Option<u64>,
111    pub snapshot_token: Option<String>,
112    pub snapshot_total_bytes: Option<u64>,
113    pub snapshot_offset: u64,
114    pub next_snapshot_offset: Option<u64>,
115    pub snapshot_complete: bool,
116    pub snapshot_path: Option<String>,
117    pub snapshot_chunk: Option<Vec<u8>>,
118    pub snapshot_hex: Option<Vec<u8>>,
119    pub metadata_binary: Option<Vec<u8>>,
120    pub metadata_json: Option<Vec<u8>>,
121    pub basebackup_available: bool,
122    pub basebackup_timeline: Option<u64>,
123    pub basebackup_start_lsn: Option<u64>,
124    pub basebackup_checkpoint_lsn: Option<u64>,
125    pub basebackup_snapshot_bytes: Option<u64>,
126    pub basebackup_snapshot_checksum: Option<u64>,
127    pub basebackup_manifest: Option<Vec<u8>>,
128    pub basebackup_chunks: Vec<BaseBackupManifestChunk>,
129    pub basebackup_chunk_ordinal: Option<u32>,
130    pub basebackup_chunk: Option<Vec<u8>>,
131}
132
133#[derive(Debug, Clone, Copy, PartialEq, Eq)]
134pub struct BaseBackupChunkPart<'a> {
135    pub ordinal: u32,
136    pub bytes: &'a [u8],
137}
138
139impl BaseBackupChunk {
140    pub fn new(replica_id: impl Into<String>, slot_restart_lsn: u64) -> Self {
141        Self {
142            snapshot_available: true,
143            replica_id: replica_id.into(),
144            slot_restart_lsn,
145            snapshot_lsn: None,
146            snapshot_token: None,
147            snapshot_total_bytes: None,
148            snapshot_offset: 0,
149            next_snapshot_offset: None,
150            snapshot_complete: false,
151            snapshot_path: None,
152            snapshot_chunk: None,
153            snapshot_hex: None,
154            metadata_binary: None,
155            metadata_json: None,
156            basebackup_available: false,
157            basebackup_timeline: None,
158            basebackup_start_lsn: None,
159            basebackup_checkpoint_lsn: None,
160            basebackup_snapshot_bytes: None,
161            basebackup_snapshot_checksum: None,
162            basebackup_manifest: None,
163            basebackup_chunks: Vec::new(),
164            basebackup_chunk_ordinal: None,
165            basebackup_chunk: None,
166        }
167    }
168
169    pub fn encode_json(&self) -> Vec<u8> {
170        let mut obj = serde_json::Map::new();
171        obj.insert(
172            "snapshot_available".to_string(),
173            JsonValue::Bool(self.snapshot_available),
174        );
175        obj.insert(
176            "replica_id".to_string(),
177            JsonValue::String(self.replica_id.clone()),
178        );
179        obj.insert(
180            "slot_restart_lsn".to_string(),
181            JsonValue::Number(self.slot_restart_lsn.into()),
182        );
183        if let Some(lsn) = self.snapshot_lsn {
184            obj.insert("snapshot_lsn".to_string(), JsonValue::Number(lsn.into()));
185        }
186        if let Some(token) = &self.snapshot_token {
187            obj.insert(
188                "snapshot_token".to_string(),
189                JsonValue::String(token.clone()),
190            );
191        }
192        if let Some(bytes) = self.snapshot_total_bytes {
193            obj.insert(
194                "snapshot_total_bytes".to_string(),
195                JsonValue::Number(bytes.into()),
196            );
197        }
198        obj.insert(
199            "snapshot_offset".to_string(),
200            JsonValue::Number(self.snapshot_offset.into()),
201        );
202        if let Some(offset) = self.next_snapshot_offset {
203            obj.insert(
204                "next_snapshot_offset".to_string(),
205                JsonValue::Number(offset.into()),
206            );
207        }
208        obj.insert(
209            "snapshot_complete".to_string(),
210            JsonValue::Bool(self.snapshot_complete),
211        );
212        if let Some(path) = &self.snapshot_path {
213            obj.insert("snapshot_path".to_string(), JsonValue::String(path.clone()));
214        }
215        if let Some(bytes) = &self.snapshot_chunk {
216            obj.insert(
217                "snapshot_chunk_hex".to_string(),
218                JsonValue::String(hex_encode(bytes)),
219            );
220        }
221        if let Some(bytes) = &self.snapshot_hex {
222            obj.insert(
223                "snapshot_hex".to_string(),
224                JsonValue::String(hex_encode(bytes)),
225            );
226        }
227        insert_opt_hex(&mut obj, "metadata_binary_hex", &self.metadata_binary);
228        insert_opt_hex(&mut obj, "metadata_json_hex", &self.metadata_json);
229
230        obj.insert(
231            "basebackup_available".to_string(),
232            JsonValue::Bool(self.basebackup_available),
233        );
234        insert_opt_u64(&mut obj, "basebackup_timeline", self.basebackup_timeline);
235        insert_opt_u64(&mut obj, "basebackup_start_lsn", self.basebackup_start_lsn);
236        insert_opt_u64(
237            &mut obj,
238            "basebackup_checkpoint_lsn",
239            self.basebackup_checkpoint_lsn,
240        );
241        insert_opt_u64(
242            &mut obj,
243            "basebackup_snapshot_bytes",
244            self.basebackup_snapshot_bytes,
245        );
246        insert_opt_u64(
247            &mut obj,
248            "basebackup_snapshot_checksum",
249            self.basebackup_snapshot_checksum,
250        );
251        if let Some(bytes) = &self.basebackup_manifest {
252            obj.insert(
253                BASEBACKUP_MANIFEST_HEX_FIELD.to_string(),
254                JsonValue::String(hex_encode(bytes)),
255            );
256        }
257        obj.insert(
258            "basebackup_chunks".to_string(),
259            JsonValue::Array(
260                self.basebackup_chunks
261                    .iter()
262                    .map(BaseBackupManifestChunk::to_json)
263                    .collect(),
264            ),
265        );
266        if let Some(ordinal) = self.basebackup_chunk_ordinal {
267            obj.insert(
268                BASEBACKUP_CHUNK_ORDINAL_FIELD.to_string(),
269                JsonValue::Number(ordinal.into()),
270            );
271        }
272        if let Some(bytes) = &self.basebackup_chunk {
273            obj.insert(
274                BASEBACKUP_CHUNK_HEX_FIELD.to_string(),
275                JsonValue::String(hex_encode(bytes)),
276            );
277        }
278        serde_json::to_vec(&JsonValue::Object(obj)).unwrap_or_default()
279    }
280
281    pub fn decode_json(bytes: &[u8]) -> Result<Self> {
282        let obj = object_from_slice(bytes)?;
283        let basebackup_chunks = match obj.get("basebackup_chunks") {
284            Some(JsonValue::Array(values)) => values
285                .iter()
286                .map(BaseBackupManifestChunk::from_json)
287                .collect::<Result<Vec<_>>>()?,
288            Some(_) => return Err(ReplicationPayloadError::InvalidField("basebackup_chunks")),
289            None => Vec::new(),
290        };
291        let basebackup_chunk_ordinal = match get_opt_u64(&obj, BASEBACKUP_CHUNK_ORDINAL_FIELD) {
292            Some(value) => Some(u32::try_from(value).map_err(|_| {
293                ReplicationPayloadError::InvalidField(BASEBACKUP_CHUNK_ORDINAL_FIELD)
294            })?),
295            None => None,
296        };
297        Ok(Self {
298            snapshot_available: get_bool_default(&obj, "snapshot_available", false),
299            replica_id: get_string(&obj, "replica_id")?,
300            slot_restart_lsn: get_u64(&obj, "slot_restart_lsn")?,
301            snapshot_lsn: get_opt_u64(&obj, "snapshot_lsn"),
302            snapshot_token: get_opt_string(&obj, "snapshot_token"),
303            snapshot_total_bytes: get_opt_u64(&obj, "snapshot_total_bytes"),
304            snapshot_offset: get_opt_u64(&obj, "snapshot_offset").unwrap_or(0),
305            next_snapshot_offset: get_opt_u64(&obj, "next_snapshot_offset"),
306            snapshot_complete: get_bool_default(&obj, "snapshot_complete", false),
307            snapshot_path: get_opt_string(&obj, "snapshot_path"),
308            snapshot_chunk: decode_opt_hex(&obj, "snapshot_chunk_hex")?,
309            snapshot_hex: decode_opt_hex(&obj, "snapshot_hex")?,
310            metadata_binary: decode_opt_hex(&obj, "metadata_binary_hex")?,
311            metadata_json: decode_opt_hex(&obj, "metadata_json_hex")?,
312            basebackup_available: get_bool_default(&obj, "basebackup_available", false),
313            basebackup_timeline: get_opt_u64(&obj, "basebackup_timeline"),
314            basebackup_start_lsn: get_opt_u64(&obj, "basebackup_start_lsn"),
315            basebackup_checkpoint_lsn: get_opt_u64(&obj, "basebackup_checkpoint_lsn"),
316            basebackup_snapshot_bytes: get_opt_u64(&obj, "basebackup_snapshot_bytes"),
317            basebackup_snapshot_checksum: get_opt_u64(&obj, "basebackup_snapshot_checksum"),
318            basebackup_manifest: decode_opt_hex(&obj, BASEBACKUP_MANIFEST_HEX_FIELD)?,
319            basebackup_chunks,
320            basebackup_chunk_ordinal,
321            basebackup_chunk: decode_opt_hex(&obj, BASEBACKUP_CHUNK_HEX_FIELD)?,
322        })
323    }
324
325    pub fn required_basebackup_manifest(&self) -> Result<Option<&[u8]>> {
326        if !self.basebackup_available {
327            return Ok(None);
328        }
329        self.basebackup_manifest
330            .as_deref()
331            .map(Some)
332            .ok_or(ReplicationPayloadError::MissingField(
333                BASEBACKUP_MANIFEST_HEX_FIELD,
334            ))
335    }
336
337    pub fn basebackup_chunk_part(&self) -> Result<Option<BaseBackupChunkPart<'_>>> {
338        match (
339            self.basebackup_chunk_ordinal,
340            self.basebackup_chunk.as_deref(),
341        ) {
342            (Some(ordinal), Some(bytes)) => Ok(Some(BaseBackupChunkPart { ordinal, bytes })),
343            (None, None) => Ok(None),
344            _ => Err(ReplicationPayloadError::InvalidField(
345                BASEBACKUP_CHUNK_PAIR_FIELD,
346            )),
347        }
348    }
349}
350
351fn insert_opt_u64(obj: &mut serde_json::Map<String, JsonValue>, field: &str, value: Option<u64>) {
352    if let Some(value) = value {
353        obj.insert(field.to_string(), JsonValue::Number(value.into()));
354    }
355}
356
357fn insert_opt_hex(
358    obj: &mut serde_json::Map<String, JsonValue>,
359    field: &str,
360    value: &Option<Vec<u8>>,
361) {
362    if let Some(bytes) = value {
363        obj.insert(field.to_string(), JsonValue::String(hex_encode(bytes)));
364    }
365}
366
367fn decode_opt_hex(
368    obj: &serde_json::Map<String, JsonValue>,
369    field: &'static str,
370) -> Result<Option<Vec<u8>>> {
371    match obj.get(field).and_then(JsonValue::as_str) {
372        Some(value) => Ok(Some(hex_decode(field, value)?)),
373        None => Ok(None),
374    }
375}
376
377#[cfg(test)]
378mod tests {
379    use super::*;
380
381    #[test]
382    fn basebackup_request_round_trips() {
383        let request = BaseBackupRequest {
384            replica_id: Some("replica-a".to_string()),
385            max_bytes: Some(64),
386            snapshot_offset: 128,
387            snapshot_token: Some("snapshot:r:1:2".to_string()),
388        };
389        assert_eq!(
390            BaseBackupRequest::decode_json(&request.encode_json()).unwrap(),
391            request
392        );
393    }
394
395    #[test]
396    fn basebackup_chunk_constructor_sets_wire_defaults() {
397        let chunk = BaseBackupChunk::new("replica-a", 7);
398        assert!(chunk.snapshot_available);
399        assert_eq!(chunk.replica_id, "replica-a");
400        assert_eq!(chunk.slot_restart_lsn, 7);
401        assert_eq!(chunk.snapshot_offset, 0);
402        assert!(!chunk.snapshot_complete);
403        assert!(!chunk.basebackup_available);
404        assert!(chunk.basebackup_chunks.is_empty());
405    }
406
407    #[test]
408    fn basebackup_chunk_round_trips_manifest_and_payload_chunk() {
409        let chunk = BaseBackupChunk {
410            snapshot_available: true,
411            replica_id: "replica-a".to_string(),
412            slot_restart_lsn: 7,
413            snapshot_lsn: Some(9),
414            snapshot_token: Some("snapshot:replica-a:9:100".to_string()),
415            snapshot_total_bytes: Some(100),
416            snapshot_offset: 0,
417            next_snapshot_offset: Some(10),
418            snapshot_complete: false,
419            snapshot_path: Some("/tmp/replica.rdb".to_string()),
420            snapshot_chunk: Some(b"snapshot".to_vec()),
421            snapshot_hex: None,
422            metadata_binary: Some(b"metadata-binary".to_vec()),
423            metadata_json: Some(b"metadata-json".to_vec()),
424            basebackup_available: true,
425            basebackup_timeline: Some(1),
426            basebackup_start_lsn: Some(0),
427            basebackup_checkpoint_lsn: Some(9),
428            basebackup_snapshot_bytes: Some(100),
429            basebackup_snapshot_checksum: Some(123),
430            basebackup_manifest: Some(b"manifest".to_vec()),
431            basebackup_chunks: vec![BaseBackupManifestChunk {
432                ordinal: 0,
433                snapshot_offset: 0,
434                bytes: 10,
435                checksum: 99,
436                relative_path: "base/part-000000.redbasepart".to_string(),
437            }],
438            basebackup_chunk_ordinal: Some(0),
439            basebackup_chunk: Some(b"basebackup".to_vec()),
440        };
441        assert_eq!(
442            BaseBackupChunk::decode_json(&chunk.encode_json()).unwrap(),
443            chunk
444        );
445    }
446
447    #[test]
448    fn basebackup_chunk_helpers_own_required_wire_field_validation() {
449        let mut chunk = BaseBackupChunk::new("replica-a", 7);
450        chunk.basebackup_available = true;
451        assert_eq!(
452            chunk.required_basebackup_manifest().unwrap_err(),
453            ReplicationPayloadError::MissingField(BASEBACKUP_MANIFEST_HEX_FIELD)
454        );
455
456        chunk.basebackup_manifest = Some(b"manifest".to_vec());
457        assert_eq!(
458            chunk.required_basebackup_manifest().unwrap(),
459            Some(&b"manifest"[..])
460        );
461
462        chunk.basebackup_chunk_ordinal = Some(2);
463        assert_eq!(
464            chunk.basebackup_chunk_part().unwrap_err(),
465            ReplicationPayloadError::InvalidField(BASEBACKUP_CHUNK_PAIR_FIELD)
466        );
467
468        chunk.basebackup_chunk = Some(b"part".to_vec());
469        assert_eq!(
470            chunk.basebackup_chunk_part().unwrap(),
471            Some(BaseBackupChunkPart {
472                ordinal: 2,
473                bytes: &b"part"[..],
474            })
475        );
476    }
477}