hlc_gen/

lib.rs

#![doc = include_str!("../README.md")]

mod epoch;
pub mod error;
mod timestamp;

pub use timestamp::HlcTimestamp;
use {
    crate::{
        epoch::CUSTOM_EPOCH,
        error::{HlcError, HlcResult},
    },
    chrono::Utc,
    parking_lot::RwLock,
    std::cmp::Ordering,
};

/// Provides the current timestamp in milliseconds since the Unix epoch.
pub trait CurrentTimestamp: Default {
    /// Returns the current timestamp in milliseconds since the Unix epoch.
    fn current_timestamp(&self) -> i64;

    /// Sets the current timestamp in milliseconds since the Unix epoch.
    fn set_current_timestamp(&self, timestamp: i64);
}

/// Implementation of the `CurrentTimestamp` trait using UTC.
#[derive(Default)]
pub struct UtcTimestamp;

impl CurrentTimestamp for UtcTimestamp {
    fn current_timestamp(&self) -> i64 {
        Utc::now().timestamp_millis()
    }

    fn set_current_timestamp(&self, _timestamp: i64) {
        unimplemented!("Setting current timestamp is not supported for UtcTimestamp");
    }
}

/// Implementation of the `CurrentTimestamp` trait using a manual timestamp.
///
/// Useful for testing purposes.
pub struct ManualTimestamp {
    /// The current timestamp in milliseconds since the Unix epoch.
    timestamp: RwLock<i64>,
}

impl Default for ManualTimestamp {
    fn default() -> Self {
        Self::new(CUSTOM_EPOCH)
    }
}

impl CurrentTimestamp for ManualTimestamp {
    fn current_timestamp(&self) -> i64 {
        let r = self.timestamp.read();
        *r
    }

    fn set_current_timestamp(&self, timestamp: i64) {
        let mut w = self.timestamp.write();
        *w = timestamp;
    }
}

impl ManualTimestamp {
    /// Creates a new `ManualTimestamp` with the specified timestamp.
    pub fn new(timestamp: i64) -> Self {
        Self {
            timestamp: RwLock::new(timestamp),
        }
    }
}

/// Hybrid Logical Clock (HLC) generator.
pub struct HlcGenerator<T: CurrentTimestamp = UtcTimestamp> {
    /// The last timestamp generated by the clock.
    state: HlcTimestamp,

    /// The maximum drift (in milliseconds) allowed between the physical clock
    /// and the wall-clock time.
    max_drift: usize,

    /// The timestamp provider used to get the current timestamp.
    ts_provider: T,
}

impl Default for HlcGenerator {
    fn default() -> Self {
        Self::new()
    }
}

impl<T: CurrentTimestamp> HlcGenerator<T> {
    /// Creates a new HLC clock without any drift.
    ///
    /// Whenever the clock running on a single node is required, there is no
    /// need to worry about drift, as no adjustments are made to the clock,
    /// i.e. [`update()`](HlcGenerator::update) is never called.
    pub fn new() -> Self {
        Self::with_max_drift(0)
    }

    /// Creates a new HLC clock with the specified maximum drift.
    pub fn with_max_drift(max_drift: usize) -> Self {
        let ts_provider = T::default();
        let state =
            HlcTimestamp::from_parts(ts_provider.current_timestamp(), 0).unwrap_or_else(|_| {
                HlcTimestamp::from_parts(CUSTOM_EPOCH, 0).expect("Invalid timestamp")
            });
        Self {
            state,
            max_drift,
            ts_provider,
        }
    }

    /// Get timestamp provider.
    ///
    /// Useful for testing purposes, where manual timestamps are used.
    pub fn ts_provider(&self) -> &T {
        &self.ts_provider
    }

    /// Current timestamp.
    ///
    /// Use [`next_timestamp()`](HlcGenerator::next_timestamp) to get the
    /// timestamp for local or send events.
    pub fn timestamp(&self) -> HlcTimestamp {
        self.state.clone()
    }

    /// Timestamp for the local or send event.
    pub fn next_timestamp(&self) -> Option<HlcTimestamp> {
        let timestamp = self.ts_provider.current_timestamp();

        self.state
            .update(move |pt, lc| {
                // Update the physical time and increment the logical count.
                if pt >= timestamp {
                    Ok((pt, lc + 1))
                } else {
                    Ok((timestamp, 0))
                }
            })
            .ok()
    }

    /// Adjust the clock based on incoming timestamp.
    ///
    /// Usually this happens when a timestamp is received from another node.
    /// An error may occur if drift is exceeded (if `max_drift` is set to 0,
    /// then such a check is ignored).
    ///
    /// Updated timestamp is returned.
    pub fn update(&self, incoming_state: &HlcTimestamp) -> HlcResult<HlcTimestamp> {
        let max_drift = self.max_drift;
        let timestamp = self.ts_provider.current_timestamp();

        self.state.update(move |pt, lc| {
            let (incoming_pt, incoming_lc) = incoming_state.parts();

            // Physical clock is ahead of both the incoming timestamp and the current state.
            if timestamp > incoming_pt && timestamp > pt {
                // Update the clock state.
                return Ok((timestamp, 0));
            }

            match incoming_pt.cmp(&pt) {
                // Incoming timestamp is ahead of the current state.
                Ordering::Greater => {
                    // Check for drift.
                    let drift = (incoming_pt - timestamp) as usize;
                    if max_drift > 0 && drift > max_drift {
                        Err(HlcError::DriftTooLarge(drift, max_drift))
                    } else {
                        // Remote timestamp is ahead of the current state. Update local state.
                        Ok((incoming_pt, incoming_lc + 1))
                    }
                }
                // Incoming timestamp is behind the current state.
                Ordering::Less => {
                    // Our timestamp is ahead of the incoming timestamp, so it remains unchanged.
                    // We only need to update the logical count.
                    Ok((pt, lc + 1))
                }
                // Timestamps are equal, so we need to use the maximum logical count for update.
                Ordering::Equal => {
                    // Timestamps are equal, so we need to use the maximum logical count for update.
                    Ok((pt, lc.max(incoming_lc) + 1))
                }
            }
        })
    }
}