Skip to main content

heliosdb_proxy/
transaction_journal.rs

1//! Transaction Journal - TR (Transaction Replay)
2//!
3//! Logs all statements within a transaction for replay after failover.
4//! Enables Oracle-grade TAF+TAC merged functionality.
5
6use super::{NodeId, ProxyError, Result};
7use std::collections::HashMap;
8use std::sync::Arc;
9use tokio::sync::RwLock;
10use uuid::Uuid;
11
12/// Journal entry for a single statement
13#[derive(Debug, Clone)]
14pub struct JournalEntry {
15    /// Entry sequence number
16    pub sequence: u64,
17    /// SQL statement text
18    pub statement: String,
19    /// Bound parameters
20    pub parameters: Vec<JournalValue>,
21    /// Result checksum (for verification after replay)
22    pub result_checksum: Option<u64>,
23    /// Number of rows affected
24    pub rows_affected: Option<u64>,
25    /// Timestamp
26    pub timestamp: chrono::DateTime<chrono::Utc>,
27    /// Statement type
28    pub statement_type: StatementType,
29    /// Execution duration (ms)
30    pub duration_ms: u64,
31}
32
33/// Serializable parameter value
34#[derive(Debug, Clone)]
35pub enum JournalValue {
36    Null,
37    Bool(bool),
38    Int64(i64),
39    Float64(f64),
40    Text(String),
41    Bytes(Vec<u8>),
42    Array(Vec<JournalValue>),
43}
44
45/// Statement type classification
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum StatementType {
48    /// SELECT query
49    Select,
50    /// INSERT statement
51    Insert,
52    /// UPDATE statement
53    Update,
54    /// DELETE statement
55    Delete,
56    /// DDL (CREATE, ALTER, DROP)
57    Ddl,
58    /// Transaction control (BEGIN, COMMIT, ROLLBACK)
59    Transaction,
60    /// SET statement
61    Set,
62    /// Other/unknown
63    Other,
64}
65
66impl StatementType {
67    /// Determine statement type from SQL
68    pub fn from_sql(sql: &str) -> Self {
69        let upper = sql.trim().to_uppercase();
70        if upper.starts_with("SELECT") {
71            StatementType::Select
72        } else if upper.starts_with("INSERT") {
73            StatementType::Insert
74        } else if upper.starts_with("UPDATE") {
75            StatementType::Update
76        } else if upper.starts_with("DELETE") {
77            StatementType::Delete
78        } else if upper.starts_with("CREATE")
79            || upper.starts_with("ALTER")
80            || upper.starts_with("DROP")
81        {
82            StatementType::Ddl
83        } else if upper.starts_with("BEGIN")
84            || upper.starts_with("COMMIT")
85            || upper.starts_with("ROLLBACK")
86            || upper.starts_with("SAVEPOINT")
87        {
88            StatementType::Transaction
89        } else if upper.starts_with("SET") {
90            StatementType::Set
91        } else {
92            StatementType::Other
93        }
94    }
95
96    /// Is this a read-only statement?
97    pub fn is_read_only(&self) -> bool {
98        matches!(self, StatementType::Select)
99    }
100
101    /// Is this a mutating statement?
102    pub fn is_mutation(&self) -> bool {
103        matches!(
104            self,
105            StatementType::Insert
106                | StatementType::Update
107                | StatementType::Delete
108                | StatementType::Ddl
109        )
110    }
111}
112
113/// Transaction journal for a single transaction
114#[derive(Debug, Clone)]
115pub struct TransactionJournalEntry {
116    /// Transaction ID
117    pub tx_id: Uuid,
118    /// Session ID
119    pub session_id: Uuid,
120    /// Node where transaction started
121    pub node_id: NodeId,
122    /// Transaction start time
123    pub started_at: chrono::DateTime<chrono::Utc>,
124    /// Start LSN (for WAL synchronization)
125    pub start_lsn: u64,
126    /// Journal entries
127    pub entries: Vec<JournalEntry>,
128    /// Current sequence
129    pub current_sequence: u64,
130    /// Is transaction active
131    pub active: bool,
132    /// Has mutations
133    pub has_mutations: bool,
134    /// Savepoints
135    pub savepoints: Vec<Savepoint>,
136}
137
138/// Savepoint information
139#[derive(Debug, Clone)]
140pub struct Savepoint {
141    /// Savepoint name
142    pub name: String,
143    /// Sequence at savepoint
144    pub sequence: u64,
145    /// Created timestamp
146    pub created_at: chrono::DateTime<chrono::Utc>,
147}
148
149impl TransactionJournalEntry {
150    /// Create a new transaction journal entry
151    pub fn new(tx_id: Uuid, session_id: Uuid, node_id: NodeId, start_lsn: u64) -> Self {
152        Self {
153            tx_id,
154            session_id,
155            node_id,
156            started_at: chrono::Utc::now(),
157            start_lsn,
158            entries: Vec::new(),
159            current_sequence: 0,
160            active: true,
161            has_mutations: false,
162            savepoints: Vec::new(),
163        }
164    }
165
166    /// Add an entry to the journal
167    pub fn add_entry(&mut self, entry: JournalEntry) {
168        if entry.statement_type.is_mutation() {
169            self.has_mutations = true;
170        }
171        self.current_sequence = entry.sequence;
172        self.entries.push(entry);
173    }
174
175    /// Create a savepoint
176    pub fn create_savepoint(&mut self, name: String) {
177        self.savepoints.push(Savepoint {
178            name,
179            sequence: self.current_sequence,
180            created_at: chrono::Utc::now(),
181        });
182    }
183
184    /// Rollback to savepoint
185    pub fn rollback_to_savepoint(&mut self, name: &str) -> Option<u64> {
186        if let Some(idx) = self.savepoints.iter().position(|s| s.name == name) {
187            let savepoint = &self.savepoints[idx];
188            let sequence = savepoint.sequence;
189
190            // Truncate entries after savepoint
191            self.entries.retain(|e| e.sequence <= sequence);
192
193            // Remove later savepoints
194            self.savepoints.truncate(idx + 1);
195
196            Some(sequence)
197        } else {
198            None
199        }
200    }
201
202    /// Get entries for replay
203    pub fn entries_for_replay(&self) -> Vec<&JournalEntry> {
204        self.entries.iter().collect()
205    }
206
207    /// Get only mutation entries
208    pub fn mutation_entries(&self) -> Vec<&JournalEntry> {
209        self.entries
210            .iter()
211            .filter(|e| e.statement_type.is_mutation())
212            .collect()
213    }
214
215    /// Calculate total size of journal
216    pub fn total_size(&self) -> usize {
217        self.entries
218            .iter()
219            .map(|e| e.statement.len() + estimate_params_size(&e.parameters))
220            .sum()
221    }
222}
223
224fn estimate_params_size(params: &[JournalValue]) -> usize {
225    params
226        .iter()
227        .map(|p| match p {
228            JournalValue::Null => 1,
229            JournalValue::Bool(_) => 1,
230            JournalValue::Int64(_) => 8,
231            JournalValue::Float64(_) => 8,
232            JournalValue::Text(s) => s.len(),
233            JournalValue::Bytes(b) => b.len(),
234            JournalValue::Array(a) => estimate_params_size(a),
235        })
236        .sum()
237}
238
239/// Transaction Journal Manager
240pub struct TransactionJournal {
241    /// Active transaction journals
242    journals: Arc<RwLock<HashMap<Uuid, TransactionJournalEntry>>>,
243    /// Maximum entries per journal
244    max_entries: usize,
245    /// Maximum journal size (bytes)
246    max_size: usize,
247    /// Global cap on the number of retained transaction journals. The data-path
248    /// write journaling records each write as its own auto-commit transaction
249    /// (a fresh tx_id, begin + log, never committed), so without a global bound
250    /// the map grows by one entry per write forever — an unbounded leak of the
251    /// full SQL of every write. When the cap is reached the oldest journals
252    /// (by start time) are evicted; replay only consults recent history.
253    max_journals: usize,
254    /// Whether journaling is enabled
255    enabled: bool,
256}
257
258impl TransactionJournal {
259    /// Create a new transaction journal manager
260    pub fn new() -> Self {
261        Self {
262            journals: Arc::new(RwLock::new(HashMap::new())),
263            max_entries: 10000,
264            max_size: 64 * 1024 * 1024, // 64MB
265            max_journals: 50_000,
266            enabled: true,
267        }
268    }
269
270    /// Configure maximum entries
271    pub fn with_max_entries(mut self, max: usize) -> Self {
272        self.max_entries = max;
273        self
274    }
275
276    /// Configure the global cap on retained journals.
277    pub fn with_max_journals(mut self, max: usize) -> Self {
278        self.max_journals = max.max(1);
279        self
280    }
281
282    /// Evict the oldest journals (by `started_at`) until at most `target_len`
283    /// remain. Called under the write lock when the global cap is hit; evicting
284    /// a batch down to a target amortizes the O(n) scan across many inserts.
285    fn evict_oldest_locked(
286        journals: &mut HashMap<Uuid, TransactionJournalEntry>,
287        target_len: usize,
288    ) {
289        if journals.len() <= target_len {
290            return;
291        }
292        let remove_count = journals.len() - target_len;
293        let mut by_time: Vec<(Uuid, chrono::DateTime<chrono::Utc>)> =
294            journals.iter().map(|(k, v)| (*k, v.started_at)).collect();
295        // Partial order is enough, but a full sort is simple and only runs when
296        // the cap is hit (once per `max_journals/10` writes).
297        by_time.sort_by_key(|(_, t)| *t);
298        for (k, _) in by_time.into_iter().take(remove_count) {
299            journals.remove(&k);
300        }
301    }
302
303    /// Configure maximum size
304    pub fn with_max_size(mut self, max: usize) -> Self {
305        self.max_size = max;
306        self
307    }
308
309    /// Enable or disable journaling
310    pub fn set_enabled(&mut self, enabled: bool) {
311        self.enabled = enabled;
312    }
313
314    /// Collect every journal entry across every active transaction
315    /// whose `timestamp` falls within the inclusive window
316    /// `[from, to]`. Results are sorted in timestamp order so the
317    /// caller can replay them chronologically regardless of which
318    /// transaction they came from.
319    ///
320    /// Used by the time-travel replay engine (`src/replay/`) to
321    /// reconstruct "what happened at the source between these two
322    /// timestamps" against a staging target.
323    pub async fn entries_in_window(
324        &self,
325        from: chrono::DateTime<chrono::Utc>,
326        to: chrono::DateTime<chrono::Utc>,
327    ) -> Vec<(Uuid, JournalEntry)> {
328        let journals = self.journals.read().await;
329        let mut out: Vec<(Uuid, JournalEntry)> = Vec::new();
330        for (tx_id, j) in journals.iter() {
331            for entry in &j.entries {
332                if entry.timestamp >= from && entry.timestamp <= to {
333                    out.push((*tx_id, entry.clone()));
334                }
335            }
336        }
337        out.sort_by_key(|(_, e)| e.timestamp);
338        out
339    }
340
341    /// Start journaling a transaction
342    pub async fn begin_transaction(
343        &self,
344        tx_id: Uuid,
345        session_id: Uuid,
346        node_id: NodeId,
347        start_lsn: u64,
348    ) -> Result<()> {
349        if !self.enabled {
350            return Ok(());
351        }
352
353        let journal = TransactionJournalEntry::new(tx_id, session_id, node_id, start_lsn);
354        let mut journals = self.journals.write().await;
355        // Bound total retained journals: the data path never commits (each write
356        // is its own tx_id), so without this the map grows forever. Evict down
357        // to 90% of the cap in one pass to amortize the sort.
358        if journals.len() >= self.max_journals {
359            let target = (self.max_journals * 9 / 10).max(1);
360            Self::evict_oldest_locked(&mut journals, target);
361        }
362        journals.insert(tx_id, journal);
363        drop(journals);
364
365        tracing::debug!("Started journaling transaction {:?}", tx_id);
366        Ok(())
367    }
368
369    /// Log a statement
370    pub async fn log_statement(
371        &self,
372        tx_id: Uuid,
373        statement: String,
374        parameters: Vec<JournalValue>,
375        result_checksum: Option<u64>,
376        rows_affected: Option<u64>,
377        duration_ms: u64,
378    ) -> Result<()> {
379        if !self.enabled {
380            return Ok(());
381        }
382
383        let mut journals = self.journals.write().await;
384        let journal = journals.get_mut(&tx_id).ok_or_else(|| {
385            ProxyError::Internal(format!("No journal for transaction {:?}", tx_id))
386        })?;
387
388        // Check limits
389        if journal.entries.len() >= self.max_entries {
390            return Err(ProxyError::Internal(
391                "Transaction journal entries limit exceeded".to_string(),
392            ));
393        }
394
395        if journal.total_size() >= self.max_size {
396            return Err(ProxyError::Internal(
397                "Transaction journal size limit exceeded".to_string(),
398            ));
399        }
400
401        let sequence = journal.current_sequence + 1;
402        let statement_type = StatementType::from_sql(&statement);
403
404        let entry = JournalEntry {
405            sequence,
406            statement,
407            parameters,
408            result_checksum,
409            rows_affected,
410            timestamp: chrono::Utc::now(),
411            statement_type,
412            duration_ms,
413        };
414
415        journal.add_entry(entry);
416
417        Ok(())
418    }
419
420    /// Create a savepoint
421    pub async fn create_savepoint(&self, tx_id: Uuid, name: String) -> Result<()> {
422        if !self.enabled {
423            return Ok(());
424        }
425
426        let mut journals = self.journals.write().await;
427        let journal = journals.get_mut(&tx_id).ok_or_else(|| {
428            ProxyError::Internal(format!("No journal for transaction {:?}", tx_id))
429        })?;
430
431        journal.create_savepoint(name);
432        Ok(())
433    }
434
435    /// Rollback to savepoint
436    pub async fn rollback_to_savepoint(&self, tx_id: Uuid, name: &str) -> Result<()> {
437        if !self.enabled {
438            return Ok(());
439        }
440
441        let mut journals = self.journals.write().await;
442        let journal = journals.get_mut(&tx_id).ok_or_else(|| {
443            ProxyError::Internal(format!("No journal for transaction {:?}", tx_id))
444        })?;
445
446        journal
447            .rollback_to_savepoint(name)
448            .ok_or_else(|| ProxyError::Internal(format!("Savepoint '{}' not found", name)))?;
449
450        Ok(())
451    }
452
453    /// Commit transaction (clear journal)
454    pub async fn commit_transaction(&self, tx_id: Uuid) -> Result<()> {
455        self.journals.write().await.remove(&tx_id);
456        tracing::debug!("Committed and cleared journal for transaction {:?}", tx_id);
457        Ok(())
458    }
459
460    /// Rollback transaction (clear journal)
461    pub async fn rollback_transaction(&self, tx_id: Uuid) -> Result<()> {
462        self.journals.write().await.remove(&tx_id);
463        tracing::debug!(
464            "Rolled back and cleared journal for transaction {:?}",
465            tx_id
466        );
467        Ok(())
468    }
469
470    /// Get journal for a transaction (for replay)
471    pub async fn get_journal(&self, tx_id: &Uuid) -> Option<TransactionJournalEntry> {
472        self.journals.read().await.get(tx_id).cloned()
473    }
474
475    /// Get active transaction count
476    pub async fn active_count(&self) -> usize {
477        self.journals.read().await.len()
478    }
479
480    /// Get statistics
481    pub async fn stats(&self) -> JournalStats {
482        let journals = self.journals.read().await;
483        let total_entries: usize = journals.values().map(|j| j.entries.len()).sum();
484        let total_size: usize = journals.values().map(|j| j.total_size()).sum();
485
486        JournalStats {
487            active_transactions: journals.len(),
488            total_entries,
489            total_size_bytes: total_size,
490            enabled: self.enabled,
491        }
492    }
493
494    /// Get all active transaction journals (for failover replay)
495    pub async fn get_all_active(&self) -> Vec<TransactionJournalEntry> {
496        self.journals.read().await.values().cloned().collect()
497    }
498
499    /// Get the maximum start LSN across all active transactions
500    /// Used to determine how far the standby needs to catch up
501    pub async fn get_max_start_lsn(&self) -> Option<u64> {
502        let journals = self.journals.read().await;
503        journals.values().map(|j| j.start_lsn).max()
504    }
505
506    /// Get transactions that started on a specific node
507    /// Useful for replaying only transactions affected by a node failure
508    pub async fn get_transactions_for_node(&self, node_id: NodeId) -> Vec<TransactionJournalEntry> {
509        self.journals
510            .read()
511            .await
512            .values()
513            .filter(|j| j.node_id == node_id)
514            .cloned()
515            .collect()
516    }
517}
518
519impl Default for TransactionJournal {
520    fn default() -> Self {
521        Self::new()
522    }
523}
524
525/// Journal statistics
526#[derive(Debug, Clone)]
527pub struct JournalStats {
528    /// Number of active transactions being journaled
529    pub active_transactions: usize,
530    /// Total journal entries across all transactions
531    pub total_entries: usize,
532    /// Total size of journals in bytes
533    pub total_size_bytes: usize,
534    /// Whether journaling is enabled
535    pub enabled: bool,
536}
537
538#[cfg(test)]
539mod tests {
540    use super::*;
541
542    #[test]
543    fn test_statement_type_detection() {
544        assert_eq!(
545            StatementType::from_sql("SELECT * FROM users"),
546            StatementType::Select
547        );
548        assert_eq!(
549            StatementType::from_sql("INSERT INTO users VALUES (1)"),
550            StatementType::Insert
551        );
552        assert_eq!(
553            StatementType::from_sql("UPDATE users SET name = 'x'"),
554            StatementType::Update
555        );
556        assert_eq!(
557            StatementType::from_sql("DELETE FROM users"),
558            StatementType::Delete
559        );
560        assert_eq!(
561            StatementType::from_sql("CREATE TABLE foo (id INT)"),
562            StatementType::Ddl
563        );
564        assert_eq!(StatementType::from_sql("BEGIN"), StatementType::Transaction);
565        assert_eq!(
566            StatementType::from_sql("SET search_path = public"),
567            StatementType::Set
568        );
569    }
570
571    #[test]
572    fn test_statement_type_properties() {
573        assert!(StatementType::Select.is_read_only());
574        assert!(!StatementType::Insert.is_read_only());
575
576        assert!(StatementType::Insert.is_mutation());
577        assert!(StatementType::Update.is_mutation());
578        assert!(!StatementType::Select.is_mutation());
579    }
580
581    #[tokio::test]
582    async fn test_journal_lifecycle() {
583        let journal = TransactionJournal::new();
584        let tx_id = Uuid::new_v4();
585        let session_id = Uuid::new_v4();
586        let node_id = NodeId::new();
587
588        // Begin transaction
589        journal
590            .begin_transaction(tx_id, session_id, node_id, 0)
591            .await
592            .unwrap();
593
594        // Log statements
595        journal
596            .log_statement(
597                tx_id,
598                "SELECT * FROM users".to_string(),
599                vec![],
600                Some(12345),
601                None,
602                10,
603            )
604            .await
605            .unwrap();
606
607        journal
608            .log_statement(
609                tx_id,
610                "INSERT INTO users (name) VALUES ($1)".to_string(),
611                vec![JournalValue::Text("test".to_string())],
612                None,
613                Some(1),
614                5,
615            )
616            .await
617            .unwrap();
618
619        // Check journal
620        let j = journal.get_journal(&tx_id).await.unwrap();
621        assert_eq!(j.entries.len(), 2);
622        assert!(j.has_mutations);
623
624        // Commit
625        journal.commit_transaction(tx_id).await.unwrap();
626        assert!(journal.get_journal(&tx_id).await.is_none());
627    }
628
629    #[tokio::test]
630    async fn test_savepoints() {
631        let journal = TransactionJournal::new();
632        let tx_id = Uuid::new_v4();
633        let session_id = Uuid::new_v4();
634        let node_id = NodeId::new();
635
636        journal
637            .begin_transaction(tx_id, session_id, node_id, 0)
638            .await
639            .unwrap();
640
641        // Log some statements
642        for i in 0..3 {
643            journal
644                .log_statement(
645                    tx_id,
646                    format!("INSERT INTO t VALUES ({})", i),
647                    vec![],
648                    None,
649                    Some(1),
650                    1,
651                )
652                .await
653                .unwrap();
654        }
655
656        // Create savepoint
657        journal
658            .create_savepoint(tx_id, "sp1".to_string())
659            .await
660            .unwrap();
661
662        // Log more
663        for i in 3..5 {
664            journal
665                .log_statement(
666                    tx_id,
667                    format!("INSERT INTO t VALUES ({})", i),
668                    vec![],
669                    None,
670                    Some(1),
671                    1,
672                )
673                .await
674                .unwrap();
675        }
676
677        let j = journal.get_journal(&tx_id).await.unwrap();
678        assert_eq!(j.entries.len(), 5);
679
680        // Rollback to savepoint
681        journal.rollback_to_savepoint(tx_id, "sp1").await.unwrap();
682
683        let j = journal.get_journal(&tx_id).await.unwrap();
684        assert_eq!(j.entries.len(), 3);
685    }
686
687    #[tokio::test]
688    async fn test_stats() {
689        let journal = TransactionJournal::new();
690        let tx_id = Uuid::new_v4();
691        let session_id = Uuid::new_v4();
692        let node_id = NodeId::new();
693
694        journal
695            .begin_transaction(tx_id, session_id, node_id, 0)
696            .await
697            .unwrap();
698        journal
699            .log_statement(tx_id, "SELECT 1".to_string(), vec![], None, None, 1)
700            .await
701            .unwrap();
702
703        let stats = journal.stats().await;
704        assert_eq!(stats.active_transactions, 1);
705        assert_eq!(stats.total_entries, 1);
706        assert!(stats.enabled);
707    }
708
709    /// The global journal cap must bound the map even when transactions are
710    /// never committed (the data-path auto-commit journaling pattern), evicting
711    /// the oldest journals rather than growing without limit.
712    #[tokio::test]
713    async fn global_cap_evicts_oldest_journals() {
714        let journal = TransactionJournal::new().with_max_journals(10);
715        let node_id = NodeId::new();
716        // Begin far more single-statement transactions than the cap, never
717        // committing any — mirrors journal_write on the query path.
718        for _ in 0..100 {
719            let tx = Uuid::new_v4();
720            journal
721                .begin_transaction(tx, Uuid::new_v4(), node_id, 0)
722                .await
723                .unwrap();
724            journal
725                .log_statement(
726                    tx,
727                    "INSERT INTO t VALUES (1)".to_string(),
728                    vec![],
729                    None,
730                    None,
731                    1,
732                )
733                .await
734                .unwrap();
735        }
736        // Bounded at the cap (not 100).
737        let stats = journal.stats().await;
738        assert!(
739            stats.active_transactions <= 10,
740            "journal map must stay within the cap, got {}",
741            stats.active_transactions
742        );
743    }
744}