367 lines
11 KiB
Rust
367 lines
11 KiB
Rust
use crate::pg_wal_parser::{ColumnValue, LogicalMessage, MessageParser, RelationInfo, TupleData};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use tokio::sync::broadcast;
|
|
|
|
/// High-level cache invalidation event emitted at transaction commit
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct CacheInvalidationEvent {
|
|
/// Transaction ID from PostgreSQL
|
|
pub xid: u32,
|
|
/// Commit LSN/watermark for checkpoint
|
|
pub commit_lsn: u64,
|
|
/// End LSN of the transaction
|
|
pub end_lsn: u64,
|
|
/// Timestamp of commit (microseconds since PG epoch)
|
|
pub timestamp: i64,
|
|
/// All operations in this transaction, emitted atomically
|
|
pub changes: Vec<Change>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct Change {
|
|
/// Table name that changed
|
|
pub table: String,
|
|
/// Schema/namespace
|
|
pub schema: String,
|
|
/// Operation type
|
|
pub operation: Operation,
|
|
/// Row data if available
|
|
pub row_data: Option<RowData>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum Operation {
|
|
Insert,
|
|
Update,
|
|
Delete,
|
|
Truncate,
|
|
}
|
|
|
|
/// Simplified row representation for broadcasting
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct RowData {
|
|
/// Mapping of column name -> value
|
|
pub columns: HashMap<String, CacheValue>,
|
|
/// Primary key values if available
|
|
pub key_values: Option<HashMap<String, CacheValue>>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum CacheValue {
|
|
Null,
|
|
String(String),
|
|
Binary(Vec<u8>),
|
|
UnchangedToasted,
|
|
}
|
|
|
|
/// Transaction buffer for a single transaction
|
|
struct TransactionBuffer {
|
|
xid: u32,
|
|
changes: Vec<Change>,
|
|
}
|
|
|
|
pub struct CacheInvalidationEmitter {
|
|
parser: MessageParser,
|
|
tx: broadcast::Sender<CacheInvalidationEvent>,
|
|
/// Buffers for active transactions, keyed by XID
|
|
txn_buffers: HashMap<u32, TransactionBuffer>,
|
|
/// Track the current transaction being processed
|
|
current_xid: Option<u32>,
|
|
}
|
|
|
|
impl CacheInvalidationEmitter {
|
|
pub fn new(tx: broadcast::Sender<CacheInvalidationEvent>) -> Self {
|
|
Self {
|
|
parser: MessageParser::new(),
|
|
tx,
|
|
txn_buffers: HashMap::new(),
|
|
current_xid: None,
|
|
}
|
|
}
|
|
|
|
/// Process a single logical replication message
|
|
/// Events are buffered until Commit, then emitted atomically
|
|
pub async fn process_message(&mut self, msg: LogicalMessage) -> Result<(), String> {
|
|
match msg {
|
|
LogicalMessage::Begin { xid, .. } => {
|
|
// Start buffering for this transaction
|
|
self.current_xid = Some(xid);
|
|
self.txn_buffers.insert(xid, TransactionBuffer {
|
|
xid,
|
|
changes: Vec::new(),
|
|
});
|
|
Ok(())
|
|
}
|
|
LogicalMessage::Relation(_) => {
|
|
// Relation updates internal state via parser
|
|
Ok(())
|
|
}
|
|
LogicalMessage::Insert {
|
|
xid,
|
|
rel_oid,
|
|
new_tuple,
|
|
} => {
|
|
let change = self.build_insert_change(rel_oid, new_tuple)?;
|
|
self.buffer_change(xid, change);
|
|
Ok(())
|
|
}
|
|
LogicalMessage::Update {
|
|
xid,
|
|
rel_oid,
|
|
key_tuple,
|
|
old_tuple: _,
|
|
new_tuple,
|
|
} => {
|
|
let change = self.build_update_change(rel_oid, new_tuple, key_tuple)?;
|
|
self.buffer_change(xid, change);
|
|
Ok(())
|
|
}
|
|
LogicalMessage::Delete {
|
|
xid,
|
|
rel_oid,
|
|
key_tuple,
|
|
old_tuple,
|
|
} => {
|
|
let change = self.build_delete_change(rel_oid, key_tuple.or(old_tuple))?;
|
|
self.buffer_change(xid, change);
|
|
Ok(())
|
|
}
|
|
LogicalMessage::Commit {
|
|
commit_lsn,
|
|
end_lsn,
|
|
timestamp,
|
|
..
|
|
} => {
|
|
// Emit the buffered transaction atomically
|
|
if let Some(xid) = self.current_xid.take() {
|
|
self.emit_transaction(xid, commit_lsn, end_lsn, timestamp)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
LogicalMessage::Truncate {
|
|
xid,
|
|
rel_oids,
|
|
..
|
|
} => {
|
|
for oid in rel_oids {
|
|
let change = self.build_truncate_change(oid)?;
|
|
self.buffer_change(xid, change);
|
|
}
|
|
Ok(())
|
|
}
|
|
LogicalMessage::Type { .. } => {
|
|
// Type information, no action needed
|
|
Ok(())
|
|
}
|
|
LogicalMessage::Message { .. } => {
|
|
// Custom messages, can be ignored for cache invalidation
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
|
|
fn buffer_change(&mut self, xid: u32, change: Change) {
|
|
if let Some(buf) = self.txn_buffers.get_mut(&xid) {
|
|
buf.changes.push(change);
|
|
}
|
|
}
|
|
|
|
fn build_insert_change(&self, rel_oid: u32, new_tuple: TupleData) -> Result<Change, String> {
|
|
let relation = self
|
|
.parser
|
|
.get_relation(rel_oid)
|
|
.ok_or_else(|| format!("Relation OID {} not found", rel_oid))?;
|
|
|
|
let row_data = self.build_row_data(relation, &new_tuple)?;
|
|
|
|
Ok(Change {
|
|
table: relation.name.clone(),
|
|
schema: relation.namespace.clone(),
|
|
operation: Operation::Insert,
|
|
row_data: Some(row_data),
|
|
})
|
|
}
|
|
|
|
fn build_update_change(
|
|
&self,
|
|
rel_oid: u32,
|
|
new_tuple: TupleData,
|
|
key_tuple: Option<TupleData>,
|
|
) -> Result<Change, String> {
|
|
let relation = self
|
|
.parser
|
|
.get_relation(rel_oid)
|
|
.ok_or_else(|| format!("Relation OID {} not found", rel_oid))?;
|
|
|
|
let mut row_data = self.build_row_data(relation, &new_tuple)?;
|
|
|
|
// If we have a key tuple, extract key values
|
|
if let Some(key) = key_tuple {
|
|
let key_values = self.extract_key_values(relation, &key)?;
|
|
row_data.key_values = Some(key_values);
|
|
}
|
|
|
|
Ok(Change {
|
|
table: relation.name.clone(),
|
|
schema: relation.namespace.clone(),
|
|
operation: Operation::Update,
|
|
row_data: Some(row_data),
|
|
})
|
|
}
|
|
|
|
fn build_delete_change(
|
|
&self,
|
|
rel_oid: u32,
|
|
key_or_old: Option<TupleData>,
|
|
) -> Result<Change, String> {
|
|
let relation = self
|
|
.parser
|
|
.get_relation(rel_oid)
|
|
.ok_or_else(|| format!("Relation OID {} not found", rel_oid))?;
|
|
|
|
let mut row_data = None;
|
|
if let Some(tuple) = key_or_old {
|
|
let mut data = self.build_row_data(relation, &tuple)?;
|
|
let key_values = self.extract_key_values(relation, &tuple)?;
|
|
data.key_values = Some(key_values);
|
|
row_data = Some(data);
|
|
}
|
|
|
|
Ok(Change {
|
|
table: relation.name.clone(),
|
|
schema: relation.namespace.clone(),
|
|
operation: Operation::Delete,
|
|
row_data,
|
|
})
|
|
}
|
|
|
|
fn build_truncate_change(&self, rel_oid: u32) -> Result<Change, String> {
|
|
let relation = self
|
|
.parser
|
|
.get_relation(rel_oid)
|
|
.ok_or_else(|| format!("Relation OID {} not found", rel_oid))?;
|
|
|
|
Ok(Change {
|
|
table: relation.name.clone(),
|
|
schema: relation.namespace.clone(),
|
|
operation: Operation::Truncate,
|
|
row_data: None,
|
|
})
|
|
}
|
|
|
|
fn build_row_data(&self, relation: &RelationInfo, tuple: &TupleData) -> Result<RowData, String> {
|
|
let mut columns = HashMap::new();
|
|
|
|
for (i, value) in tuple.values.iter().enumerate() {
|
|
if i >= relation.columns.len() {
|
|
return Err(format!(
|
|
"Tuple has more values than relation has columns: {} vs {}",
|
|
tuple.values.len(),
|
|
relation.columns.len()
|
|
));
|
|
}
|
|
|
|
let col = &relation.columns[i];
|
|
let cache_value = match value {
|
|
ColumnValue::Null => CacheValue::Null,
|
|
ColumnValue::UnchangedToasted => CacheValue::UnchangedToasted,
|
|
ColumnValue::Text(s) => CacheValue::String(s.clone()),
|
|
ColumnValue::Binary(b) => CacheValue::Binary(b.clone()),
|
|
};
|
|
|
|
columns.insert(col.name.clone(), cache_value);
|
|
}
|
|
|
|
Ok(RowData {
|
|
columns,
|
|
key_values: None,
|
|
})
|
|
}
|
|
|
|
fn extract_key_values(
|
|
&self,
|
|
relation: &RelationInfo,
|
|
key_tuple: &TupleData,
|
|
) -> Result<HashMap<String, CacheValue>, String> {
|
|
let mut key_values = HashMap::new();
|
|
|
|
for (i, value) in key_tuple.values.iter().enumerate() {
|
|
if i >= relation.columns.len() {
|
|
break; // Key might have fewer columns
|
|
}
|
|
|
|
let col = &relation.columns[i];
|
|
if col.flags & 1 != 0 {
|
|
// This column is part of the key
|
|
let cache_value = match value {
|
|
ColumnValue::Null => CacheValue::Null,
|
|
ColumnValue::UnchangedToasted => CacheValue::UnchangedToasted,
|
|
ColumnValue::Text(s) => CacheValue::String(s.clone()),
|
|
ColumnValue::Binary(b) => CacheValue::Binary(b.clone()),
|
|
};
|
|
|
|
key_values.insert(col.name.clone(), cache_value);
|
|
}
|
|
}
|
|
|
|
Ok(key_values)
|
|
}
|
|
|
|
/// Emit a buffered transaction atomically
|
|
fn emit_transaction(&mut self, xid: u32, commit_lsn: u64, end_lsn: u64, timestamp: i64) -> Result<(), String> {
|
|
if let Some(buf) = self.txn_buffers.remove(&xid) {
|
|
let event = CacheInvalidationEvent {
|
|
xid,
|
|
commit_lsn,
|
|
end_lsn,
|
|
timestamp,
|
|
changes: buf.changes,
|
|
};
|
|
|
|
self.tx.send(event).ok(); // ignore if no receivers
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Parse raw replication data bytes and process messages
|
|
/// This is the main entry point from pgwire-replication
|
|
pub async fn process_raw(&mut self, data: &[u8]) -> Result<(), String> {
|
|
let mut buf = bytes::BytesMut::from(&data[..]);
|
|
|
|
while !buf.is_empty() {
|
|
if let Some(msg) = self.parser.parse_message(&mut buf)? {
|
|
self.process_message(msg).await?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Get statistics about buffered transactions
|
|
pub fn buffered_transaction_count(&self) -> usize {
|
|
self.txn_buffers.len()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_emitter_creation() {
|
|
let (tx, _rx) = broadcast::channel(100);
|
|
let _emitter = CacheInvalidationEmitter::new(tx);
|
|
}
|
|
|
|
#[test]
|
|
fn test_transaction_buffering() {
|
|
let (tx, _rx) = broadcast::channel(100);
|
|
let emitter = CacheInvalidationEmitter::new(tx);
|
|
|
|
// Changes should be buffered per transaction until commit
|
|
assert_eq!(emitter.buffered_transaction_count(), 0);
|
|
}
|
|
}
|