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};
7use super::{catchup::CatchupModeReply, DEFAULT_REPLICATION_TERM};
8
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct WalStreamOpen {
11 pub since_lsn: u64,
12 pub max_count: usize,
13 pub replica_id: Option<String>,
14 pub installed_snapshot_token: Option<String>,
15 pub term: u64,
16 pub await_data: bool,
17 pub await_timeout_ms: u64,
18}
19
20impl WalStreamOpen {
21 pub fn encode_json(&self) -> Vec<u8> {
22 self.encode_json_with_endpoints(None)
23 }
24
25 pub fn encode_json_with_endpoints(
26 &self,
27 endpoints: Option<&crate::routing::RoutingEndpoints>,
28 ) -> Vec<u8> {
29 let mut obj = serde_json::Map::new();
30 obj.insert(
31 "since_lsn".to_string(),
32 JsonValue::Number(self.since_lsn.into()),
33 );
34 obj.insert(
35 "max_count".to_string(),
36 JsonValue::Number((self.max_count as u64).into()),
37 );
38 if let Some(replica_id) = &self.replica_id {
39 obj.insert(
40 "replica_id".to_string(),
41 JsonValue::String(replica_id.clone()),
42 );
43 }
44 if let Some(token) = &self.installed_snapshot_token {
45 obj.insert(
46 "installed_snapshot_token".to_string(),
47 JsonValue::String(token.clone()),
48 );
49 }
50 obj.insert("term".to_string(), JsonValue::Number(self.term.into()));
51 obj.insert("await_data".to_string(), JsonValue::Bool(self.await_data));
52 obj.insert(
53 "await_timeout_ms".to_string(),
54 JsonValue::Number(self.await_timeout_ms.into()),
55 );
56 if let Some(endpoints) = endpoints {
57 obj.insert("endpoints".to_string(), serde_json::json!(endpoints));
58 }
59 serde_json::to_vec(&JsonValue::Object(obj)).unwrap_or_default()
60 }
61
62 pub fn decode_json(bytes: &[u8]) -> Result<Self> {
63 Self::decode_json_with_endpoints(bytes).map(|(open, _)| open)
64 }
65
66 pub fn decode_json_with_endpoints(
67 bytes: &[u8],
68 ) -> Result<(Self, Option<crate::routing::RoutingEndpoints>)> {
69 let obj = object_from_slice(bytes)?;
70 let endpoints = obj.get("endpoints").map(crate::routing::routing_endpoints);
71 let max_count = get_opt_u64(&obj, "max_count").unwrap_or(1000);
72 let open = Self {
73 since_lsn: get_opt_u64(&obj, "since_lsn").unwrap_or(0),
74 max_count: usize::try_from(max_count)
75 .map_err(|_| ReplicationPayloadError::InvalidField("max_count"))?,
76 replica_id: get_opt_string(&obj, "replica_id"),
77 installed_snapshot_token: get_opt_string(&obj, "installed_snapshot_token"),
78 term: get_opt_u64(&obj, "term").unwrap_or(DEFAULT_REPLICATION_TERM),
79 await_data: get_bool_default(&obj, "await_data", false),
80 await_timeout_ms: get_opt_u64(&obj, "await_timeout_ms").unwrap_or(30_000),
81 };
82 Ok((open, endpoints))
83 }
84}
85
86#[derive(Debug, Clone, PartialEq, Eq)]
87pub struct WalStreamRecord {
88 pub lsn: u64,
89 pub data: Vec<u8>,
90}
91
92impl WalStreamRecord {
93 fn to_json(&self) -> JsonValue {
94 let mut obj = serde_json::Map::new();
95 obj.insert("lsn".to_string(), JsonValue::Number(self.lsn.into()));
96 obj.insert(
97 "data".to_string(),
98 JsonValue::String(hex_encode(&self.data)),
99 );
100 JsonValue::Object(obj)
101 }
102
103 fn from_json(value: &JsonValue) -> Result<Self> {
104 let obj = value
105 .as_object()
106 .ok_or(ReplicationPayloadError::InvalidField("records"))?;
107 let data_hex = get_string(obj, "data")?;
108 Ok(Self {
109 lsn: get_u64(obj, "lsn")?,
110 data: hex_decode("data", &data_hex)?,
111 })
112 }
113}
114
115#[derive(Debug, Clone, PartialEq, Eq)]
116pub struct WalStreamChunk {
117 pub records: Vec<WalStreamRecord>,
118 pub current_lsn: u64,
119 pub oldest_available_lsn: Option<u64>,
120 pub partial_resync: bool,
121 pub partial_resync_count: u64,
122 pub needs_rebootstrap: bool,
123 pub invalidation_reason: Option<String>,
124 pub catchup: Option<CatchupModeReply>,
125}
126
127impl WalStreamChunk {
128 pub fn encode_json(&self) -> Vec<u8> {
129 self.encode_json_with_endpoints(None)
130 }
131
132 pub fn encode_json_with_endpoints(
133 &self,
134 endpoints: Option<&crate::routing::RoutingEndpoints>,
135 ) -> Vec<u8> {
136 let mut obj = serde_json::Map::new();
137 obj.insert(
138 "records".to_string(),
139 JsonValue::Array(self.records.iter().map(WalStreamRecord::to_json).collect()),
140 );
141 obj.insert(
142 "current_lsn".to_string(),
143 JsonValue::Number(self.current_lsn.into()),
144 );
145 if let Some(lsn) = self.oldest_available_lsn {
146 obj.insert(
147 "oldest_available_lsn".to_string(),
148 JsonValue::Number(lsn.into()),
149 );
150 }
151 obj.insert(
152 "partial_resync".to_string(),
153 JsonValue::Bool(self.partial_resync),
154 );
155 obj.insert(
156 "partial_resync_count".to_string(),
157 JsonValue::Number(self.partial_resync_count.into()),
158 );
159 obj.insert(
160 "needs_rebootstrap".to_string(),
161 JsonValue::Bool(self.needs_rebootstrap),
162 );
163 if let Some(reason) = &self.invalidation_reason {
164 obj.insert(
165 "invalidation_reason".to_string(),
166 JsonValue::String(reason.clone()),
167 );
168 }
169 if let Some(catchup) = &self.catchup {
170 let catchup_obj = object_from_slice(&catchup.encode_json())
171 .expect("CatchupModeReply emits a JSON object");
172 for (key, value) in catchup_obj {
173 obj.insert(key, value);
174 }
175 }
176 if let Some(endpoints) = endpoints {
177 obj.insert(
178 "primary_endpoints".to_string(),
179 serde_json::json!(endpoints),
180 );
181 }
182 serde_json::to_vec(&JsonValue::Object(obj)).unwrap_or_default()
183 }
184
185 pub fn decode_json(bytes: &[u8]) -> Result<Self> {
186 Self::decode_json_with_endpoints(bytes).map(|(chunk, _)| chunk)
187 }
188
189 pub fn decode_json_with_endpoints(
190 bytes: &[u8],
191 ) -> Result<(Self, Option<crate::routing::RoutingEndpoints>)> {
192 let obj = object_from_slice(bytes)?;
193 let endpoints = obj
194 .get("primary_endpoints")
195 .map(crate::routing::routing_endpoints);
196 let records = match obj.get("records") {
197 Some(JsonValue::Array(values)) => values
198 .iter()
199 .map(WalStreamRecord::from_json)
200 .collect::<Result<Vec<_>>>()?,
201 Some(_) => return Err(ReplicationPayloadError::InvalidField("records")),
202 None => Vec::new(),
203 };
204 let chunk = Self {
205 records,
206 current_lsn: get_u64(&obj, "current_lsn")?,
207 oldest_available_lsn: get_opt_u64(&obj, "oldest_available_lsn"),
208 partial_resync: get_bool_default(&obj, "partial_resync", false),
209 partial_resync_count: get_opt_u64(&obj, "partial_resync_count").unwrap_or(0),
210 needs_rebootstrap: get_bool_default(&obj, "needs_rebootstrap", false),
211 invalidation_reason: get_opt_string(&obj, "invalidation_reason"),
212 catchup: CatchupModeReply::from_wal_rebootstrap_object(&obj)?,
213 };
214 Ok((chunk, endpoints))
215 }
216}
217
218#[derive(Debug, Clone, PartialEq, Eq)]
219pub struct WalStreamAck {
220 pub replica_id: String,
221 pub applied_lsn: u64,
222 pub durable_lsn: u64,
223 pub apply_errors_total: u64,
224 pub divergence_total: u64,
225}
226
227impl WalStreamAck {
228 pub fn encode_json(&self) -> Vec<u8> {
229 let mut obj = serde_json::Map::new();
230 obj.insert(
231 "replica_id".to_string(),
232 JsonValue::String(self.replica_id.clone()),
233 );
234 obj.insert(
235 "applied_lsn".to_string(),
236 JsonValue::Number(self.applied_lsn.into()),
237 );
238 obj.insert(
239 "durable_lsn".to_string(),
240 JsonValue::Number(self.durable_lsn.into()),
241 );
242 obj.insert(
243 "apply_errors_total".to_string(),
244 JsonValue::Number(self.apply_errors_total.into()),
245 );
246 obj.insert(
247 "divergence_total".to_string(),
248 JsonValue::Number(self.divergence_total.into()),
249 );
250 serde_json::to_vec(&JsonValue::Object(obj)).unwrap_or_default()
251 }
252
253 pub fn decode_json(bytes: &[u8]) -> Result<Self> {
254 let obj = object_from_slice(bytes)?;
255 let applied_lsn = get_u64(&obj, "applied_lsn")?;
256 Ok(Self {
257 replica_id: get_string(&obj, "replica_id")?,
258 applied_lsn,
259 durable_lsn: get_opt_u64(&obj, "durable_lsn").unwrap_or(applied_lsn),
260 apply_errors_total: get_opt_u64(&obj, "apply_errors_total").unwrap_or(0),
261 divergence_total: get_opt_u64(&obj, "divergence_total").unwrap_or(0),
262 })
263 }
264}
265
266#[derive(Debug, Clone, PartialEq, Eq)]
267pub struct WalStreamAckReply {
268 pub ok: bool,
269 pub replica_id: String,
270 pub applied_lsn: u64,
271 pub durable_lsn: u64,
272 pub apply_errors_total: u64,
273 pub divergence_total: u64,
274}
275
276impl WalStreamAckReply {
277 pub fn from_ack(ack: &WalStreamAck) -> Self {
278 Self {
279 ok: true,
280 replica_id: ack.replica_id.clone(),
281 applied_lsn: ack.applied_lsn,
282 durable_lsn: ack.durable_lsn,
283 apply_errors_total: ack.apply_errors_total,
284 divergence_total: ack.divergence_total,
285 }
286 }
287
288 pub fn encode_json(&self) -> Vec<u8> {
289 let mut obj = serde_json::Map::new();
290 obj.insert("ok".to_string(), JsonValue::Bool(self.ok));
291 obj.insert(
292 "replica_id".to_string(),
293 JsonValue::String(self.replica_id.clone()),
294 );
295 obj.insert(
296 "applied_lsn".to_string(),
297 JsonValue::Number(self.applied_lsn.into()),
298 );
299 obj.insert(
300 "durable_lsn".to_string(),
301 JsonValue::Number(self.durable_lsn.into()),
302 );
303 obj.insert(
304 "apply_errors_total".to_string(),
305 JsonValue::Number(self.apply_errors_total.into()),
306 );
307 obj.insert(
308 "divergence_total".to_string(),
309 JsonValue::Number(self.divergence_total.into()),
310 );
311 serde_json::to_vec(&JsonValue::Object(obj)).unwrap_or_default()
312 }
313
314 pub fn decode_json(bytes: &[u8]) -> Result<Self> {
315 let obj = object_from_slice(bytes)?;
316 Ok(Self {
317 ok: get_bool_default(&obj, "ok", false),
318 replica_id: get_string(&obj, "replica_id")?,
319 applied_lsn: get_u64(&obj, "applied_lsn")?,
320 durable_lsn: get_u64(&obj, "durable_lsn")?,
321 apply_errors_total: get_opt_u64(&obj, "apply_errors_total").unwrap_or(0),
322 divergence_total: get_opt_u64(&obj, "divergence_total").unwrap_or(0),
323 })
324 }
325}
326
327#[cfg(test)]
328mod tests {
329 use super::*;
330 use crate::replication::CatchupMode;
331
332 #[test]
333 fn wal_stream_open_round_trips() {
334 let open = WalStreamOpen {
335 since_lsn: 10,
336 max_count: 128,
337 replica_id: Some("replica-a".to_string()),
338 installed_snapshot_token: Some("snapshot-proof".to_string()),
339 term: 7,
340 await_data: true,
341 await_timeout_ms: 5000,
342 };
343 assert_eq!(
344 WalStreamOpen::decode_json(&open.encode_json()).unwrap(),
345 open
346 );
347 let endpoints = crate::routing::RoutingEndpoints::from([(
348 "https".to_string(),
349 "https://replica:8443".to_string(),
350 )]);
351 let advertised = open.encode_json_with_endpoints(Some(&endpoints));
352 assert_eq!(
353 WalStreamOpen::decode_json(&advertised).expect("legacy fields"),
354 open
355 );
356 assert_eq!(
357 WalStreamOpen::decode_json_with_endpoints(&advertised).expect("endpoint metadata"),
358 (open.clone(), Some(endpoints))
359 );
360 assert_eq!(
361 WalStreamOpen::decode_json_with_endpoints(&open.encode_json())
362 .expect("old peer")
363 .1,
364 None
365 );
366 }
367
368 #[test]
369 fn wal_stream_open_defaults_missing_term_to_legacy_term() {
370 let open = WalStreamOpen::decode_json(
371 br#"{"since_lsn":10,"max_count":128,"replica_id":"replica-a"}"#,
372 )
373 .unwrap();
374 assert_eq!(open.term, DEFAULT_REPLICATION_TERM);
375 assert_eq!(open.installed_snapshot_token, None);
376 }
377
378 #[test]
379 fn wal_stream_chunk_round_trips_records_and_rebootstrap_hint() {
380 let chunk = WalStreamChunk {
381 records: vec![WalStreamRecord {
382 lsn: 11,
383 data: b"record".to_vec(),
384 }],
385 current_lsn: 12,
386 oldest_available_lsn: Some(9),
387 partial_resync: true,
388 partial_resync_count: 3,
389 needs_rebootstrap: true,
390 invalidation_reason: Some("retention".to_string()),
391 catchup: Some(CatchupModeReply {
392 mode: CatchupMode::BaseBackupThenWal,
393 available_from_lsn: Some(9),
394 replica_lsn: None,
395 reason: Some("retention".to_string()),
396 }),
397 };
398 assert_eq!(
399 WalStreamChunk::decode_json(&chunk.encode_json()).unwrap(),
400 chunk
401 );
402 let endpoints = crate::routing::RoutingEndpoints::from([(
403 "grpc".to_string(),
404 "grpc://primary:55055".to_string(),
405 )]);
406 let advertised = chunk.encode_json_with_endpoints(Some(&endpoints));
407 assert_eq!(
408 WalStreamChunk::decode_json(&advertised).expect("legacy fields"),
409 chunk
410 );
411 assert_eq!(
412 WalStreamChunk::decode_json_with_endpoints(&advertised).expect("endpoint metadata"),
413 (chunk.clone(), Some(endpoints))
414 );
415 assert_eq!(
416 WalStreamChunk::decode_json_with_endpoints(&chunk.encode_json())
417 .expect("old peer")
418 .1,
419 None
420 );
421 }
422
423 #[test]
424 fn wal_ack_defaults_durable_to_applied() {
425 let ack = WalStreamAck::decode_json(br#"{"replica_id":"r","applied_lsn":7}"#).unwrap();
426 assert_eq!(ack.durable_lsn, 7);
427 assert_eq!(ack.apply_errors_total, 0);
428 }
429
430 #[test]
431 fn wal_ack_reply_round_trips() {
432 let ack = WalStreamAck {
433 replica_id: "replica-a".to_string(),
434 applied_lsn: 11,
435 durable_lsn: 10,
436 apply_errors_total: 2,
437 divergence_total: 1,
438 };
439 let reply = WalStreamAckReply::from_ack(&ack);
440
441 assert_eq!(
442 WalStreamAckReply::decode_json(&reply.encode_json()).unwrap(),
443 reply
444 );
445 }
446}