feat: Implement rollback buffer (#49)

pallas-miniprotocols/src/chainsync/buffer.rs

@@ -0,0 +1,193 @@
+use std::collections::{vec_deque::Iter, VecDeque};
+
+use crate::Point;
+
+/// A memory buffer to handle chain rollbacks
+///
+/// This structure is intended to facilitate the process of managing rollbacks
+/// in a chain sync process. The goal is to keep points in memory until they
+/// reach a certain depth (# of confirmations). If a rollback happens, the
+/// buffer will try to find the intersection, clear the orphaned points and keep
+/// the remaining still in memory. Further forward rolls will accumulate from
+/// the intersection.
+///
+/// It works by keeping a `VecDeque` data structure of points, where
+/// roll-forward operations accumulate at the end of the deque and retrieving
+/// confirmed points means to pop from the front of the deque.
+///
+/// Notice that it works by keeping track of points, not blocks. It is meant to
+/// be used as a lightweight index where blocks can then be retrieved from a
+/// more suitable memory structure / persistent storage.
+#[derive(Debug)]
+pub struct RollbackBuffer {
+    points: VecDeque<Point>,
+}
+
+impl Default for RollbackBuffer {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+pub enum RollbackEffect {
+    Handled,
+    OutOfScope,
+}
+
+impl RollbackBuffer {
+    pub fn new() -> Self {
+        Self {
+            points: VecDeque::new(),
+        }
+    }
+
+    /// Adds a new point to the back of the buffer
+    pub fn roll_forward(&mut self, point: Point) {
+        self.points.push_back(point);
+    }
+
+    /// Retrieves all points above or equal a certain depth
+    pub fn pop_with_depth(&mut self, min_depth: usize) -> Vec<Point> {
+        match self.points.len().checked_sub(min_depth) {
+            Some(ready) => self.points.drain(0..ready).collect(),
+            None => vec![],
+        }
+    }
+
+    /// Find the position of a point within the buffer
+    pub fn position(&self, point: &Point) -> Option<usize> {
+        self.points.iter().position(|p| p.eq(point))
+    }
+
+    /// Iterates over the contents of the buffer
+    pub fn peek(&self) -> Iter<Point> {
+        self.points.iter()
+    }
+
+    /// Returns the size of the buffer (number of points)
+    pub fn size(&self) -> usize {
+        self.points.len()
+    }
+
+    /// Returns the newest point in the buffer
+    pub fn latest(&self) -> Option<&Point> {
+        self.points.back()
+    }
+
+    /// Returns the oldest point in the buffer
+    pub fn oldest(&self) -> Option<&Point> {
+        self.points.front()
+    }
+
+    /// Unwind the buffer up to a certain point, clearing orphaned items
+    ///
+    /// If the buffer contains the rollback point, we can safely discard from
+    /// the back and return Ok. If the rollback point is outside the scope of
+    /// the buffer, we clear the whole buffer and notify a failure
+    /// in the rollback process.
+    pub fn roll_back(&mut self, point: &Point) -> RollbackEffect {
+        if let Some(x) = self.position(point) {
+            self.points.truncate(x + 1);
+            RollbackEffect::Handled
+        } else {
+            self.points.clear();
+            RollbackEffect::OutOfScope
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{chainsync::RollbackEffect, Point};
+
+    use super::RollbackBuffer;
+
+    fn dummy_point(i: u64) -> Point {
+        Point::new(i, i.to_le_bytes().to_vec())
+    }
+
+    fn build_filled_buffer(n: usize) -> RollbackBuffer {
+        let mut buffer = RollbackBuffer::new();
+
+        for i in 0..n {
+            let point = dummy_point(i as u64);
+            buffer.roll_forward(point);
+        }
+
+        buffer
+    }
+
+    #[test]
+    fn roll_forward_accumulates_points() {
+        let buffer = build_filled_buffer(3);
+
+        assert!(matches!(buffer.position(&dummy_point(0)), Some(0)));
+        assert!(matches!(buffer.position(&dummy_point(1)), Some(1)));
+        assert!(matches!(buffer.position(&dummy_point(2)), Some(2)));
+
+        assert_eq!(buffer.oldest().unwrap(), &dummy_point(0));
+        assert_eq!(buffer.latest().unwrap(), &dummy_point(2));
+    }
+
+    #[test]
+    fn pop_from_valid_depth_works() {
+        let mut buffer = build_filled_buffer(5);
+
+        let ready = buffer.pop_with_depth(2);
+
+        assert_eq!(dummy_point(0), ready[0]);
+        assert_eq!(dummy_point(1), ready[1]);
+        assert_eq!(dummy_point(2), ready[2]);
+
+        assert_eq!(ready.len(), 3);
+
+        assert_eq!(buffer.oldest().unwrap(), &dummy_point(3));
+        assert_eq!(buffer.latest().unwrap(), &dummy_point(4));
+    }
+
+    #[test]
+    fn pop_from_excessive_depth_returns_empty() {
+        let mut buffer = build_filled_buffer(6);
+
+        let ready = buffer.pop_with_depth(10);
+
+        assert_eq!(ready.len(), 0);
+
+        assert_eq!(buffer.oldest().unwrap(), &dummy_point(0));
+        assert_eq!(buffer.latest().unwrap(), &dummy_point(5));
+    }
+
+    #[test]
+    fn roll_back_within_scope_works() {
+        let mut buffer = build_filled_buffer(6);
+
+        let result = buffer.roll_back(&dummy_point(2));
+
+        assert!(matches!(result, RollbackEffect::Handled));
+
+        assert_eq!(buffer.size(), 3);
+        assert_eq!(buffer.oldest().unwrap(), &dummy_point(0));
+        assert_eq!(buffer.latest().unwrap(), &dummy_point(2));
+
+        let remaining = buffer.pop_with_depth(0);
+
+        assert_eq!(dummy_point(0), remaining[0]);
+        assert_eq!(dummy_point(1), remaining[1]);
+        assert_eq!(dummy_point(2), remaining[2]);
+
+        assert_eq!(remaining.len(), 3);
+    }
+
+    #[test]
+    fn roll_back_outside_scope_works() {
+        let mut buffer = build_filled_buffer(6);
+
+        let result = buffer.roll_back(&dummy_point(100));
+
+        assert!(matches!(result, RollbackEffect::OutOfScope));
+
+        assert_eq!(buffer.size(), 0);
+        assert_eq!(buffer.oldest(), None);
+        assert_eq!(buffer.latest(), None);
+    }
+}

pallas-miniprotocols/src/chainsync/mod.rs

@@ -1,7 +1,9 @@
+mod buffer;
 mod clients;
 mod codec;
 mod protocol;
 
+pub use buffer::*;
 pub use clients::*;
 pub use codec::*;
 pub use protocol::*;

pallas-miniprotocols/src/chainsync/protocol.rs

@@ -2,7 +2,7 @@ use std::{fmt::Debug, ops::Deref};
 
 use crate::common::Point;
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Tip(pub Point, pub u64);
 
 #[derive(Debug, PartialEq, Clone)]

pallas-miniprotocols/src/common.rs

@@ -7,7 +7,7 @@ pub const TESTNET_MAGIC: u64 = 1097911063;
 pub const MAINNET_MAGIC: u64 = 764824073;
 
 /// A point within a chain
-#[derive(Clone)]
+#[derive(Clone, Eq, PartialEq)]
 pub struct Point(pub u64, pub Vec<u8>);
 
 impl Debug for Point {
@@ -18,3 +18,9 @@ impl Debug for Point {
             .finish()
     }
 }
+
+impl Point {
+    pub fn new(slot: u64, hash: Vec<u8>) -> Self {
+        Point(slot, hash)
+    }
+}