use crate::sync::version_vector::VersionVector;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum Operation {
Insert,
Update,
Delete,
CreateCollection,
DeleteCollection,
TruncateCollection,
CreateDatabase,
DeleteDatabase,
PutBlobChunk,
DeleteBlob,
ColumnarInsert,
ColumnarDelete,
ColumnarCreateCollection,
ColumnarDropCollection,
ColumnarTruncate,
CreateIndex,
DropIndex,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SyncEntry {
pub sequence: u64,
pub origin_node: String,
pub origin_sequence: u64,
pub hlc_ts: u64,
pub hlc_count: u32,
pub database: String,
pub collection: String,
pub operation: Operation,
pub document_key: String,
#[serde(with = "serde_bytes")]
pub document_data: Option<Vec<u8>>,
pub shard_id: Option<u16>,
pub version_vector: Option<VersionVector>,
pub parent_vectors: Vec<VersionVector>,
pub is_delta: bool,
#[serde(with = "serde_bytes")]
pub delta_data: Option<Vec<u8>>,
pub session_id: Option<String>,
pub device_id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ShardConfig {
pub num_shards: u16,
pub replication_factor: u16,
pub shard_key: String,
}
impl ShardConfig {
pub fn new(num_shards: u16, replication_factor: u16) -> Self {
Self {
num_shards,
replication_factor,
shard_key: "_key".to_string(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShardAssignment {
pub shard_id: u16,
pub owner: String,
pub replicas: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct NodeStats {
pub cpu_usage: f32,
pub memory_used: u64,
pub disk_used: u64,
pub document_count: u64,
pub collections_count: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SyncMessage {
AuthChallenge {
challenge: Vec<u8>,
timestamp: u64,
nonce: Vec<u8>,
},
AuthResponse { hmac: Vec<u8> },
AuthResult { success: bool, message: String },
IncrementalSyncRequest {
from_node: String,
after_sequence: u64,
max_batch_bytes: u32,
},
FullSyncRequest { from_node: String },
FullSyncStart {
total_databases: u32,
total_collections: u32,
total_documents: u64,
},
FullSyncDatabase { name: String },
FullSyncCollection {
database: String,
name: String,
shard_config: Option<ShardConfig>,
collection_type: Option<String>,
},
FullSyncDocuments {
database: String,
collection: String,
data: Vec<u8>,
compressed: bool,
doc_count: u32,
},
FullSyncComplete { final_sequence: u64 },
SyncBatch {
entries: Vec<SyncEntry>,
has_more: bool,
current_sequence: u64,
compressed: bool,
},
Heartbeat {
node_id: String,
sequence: u64,
stats: NodeStats,
},
HeartbeatAck { node_id: String },
NodeJoin {
node_id: String,
address: String,
http_address: String,
},
NodeLeave { node_id: String },
NodeDead { node_id: String },
ShardRebalance {
database: String,
collection: String,
assignments: Vec<ShardAssignment>,
},
ClientRegisterSession {
device_id: String,
api_key: String,
filter_query: Option<String>,
subscriptions: Vec<String>,
},
ClientSessionRegistered {
session_id: String,
server_vector: VersionVector,
supports_delta_sync: bool,
supports_crdt: bool,
},
ClientPullRequest {
session_id: String,
client_vector: VersionVector,
limit: Option<usize>,
},
ClientPullResponse {
changes: Vec<SyncEntry>,
server_vector: VersionVector,
has_more: bool,
conflicts: Vec<ConflictEntry>,
},
ClientPushRequest {
session_id: String,
changes: Vec<SyncEntry>,
client_vector: VersionVector,
},
ClientPushResponse {
server_vector: VersionVector,
conflicts: Vec<ConflictEntry>,
accepted: usize,
rejected: usize,
},
ClientSyncAck {
session_id: String,
applied_vector: VersionVector,
},
ClientSubscribe {
session_id: String,
collections: Vec<String>,
},
ClientUnsubscribe {
session_id: String,
collections: Vec<String>,
},
ClientNotifyChanges {
session_id: String,
has_changes: bool,
collections: Vec<String>,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConflictEntry {
pub document_key: String,
pub collection: String,
pub local_vector: VersionVector,
pub remote_vector: VersionVector,
pub local_data: Option<Vec<u8>>,
pub remote_data: Option<Vec<u8>>,
pub detected_at: u64,
}
pub fn encode_documents(batch: &[serde_json::Value]) -> Result<Vec<u8>, String> {
serde_json::to_vec(batch).map_err(|e| format!("encoding {} document(s): {e}", batch.len()))
}
pub fn decode_documents(data: &[u8]) -> Result<Vec<serde_json::Value>, String> {
serde_json::from_slice(data).map_err(|e| format!("decoding a document batch: {e}"))
}
impl SyncMessage {
pub fn encode(&self) -> Vec<u8> {
let payload = bincode::serialize(self).expect("Failed to serialize SyncMessage");
let len = payload.len() as u32;
let mut result = Vec::with_capacity(5 + payload.len());
result.push(0);
result.extend_from_slice(&len.to_be_bytes());
result.extend(payload);
result
}
pub fn decode(bytes: &[u8]) -> Result<Self, bincode::Error> {
bincode::deserialize(bytes)
}
pub const HEADER_LEN: usize = 5;
pub fn decode_frame(frame: &[u8]) -> Result<Self, bincode::Error> {
use bincode::ErrorKind;
if frame.len() < Self::HEADER_LEN {
return Err(Box::new(ErrorKind::Custom(format!(
"frame is {} bytes, shorter than the {}-byte header",
frame.len(),
Self::HEADER_LEN
))));
}
let declared = u32::from_be_bytes([frame[1], frame[2], frame[3], frame[4]]) as usize;
let body = &frame[Self::HEADER_LEN..];
if body.len() != declared {
return Err(Box::new(ErrorKind::Custom(format!(
"frame declares {declared} bytes of payload and carries {}",
body.len()
))));
}
Self::decode(body)
}
}
impl SyncEntry {
#[allow(clippy::too_many_arguments)]
pub fn new(
sequence: u64,
origin_node: String,
origin_sequence: u64,
hlc_ts: u64,
hlc_count: u32,
database: String,
collection: String,
operation: Operation,
document_key: String,
document_data: Option<Vec<u8>>,
shard_id: Option<u16>,
) -> Self {
Self {
sequence,
origin_node,
origin_sequence,
hlc_ts,
hlc_count,
database,
collection,
operation,
document_key,
document_data,
shard_id,
version_vector: None,
parent_vectors: Vec::new(),
is_delta: false,
delta_data: None,
session_id: None,
device_id: None,
}
}
pub fn with_version_vector(
origin_node: String,
database: String,
collection: String,
operation: Operation,
document_key: String,
document_data: Option<Vec<u8>>,
version_vector: VersionVector,
) -> Self {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
Self {
sequence: 0,
origin_node: origin_node.clone(),
origin_sequence: 0,
hlc_ts: now,
hlc_count: 0,
database,
collection,
operation,
document_key,
document_data,
shard_id: None,
version_vector: Some(version_vector),
parent_vectors: Vec::new(),
is_delta: false,
delta_data: None,
session_id: None,
device_id: Some(origin_node),
}
}
pub fn set_version_vector(&mut self, vector: VersionVector) {
self.version_vector = Some(vector);
}
pub fn add_parent_vector(&mut self, vector: VersionVector) {
self.parent_vectors.push(vector);
}
pub fn set_delta(&mut self, patch_data: Vec<u8>) {
self.is_delta = true;
self.delta_data = Some(patch_data);
}
pub fn effective_vector(&self) -> Option<VersionVector> {
self.version_vector.clone()
}
}
pub fn compute_shard_id(key: &str, num_shards: u16) -> u16 {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut hasher = DefaultHasher::new();
key.hash(&mut hasher);
(hasher.finish() % num_shards as u64) as u16
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_sync_message_encode_decode() {
let msg = SyncMessage::Heartbeat {
node_id: "node1".to_string(),
sequence: 42,
stats: NodeStats::default(),
};
let encoded = msg.encode();
let decoded = SyncMessage::decode_frame(&encoded).unwrap();
match decoded {
SyncMessage::Heartbeat {
node_id, sequence, ..
} => {
assert_eq!(node_id, "node1");
assert_eq!(sequence, 42);
}
_ => panic!("Wrong message type"),
}
}
#[test]
fn test_compute_shard_id() {
let shard = compute_shard_id("doc123", 8);
assert!(shard < 8);
assert_eq!(compute_shard_id("doc123", 8), shard);
}
#[test]
fn a_document_batch_survives_a_round_trip() {
let batch = vec![
serde_json::json!({"_key": "a", "n": 1, "nested": {"deep": [1, 2, 3]}}),
serde_json::json!({"_key": "b", "text": "héllo", "flag": true, "nil": null}),
serde_json::json!({"_key": "c", "float": 1.5, "big": 9007199254740991i64}),
];
let encoded = encode_documents(&batch).expect("encodes");
let decoded = decode_documents(&encoded).expect("decodes");
assert_eq!(decoded, batch);
}
#[test]
fn an_empty_batch_round_trips_as_empty() {
let encoded = encode_documents(&[]).expect("encodes");
assert!(decode_documents(&encoded).expect("decodes").is_empty());
}
#[test]
fn garbage_fails_to_decode_rather_than_yielding_no_documents() {
assert!(decode_documents(b"\x00\x01\x02not json").is_err());
}
#[test]
fn the_system_collections_a_second_node_needs_survive_the_round_trip() {
let admins = vec![serde_json::json!({
"_key": "admin",
"username": "admin",
"password_hash": "$argon2id$v=19$m=19456,t=2,p=1$c2FsdA$aGFzaA",
"roles": ["admin"],
"created_at": 1785000000
})];
let decoded = decode_documents(&encode_documents(&admins).unwrap()).unwrap();
assert_eq!(decoded, admins);
assert_eq!(decoded[0]["username"], "admin");
}
#[test]
fn a_frame_shorter_than_its_own_header_is_refused() {
let error = SyncMessage::decode_frame(&[0, 0, 0])
.unwrap_err()
.to_string();
assert!(error.contains("shorter than"), "{error}");
}
#[test]
fn a_frame_that_lies_about_its_length_is_refused() {
let mut frame = SyncMessage::FullSyncComplete { final_sequence: 7 }.encode();
frame.push(0xff);
let error = SyncMessage::decode_frame(&frame).unwrap_err().to_string();
assert!(error.contains("declares"), "{error}");
}
#[test]
fn the_header_length_matches_what_encode_writes() {
let frame = SyncMessage::FullSyncComplete { final_sequence: 1 }.encode();
let declared = u32::from_be_bytes([frame[1], frame[2], frame[3], frame[4]]) as usize;
assert_eq!(frame.len(), SyncMessage::HEADER_LEN + declared);
}
}