use crate::ha::protocol::{
SyncAck, SyncDataBegin, SyncDataChunk, SyncDataEnd, SyncRequest, SyncType, MAX_CHUNK_DATA_SIZE,
};
use crate::ha::{HAError, Result, SyncState};
use crate::pubsub;
use crate::pubsub::topics::{
SYNC_ACK_TOPIC, SYNC_DATA_BEGIN_TOPIC, SYNC_DATA_CHUNK_TOPIC, SYNC_DATA_END_TOPIC,
SYNC_REQUEST_TOPIC,
};
use crate::transaction::LogItem;
use alloc::vec::Vec;
#[cfg(feature = "log")]
use crate::log::{debug, error, info, warn};
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum SyncHandlerState {
Idle,
Syncing,
Completed,
Failed,
}
impl From<SyncHandlerState> for SyncState {
fn from(state: SyncHandlerState) -> Self {
match state {
SyncHandlerState::Idle => SyncState::Idle,
SyncHandlerState::Syncing => SyncState::Syncing,
SyncHandlerState::Completed => SyncState::Synced,
SyncHandlerState::Failed => SyncState::Failed,
}
}
}
pub struct SyncHandler {
state: SyncHandlerState,
current_request: Option<SyncRequest>,
chunks_sent: u32,
bytes_sent: u64,
lock: u32,
}
impl SyncHandler {
pub fn new() -> Self {
Self {
state: SyncHandlerState::Idle,
current_request: None,
chunks_sent: 0,
bytes_sent: 0,
lock: 0,
}
}
pub fn init(&mut self) -> Result<()> {
#[cfg(feature = "log")]
debug!("SyncHandler: Initializing and subscribing to SYNC_REQUEST_TOPIC");
pubsub::subscribe(SYNC_REQUEST_TOPIC, Self::handle_sync_request_callback)
.map_err(|_| HAError::InitFailed)?;
pubsub::subscribe(SYNC_ACK_TOPIC, Self::handle_sync_ack_callback)
.map_err(|_| HAError::InitFailed)?;
#[cfg(feature = "log")]
info!("SyncHandler: Successfully initialized");
Ok(())
}
fn handle_sync_request_callback(topic_id: u16, data: &[u8]) -> bool {
if topic_id != SYNC_REQUEST_TOPIC {
return false;
}
#[cfg(feature = "log")]
debug!(
"SyncHandler: Received sync request, data len: {}",
data.len()
);
let request = match SyncRequest::decode(data) {
Some(req) => req,
None => {
#[cfg(feature = "log")]
error!("SyncHandler: Failed to decode sync request");
return false;
}
};
#[cfg(feature = "log")]
info!(
"SyncHandler: Sync request from slave {}, type: {:?}",
request.slave_id, request.sync_type
);
match request.sync_type {
SyncType::Full => {
Self::process_full_sync_request(request);
}
SyncType::Incremental => {
Self::process_incremental_sync_request(request);
}
}
true
}
fn handle_sync_ack_callback(topic_id: u16, data: &[u8]) -> bool {
if topic_id != SYNC_ACK_TOPIC {
return false;
}
let ack = match SyncAck::decode(data) {
Some(a) => a,
None => {
#[cfg(feature = "log")]
warn!("SyncHandler: Failed to decode sync ack");
return false;
}
};
#[cfg(feature = "log")]
info!(
"SyncHandler: Received sync ack from slave {}, success: {}, chunks: {}",
ack.slave_id, ack.success, ack.chunks_received
);
true
}
fn process_full_sync_request(request: SyncRequest) {
#[cfg(feature = "log")]
debug!(
"SyncHandler: Processing full sync request from slave {}",
request.slave_id
);
let snapshot_data = match Self::create_database_snapshot() {
Ok(data) => data,
Err(e) => {
#[cfg(feature = "log")]
error!("SyncHandler: Failed to create snapshot: {:?}", e);
return;
}
};
#[cfg(feature = "log")]
info!(
"SyncHandler: Created snapshot, size: {} bytes",
snapshot_data.len()
);
if let Err(e) = Self::send_snapshot_chunks(&snapshot_data) {
#[cfg(feature = "log")]
error!("SyncHandler: Failed to send snapshot chunks: {:?}", e);
}
}
fn process_incremental_sync_request(request: SyncRequest) {
#[cfg(feature = "log")]
debug!(
"SyncHandler: Processing incremental sync request from slave {}, last_log_index: {}",
request.slave_id, request.last_log_index
);
let wal_data = match Self::get_wal_logs_since(request.last_log_index) {
Ok(data) => data,
Err(e) => {
#[cfg(feature = "log")]
error!("SyncHandler: Failed to get WAL logs: {:?}", e);
return;
}
};
#[cfg(feature = "log")]
info!(
"SyncHandler: Retrieved WAL logs, size: {} bytes",
wal_data.len()
);
if let Err(e) = Self::send_wal_chunks(&wal_data) {
#[cfg(feature = "log")]
error!("SyncHandler: Failed to send WAL chunks: {:?}", e);
}
}
fn create_database_snapshot() -> Result<Vec<u8>> {
let mut snapshot = Vec::new();
let db = unsafe { crate::get_global_db() }.ok_or(HAError::SyncFailed)?;
unsafe {
let table_count = db.tables.len() as u8;
snapshot.push(table_count);
for (table_id, table_opt) in db.tables.iter().enumerate() {
if let Some(table) = table_opt {
let table_name = &table.def.name;
let name_bytes = table_name.as_bytes();
snapshot.push(name_bytes.len() as u8);
snapshot.extend_from_slice(name_bytes);
snapshot.extend_from_slice(&(table.record_size as u32).to_le_bytes());
snapshot.extend_from_slice(&(table.record_count as u32).to_le_bytes());
snapshot.extend_from_slice(&(table.def.max_records as u32).to_le_bytes());
snapshot.push(table.def.fields.len() as u8);
for field in &table.def.fields {
let field_name_bytes = field.name.as_bytes();
snapshot.push(field_name_bytes.len() as u8);
snapshot.extend_from_slice(field_name_bytes);
snapshot.push(field.data_type as u8);
snapshot.extend_from_slice(&(field.offset as u16).to_le_bytes());
let dimension = field
.vector_metadata
.as_ref()
.map(|vm| vm.dimension)
.unwrap_or(0);
snapshot.extend_from_slice(&dimension.to_le_bytes());
}
snapshot.push(table.def.primary_key.len() as u8);
for &pk_idx in &table.def.primary_key {
snapshot.push(pk_idx as u8);
}
for record_id in 0..table.def.max_records {
let status_ptr = table.get_status_ptr(record_id);
if (*status_ptr).status == crate::types::RecordStatus::Used {
snapshot.push(1);
snapshot.extend_from_slice(&(record_id as u32).to_le_bytes());
let record_ptr = table.get_record_ptr(record_id);
let record_data =
core::slice::from_raw_parts(record_ptr, table.record_size);
snapshot.extend_from_slice(record_data);
}
}
snapshot.push(0);
#[cfg(feature = "log")]
debug!(
"SyncHandler: Snapshotted table '{}' (id: {}), {} records",
table_name, table_id, table.record_count
);
}
}
}
Ok(snapshot)
}
fn get_wal_logs_since(_last_log_index: u32) -> Result<Vec<u8>> {
#[cfg(feature = "log")]
warn!("SyncHandler: WAL log retrieval not yet implemented");
Ok(Vec::new())
}
fn send_snapshot_chunks(data: &[u8]) -> Result<()> {
let total_size = data.len() as u64;
let chunk_count = (total_size as usize).div_ceil(MAX_CHUNK_DATA_SIZE) as u32;
let table_count = if !data.is_empty() { data[0] } else { 0 };
#[cfg(feature = "log")]
info!(
"SyncHandler: Sending {} chunks, total size: {} bytes, {} tables",
chunk_count, total_size, table_count
);
let begin = SyncDataBegin::new_snapshot(total_size, chunk_count, table_count);
let begin_data = begin.encode();
pubsub::publish(SYNC_DATA_BEGIN_TOPIC, &begin_data).map_err(|_| HAError::SyncFailed)?;
let mut offset = 0;
let mut chunk_index = 0;
while offset < data.len() {
let chunk_end = core::cmp::min(offset + MAX_CHUNK_DATA_SIZE, data.len());
let chunk_data = &data[offset..chunk_end];
let chunk = SyncDataChunk::new(chunk_index, chunk_data);
let encoded = chunk.encode();
pubsub::publish(SYNC_DATA_CHUNK_TOPIC, &encoded).map_err(|_| HAError::SyncFailed)?;
offset = chunk_end;
chunk_index += 1;
#[cfg(feature = "std")]
std::thread::sleep(std::time::Duration::from_micros(100));
}
let end = SyncDataEnd::new(chunk_count, 0); let end_data = end.encode();
pubsub::publish(SYNC_DATA_END_TOPIC, &end_data).map_err(|_| HAError::SyncFailed)?;
#[cfg(feature = "log")]
info!("SyncHandler: Completed sending {} chunks", chunk_count);
Ok(())
}
fn send_wal_chunks(data: &[u8]) -> Result<()> {
let total_size = data.len() as u64;
let chunk_count = (total_size as usize).div_ceil(MAX_CHUNK_DATA_SIZE) as u32;
let log_count = (total_size / core::mem::size_of::<LogItem>() as u64) as u32;
#[cfg(feature = "log")]
info!(
"SyncHandler: Sending {} WAL chunks, total size: {} bytes, ~{} logs",
chunk_count, total_size, log_count
);
let begin = SyncDataBegin::new_wal(total_size, chunk_count, log_count);
let begin_data = begin.encode();
pubsub::publish(SYNC_DATA_BEGIN_TOPIC, &begin_data).map_err(|_| HAError::SyncFailed)?;
let mut offset = 0;
let mut chunk_index = 0;
while offset < data.len() {
let chunk_end = core::cmp::min(offset + MAX_CHUNK_DATA_SIZE, data.len());
let chunk_data = &data[offset..chunk_end];
let chunk = SyncDataChunk::new(chunk_index, chunk_data);
let encoded = chunk.encode();
pubsub::publish(SYNC_DATA_CHUNK_TOPIC, &encoded).map_err(|_| HAError::SyncFailed)?;
offset = chunk_end;
chunk_index += 1;
#[cfg(feature = "std")]
std::thread::sleep(std::time::Duration::from_micros(100));
}
let end = SyncDataEnd::new(chunk_count, 0);
let end_data = end.encode();
pubsub::publish(SYNC_DATA_END_TOPIC, &end_data).map_err(|_| HAError::SyncFailed)?;
#[cfg(feature = "log")]
info!("SyncHandler: Completed sending {} WAL chunks", chunk_count);
Ok(())
}
pub fn shutdown(&mut self) -> Result<()> {
self.state = SyncHandlerState::Idle;
self.current_request = None;
self.chunks_sent = 0;
self.bytes_sent = 0;
Ok(())
}
pub fn get_state(&self) -> SyncHandlerState {
self.state
}
}
impl Default for SyncHandler {
fn default() -> Self {
Self::new()
}
}