time_key_stream_set/

lib.rs

use serde::{Deserialize, Serialize};
use std::{
    fmt::{Display, Formatter},
    path::PathBuf,
    time::Duration,
};
use stream_set::TkStreamSetState;
use tempdir::TempDir;

pub mod stream_set;
pub mod time_key;

pub mod prelude {
    pub use crate::time_key::{TimeKey, TimeKeyItem, UserDeviceTimeKey};
    pub use crate::{MemoryLimit, TkStreamSet, TkStreamSetBuilder, TkssError, TkssResult};
}

///
/// A `TkStreamSet` is a data structure that is primarily backed on disk.
/// The data structure will only be temporarily loaded into memory with
/// pages cached in memory as needed, but limited by the `MemoryLimit` setting.
///
/// The `MemoryLimit` is a soft limit that is attempted to be enforced with best
/// effort but may be exceeded in some cases by a relatively small margin, ~5%.
///

#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub enum MemoryLimit {
    /// 1 GiB
    VeryLow,
    /// 8 GiB
    Low,
    /// 32 GiB
    Medium,
    /// 128 GiB
    High,
    /// 1 TiB
    VeryHigh,
    /// No enforced limit
    Unlimited,
}

///
/// A `TkStreamSet` is the top-level context for
/// a _time key_ stream set used for deduplication of a
/// mostly continuous stream of time-keyed data.
///
#[derive(Debug)]
pub struct TkStreamSet {
    memory_limit: MemoryLimit,
    working_dir: PathBuf,
    segment_time_interval: Duration,
    segment_retention_period: Duration,
    state: TkStreamSetState,
}

///
/// A `TkStreamSetBuilder` is used to create a `TkStreamSet` with clear configuration.
///
#[derive(Default)]
pub struct TkStreamSetBuilder {
    memory_limit: Option<MemoryLimit>,
    working_dir: Option<PathBuf>,
    segment_time_interval: Option<Duration>,
    segment_retention_period: Option<Duration>,
}

impl TkStreamSetBuilder {
    ///
    /// Create a new instance of the builder
    /// with default configuration.
    ///
    pub fn new() -> TkStreamSetBuilder {
        TkStreamSetBuilder::default()
    }

    ///
    /// Create an instance while consuming the builder
    ///
    /// Default the rest of the configuration.
    ///
    pub async fn build(self) -> TkssResult<TkStreamSet> {
        let memory_limit = self.memory_limit.unwrap_or(MemoryLimit::VeryLow);
        let working_dir = self.working_dir.unwrap_or_else(|| {
            let tmp = TempDir::new("tkstreamset").unwrap();
            tmp.into_path()
        });
        let segment_time_interval = self
            .segment_time_interval
            .unwrap_or(Duration::from_secs(60 * 60));
        let segment_retention_period = self
            .segment_retention_period
            .unwrap_or(Duration::from_secs(60 * 60 * 24 * 7));
        let mut sset = TkStreamSet {
            memory_limit,
            working_dir,
            segment_time_interval,
            segment_retention_period,
            state: TkStreamSetState::default(),
        };
        sset.initialize().await?;
        Ok(sset)
    }

    ///
    /// Set the memory limit for the stream set.
    ///
    /// If not set, the limit will be set to `MemoryLimit::VeryLow`.
    /// If set, the limit will be enforced with best effort.
    /// The limit is a soft limit and may be exceeded in some cases
    /// by a relatively small margin, ~5%.
    ///
    pub fn with_memory_limit(self, limit: MemoryLimit) -> TkStreamSetBuilder {
        TkStreamSetBuilder {
            memory_limit: Some(limit),
            ..self
        }
    }

    ///
    /// Set the working directory for the stream set.
    ///
    /// If not set, a temporary directory will be created
    /// and cleaned up on process exit.
    /// If set, the directory must exist and be writable.
    ///
    pub fn with_working_dir(self, dir: PathBuf) -> TkStreamSetBuilder {
        TkStreamSetBuilder {
            working_dir: Some(dir),
            ..self
        }
    }

    ///
    /// Set the segment time interval for the stream set.
    /// This determines the size of the time-keyed segments that need
    /// to be loaded into memory to perform deduplication of a batch insert.
    ///
    /// If not set, the interval will be set to 1 hour.
    /// If set, the interval will be enforced.
    ///
    pub fn with_segment_time_interval(self, interval: Duration) -> TkStreamSetBuilder {
        TkStreamSetBuilder {
            segment_time_interval: Some(interval),
            ..self
        }
    }

    ///
    /// Set the segment retention period for the stream set.
    /// This determines how long time-keyed segments are kept on disk (even if inactive or small).
    /// This period is enforced while inserting data into the stream set.
    ///
    /// If not set, the retention period will be set to 7 days.
    /// If set, the retention period will be enforced.
    ///
    pub fn with_segment_retention_period(self, period: Duration) -> TkStreamSetBuilder {
        TkStreamSetBuilder {
            segment_retention_period: Some(period),
            ..self
        }
    }
}

pub type TkssResult<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;

#[derive(Debug, Clone)]
pub enum TkssError {
    Unknown(String),
    InvalidDirectory(PathBuf),
    InvalidHeader(PathBuf),
    InvalidLength(PathBuf),
}

impl Display for TkssError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            TkssError::Unknown(s) => write!(f, "Unknown error: {}", s),
            TkssError::InvalidDirectory(p) => write!(f, "Invalid directory: {}", p.display()),
            TkssError::InvalidHeader(p) => write!(f, "Invalid header: {}", p.display()),
            TkssError::InvalidLength(p) => write!(f, "Invalid length: {}", p.display()),
        }
    }
}

impl std::error::Error for TkssError {}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn can_build_instance() {
        let t = TkStreamSetBuilder::new().build().await.unwrap();
        assert_eq!(t.memory_limit, MemoryLimit::VeryLow);

        let t = TkStreamSetBuilder::new()
            .with_memory_limit(MemoryLimit::Medium)
            .build()
            .await
            .unwrap();
        assert_eq!(t.memory_limit, MemoryLimit::Medium);

        let t = TkStreamSetBuilder::new()
            .with_memory_limit(MemoryLimit::High)
            .with_working_dir(PathBuf::from("/tmp"))
            .build()
            .await
            .unwrap();
        assert_eq!(t.memory_limit, MemoryLimit::High);
        assert_eq!(t.working_dir, PathBuf::from("/tmp"));
    }
}