time_key_stream_set/

stream_set.rs

use std::{
    cell::RefCell,
    cmp::min,
    collections::{BTreeMap, BTreeSet, HashSet},
    fmt::{Debug, Formatter},
    mem::{replace, size_of, swap, take},
    num::NonZeroUsize,
    path::{Path, PathBuf},
    rc::Rc,
    sync::atomic::{AtomicBool, AtomicIsize, AtomicUsize, Ordering},
    time::SystemTime,
};

use async_compression::tokio::{bufread::ZstdDecoder, write::ZstdEncoder};
use futures::{stream, StreamExt};
use lru::LruCache;
use tokio::{
    fs::{create_dir_all, remove_file, File, OpenOptions},
    io::{AsyncReadExt, AsyncWriteExt, BufReader},
    sync::Mutex,
    task::spawn_blocking,
};
use tsz_compress::prelude::*;

use crate::{prelude::*, time_key::TimeKeyItem};

static HOT_SEGMENTS: AtomicIsize = AtomicIsize::new(0);
struct DeferHotSegment(u64);
impl Drop for DeferHotSegment {
    fn drop(&mut self) {
        let hot_segments = HOT_SEGMENTS.fetch_sub(1, Ordering::Relaxed);
        println!("Dehydrating segment ({}): {:?}", hot_segments, self.0);
    }
}

///
/// Implement the stream set interface
///
unsafe impl Send for TkStreamSet {}
unsafe impl Sync for TkStreamSet {}
impl TkStreamSet {
    pub(crate) async fn initialize(&mut self) -> TkssResult<()> {
        // Confirm the working directory exists
        if !self.working_dir.exists() {
            create_dir_all(self.working_dir.clone())
                .await
                .map_err(|_| TkssError::InvalidDirectory(self.working_dir.clone()))?;
        }

        println!("Initializing stream set in {}", self.working_dir.display());

        self.state.segment_interval_us = self.segment_time_interval.as_micros() as u64;
        self.state.retention_us = self.segment_retention_period.as_micros() as u64;
        self.state.max_memory_usage = match self.memory_limit {
            MemoryLimit::VeryLow => 1 * 1024 * 1024 * 1024,
            MemoryLimit::Low => 8 * 1024 * 1024 * 1024,
            MemoryLimit::Medium => 32 * 1024 * 1024 * 1024,
            MemoryLimit::High => 128 * 1024 * 1024 * 1024,
            MemoryLimit::VeryHigh => 1024 * 1024 * 1024 * 1024,
            MemoryLimit::Unlimited => usize::MAX,
        };

        let max_segments = self.state.max_memory_usage as f64 * 0.8 / (256 * 1024 * 1024) as f64;
        let max_segments = min(1024, max_segments as usize);
        println!("Using {} segments in Lru Cache", max_segments);
        self.state.segment_cache =
            Mutex::new(LruCache::new(NonZeroUsize::new(max_segments).unwrap()));

        Ok(())
    }

    pub async fn insert<I, T>(&self, keys: I) -> TkssResult<Vec<bool>>
    where
        I: Iterator<Item = T>,
        T: TimeKey,
    {
        // Find the segments that need to be loaded into memory
        let prune_segment = (SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_micros() as u64
            - self.state.retention_us)
            / self.state.segment_interval_us
            * self.state.segment_interval_us;
        let mut segmented_keys = keys
            .map(|k| {
                (
                    k,
                    k.timestamp_us() / self.state.segment_interval_us
                        * self.state.segment_interval_us,
                )
            })
            .collect::<Vec<_>>();
        segmented_keys.sort_by_key(|(_, s)| *s);
        if segmented_keys.is_empty() {
            return Ok(vec![]);
        }

        // We now hold exclusive access over the segment index for the whole stream set
        let mut segment_idx = self.state.segment_idx.lock().await;

        // Split off the old segments that are no longer needed
        // Put the segments that we actually want back in the index
        let mut prune = segment_idx.split_off(&prune_segment);
        swap(&mut *segment_idx, &mut prune);
        stream::iter(prune.into_iter())
            .for_each_concurrent(num_cpus::get(), |(_, v)| async move {
                // If deletion fails, we just ignore it
                let _ = v
                    .as_ref()
                    .borrow()
                    .delete_segment()
                    .await
                    .and_then(|bytes| {
                        self.state
                            .estimated_memory_used
                            .fetch_sub(bytes, std::sync::atomic::Ordering::Relaxed);
                        Ok(())
                    });
            })
            .await;

        // These are the segments that need to be in memory right now
        let first_segment = segmented_keys[0].1;
        let last_segment = segmented_keys.last().unwrap().1;
        let segments = segment_idx
            .range(first_segment..=last_segment)
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect::<Vec<_>>();

        // Create the segments if they do not exist
        let mut new_segments = false;
        let new_segment_set = segmented_keys
            .iter()
            .filter(|(_, s)| {
                *s >= prune_segment
                    && !segments
                        .iter()
                        .any(|seg| seg.1.borrow().spans_key(*s, self.state.segment_interval_us))
            })
            .map(|(_, s)| *s)
            .collect::<HashSet<_>>();

        // Inserting new segments may evict segments from the LRU cache that are still in use ... (thrashing)
        let mut cache = self.state.segment_cache.lock().await;
        let uniq_keys = segmented_keys
            .iter()
            .map(|(_, s)| s)
            .collect::<HashSet<_>>();
        uniq_keys.iter().for_each(|s| cache.promote(s));
        for s in new_segment_set {
            new_segments = true;
            let segment = Rc::new(RefCell::new(Segment::new(false, &self.working_dir, s)));
            segment_idx.insert(s, segment.clone());
            let stale_segment = cache.push(s, segment);

            if let Some((stale_k, stale_segment)) = stale_segment {
                println!("New segment, Evicting segment {} from memory", stale_k);
                debug_assert!(stale_segment.borrow().hot.load(Ordering::SeqCst));
                if uniq_keys.contains(&stale_k) {
                    tracing::error!("Warning segment cache too small ({}), cannot evict required segment, {} of {:?}", cache.len(), stale_k, uniq_keys );
                    println!("Warning segment cache too small ({}), cannot evict required segment, {} of {:?}", cache.len(), stale_k, uniq_keys );
                    let usage = stale_segment.as_ref().borrow().dehydrate_segment().await?;
                    self.state
                        .estimated_memory_used
                        .fetch_sub(usage, std::sync::atomic::Ordering::Relaxed);
                }
            }
        }

        // Find and lock the segments again, if we created new ones
        let segments = if !new_segments {
            segments
        } else {
            segment_idx
                .range(first_segment..=last_segment)
                .map(|(k, v)| (k.clone(), v.clone()))
                .collect::<Vec<_>>()
        };

        // Bring the segments into memory and flush stale segments to disk if we are over the memory limit
        let required_segments = segments.iter().map(|(k, _)| *k).collect::<Vec<_>>();
        let mut cannot_evict_more = false;
        for (_k, s) in segments.iter() {
            // Evict data from memory if we cannot fit the new data in memory
            if cannot_evict_more {
                break;
            }
            // println!(
            //     "Estimated memory usage {} > {}?",
            //     self.state.estimated_memory_used.load(Ordering::Relaxed),
            //     self.state.max_memory_usage
            // );
            while self.state.estimated_memory_used.load(Ordering::Relaxed)
                > self.state.max_memory_usage
                && !cache.is_empty()
            {
                println!(
                    "High memory pressure {} > {}, Evicting segment from memory",
                    self.state.estimated_memory_used.load(Ordering::Relaxed),
                    self.state.max_memory_usage
                );
                // Every item in this insertion was already promoted
                let stale_segment = cache.pop_lru();
                if let Some((stale_k, stale_segment)) = stale_segment {
                    if required_segments.contains(&stale_k) {
                        println!("Cannot evict required segment");
                        cache.push(stale_k, stale_segment);
                        cannot_evict_more = true;
                        break;
                    } else {
                        debug_assert!(stale_segment.borrow().hot.load(Ordering::SeqCst));
                        let usage = stale_segment.as_ref().borrow().dehydrate_segment().await?;
                        self.state
                            .estimated_memory_used
                            .fetch_sub(usage, std::sync::atomic::Ordering::Relaxed);
                    }
                }
            }

            // Load the segment into memory if it is not already
            let seg = s.as_ref().borrow();
            if !seg.hot() {
                let usage = seg.hydrate_segment().await?;
                self.state
                    .estimated_memory_used
                    .fetch_add(usage, std::sync::atomic::Ordering::Relaxed);
            }
        }

        // We now need to hold exclusive access over the segments that are required for this insert
        let mut _guards = Vec::with_capacity(segments.len());
        for (k, s) in segments.iter() {
            let b = s.as_ref().borrow();
            _guards.push((*k, b));
        }
        let mut guards = stream::iter(_guards.iter())
            .map(|(k, b)| async move { (*k, b.idx.lock().await) })
            .buffer_unordered(num_cpus::get())
            .collect::<Vec<_>>()
            .await;

        // We no longer need to hold exclusive access over the cache and the segment index
        drop(cache);
        drop(segment_idx);

        // Insert the keys into the segments
        let deduped = segmented_keys
            .into_iter()
            .map(|(k, s)| {
                if s < prune_segment {
                    return false;
                }
                let Some(idx) = guards.iter_mut().find(|seg| seg.0 == s) else {
                    tracing::error!(
                        "Pseudo-panic: Expected segment {} in {:?}",
                        s,
                        guards.iter().map(|(k, _)| k).collect::<Vec<_>>()
                    );
                    println!(
                        "Pseudo-panic: Expected segment {} in {:?}",
                        s,
                        guards.iter().map(|(k, _)| k).collect::<Vec<_>>()
                    );
                    return false;
                };
                idx.1.insert(k.key())
            })
            .collect::<Vec<_>>();

        // Dehydrate any evicted segments not yet dehydrated
        drop(guards);
        drop(_guards);

        // Update memory usage
        let new_keys = deduped.iter().filter(|b| **b).count();
        self.state.estimated_memory_used.fetch_add(
            (new_keys as f32 * size_of::<u128>() as f32 * 2.1) as usize,
            Ordering::Relaxed,
        );

        Ok(deduped)
    }
}

pub(crate) struct TkStreamSetState {
    ///
    /// Cached interval for bucketing segments
    ///
    segment_interval_us: u64,

    ///
    /// Cached interval for deleting segments
    /// Segments older than (now - retention_us) will be deleted
    ///
    retention_us: u64,

    ///
    /// Cached maximum memory usage in bytes
    ///
    max_memory_usage: usize,

    ///
    /// An ongoing estimate of the amount of memory used by the stream set
    /// in bytes.
    ///
    estimated_memory_used: AtomicUsize,

    ///
    /// This keeps an ordered map of the time-keyed segments based on
    /// the first timestamp in the segment.
    ///
    segment_idx: Mutex<BTreeMap<u64, Rc<RefCell<Segment>>>>,

    ///
    /// This keeps an LRU cache for segments in memory
    ///
    segment_cache: Mutex<LruCache<u64, Rc<RefCell<Segment>>>>,
}

impl Default for TkStreamSetState {
    fn default() -> Self {
        TkStreamSetState {
            segment_interval_us: 0,
            retention_us: 0,
            max_memory_usage: 0,
            estimated_memory_used: 0.into(),
            segment_idx: Mutex::new(BTreeMap::new()),
            segment_cache: Mutex::new(LruCache::new(NonZeroUsize::new(16).unwrap())),
        }
    }
}

impl Debug for TkStreamSetState {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TkStreamSetState")
            .field("segment_interval_us", &self.segment_interval_us)
            .finish()
    }
}

///
/// A `Segment` represents the smallest unit of data that is read/written to disk.
/// They represent a unique interval of time, and are bucketed by the first timestamp.
///
/// It is recommended for efficiency that the segment interval is sized such that
/// segments are not too large, but also not too small. A good rule of thumb is to
/// target a segment size of 256MiB or less for optimal log2(n) insertion performance
/// and reduced impact of single-threaded tsz_compress compaction.
///
/// The `hot` flag indicates whether the segment is currently in memory or not. Where
/// segments are held in an LRU cache sized based on the `max_memory_usage` parameter
/// and the `segment_interval_us` parameter to hold ~256MiB of segments in memory.
///
#[allow(unused)]
pub(crate) struct Segment {
    halt: AtomicBool,
    hot: AtomicBool,
    working_dir: PathBuf,
    timestamp_us: u64,
    idx: Mutex<BTreeSet<TimeKeyItem>>,
}

impl Segment {
    pub fn new(halt: bool, working_dir: &Path, timestamp_us: u64) -> Self {
        Segment {
            halt: halt.into(),
            hot: false.into(),
            working_dir: working_dir.to_path_buf(),
            timestamp_us,
            idx: Mutex::new(BTreeSet::new()),
        }
    }

    pub fn spans_key(&self, timestamp_us: u64, interval_us: u64) -> bool {
        self.timestamp_us <= timestamp_us && timestamp_us < self.timestamp_us + interval_us
    }

    pub fn hot(&self) -> bool {
        self.hot.load(Ordering::Relaxed)
    }

    fn path(&self) -> PathBuf {
        self.working_dir.join(format!("{}.bin", self.timestamp_us))
    }

    #[tracing::instrument(skip(self), ret)]
    pub async fn hydrate_segment(&self) -> TkssResult<usize> {
        let mut idx = self.idx.lock().await;
        let hot = self.hot.swap(true, Ordering::Relaxed);
        if hot {
            // No change in size, if already hot
            return Ok(0);
        }

        let hot_segments = HOT_SEGMENTS.fetch_add(1, Ordering::Relaxed);

        println!(
            "Hydrating segment ({}): {:?}",
            hot_segments, self.timestamp_us
        );

        // If the path doesn't exist, it wasn't written to disk
        let path = self.path();
        if !path.exists() {
            return Ok(0);
        }

        // Read all the bits from disk into a Vec<u8>
        let file = BufReader::new(File::open(&path).await?);
        let mut zfile = ZstdDecoder::new(file);
        let mut compressed = Vec::new();
        zfile.read_to_end(&mut compressed).await?;
        let bytes = compressed.len();

        // Expect the header to be "TKSS"
        let header = "TKSS".as_bytes();
        if &compressed[0..header.len()] != header {
            Err(TkssError::InvalidHeader(path.clone()))?;
        }
        let compressed = &compressed[header.len()..];

        // The next 8 bytes are the u64 bit length of the compressed data
        let length = u64::from_be_bytes(compressed[..8].try_into()?);
        if (length + 7) as isize / 8 < compressed.len() as isize - (header.len() + 8) as isize {
            println!("Length: {} ({} bytes)", length, (length + 7) as usize / 8);
            Err(TkssError::InvalidLength(path.clone()))?;
        }
        let compressed = &compressed[8..];
        let bits = BitBufferSlice::from_slice(compressed);
        let bits = bits.split_at(length as usize).0;
        let mut decompressor = Decompressor::new(bits);

        println!(
            "Read segment: {:?} ({} + 8 ({}) + {} = {} bytes)",
            self.timestamp_us,
            header.len(),
            length,
            compressed.len(),
            bytes
        );

        // Decompress the data directly into the index
        idx.extend(
            decompressor
                .decompress()
                .filter_map(|e: Result<TimeKeyItem, _>| match e {
                    Ok(e) => Some(e),
                    Err(e) => {
                        tracing::error!("Error decompressing segment: {:?}", e);
                        None
                    }
                }),
        );

        // Return the estimated size of the segment based on the number of items
        let size_estimate = (idx.len() as f64 * size_of::<u128>() as f64 * 2.1) as usize;
        println!("Estimated size: {}", size_estimate);
        Ok(size_estimate)
    }

    #[tracing::instrument(skip(self), ret)]
    pub async fn dehydrate_segment(&self) -> TkssResult<usize> {
        let mut idx = self.idx.lock().await;
        let hot = self.hot.swap(false, Ordering::Relaxed);
        if !hot {
            // No change in size, if already cold
            return Ok(0);
        }

        let _hot_segment = DeferHotSegment(self.timestamp_us);

        // Compress the data
        let pts = idx.len();
        let size_estimate = (idx.len() as f64 * size_of::<u128>() as f64 * 2.1) as usize;
        let elems = take(&mut *idx);
        let compressed = spawn_blocking(move || {
            let mut compressor = Compressor::new();
            elems.into_iter().for_each(|e| compressor.compress(e));
            let mut compressed = compressor.finish();
            compressed.shrink_to_fit();
            compressed.set_uninitialized(false);
            compressed
        })
        .await?;

        // Write the data to disk
        let header = "TKSS".as_bytes();
        let length = compressed.len() as u64;
        let compressed = compressed.into_vec();
        let file = OpenOptions::new()
            .write(true)
            .create(true)
            .open(&self.path())
            .await?;
        let mut zfile = ZstdEncoder::new(file);

        // Write the header, the bits length as a u64, and then the bits
        zfile.write_all(header).await?;
        zfile.write_u64(length).await?;
        zfile.write_all(&compressed).await?;

        // Flush the file and drop it, print the file size
        zfile.shutdown().await?;
        let bytes = zfile.into_inner().metadata().await?.len();
        println!(
            "Wrote segment: {:?} ({} + 8 ({}) + {} = {} bytes / {} bytes) ({:.2}x compression)",
            self.timestamp_us,
            header.len(),
            length,
            compressed.len(),
            header.len() + 8 + compressed.len(),
            bytes,
            (pts as f64 * size_of::<u128>() as f64) / (bytes as f64)
        );

        // Return the size of the data released
        Ok(size_estimate)
    }

    #[tracing::instrument(skip(self), ret)]
    pub async fn delete_segment(&self) -> TkssResult<usize> {
        let mut idx = self.idx.lock().await;
        let hot = self.hot.swap(false, Ordering::Relaxed);
        self.halt.store(true, Ordering::Relaxed);
        let pts = if !hot {
            0
        } else {
            HOT_SEGMENTS.fetch_sub(1, Ordering::Relaxed);
            idx.len()
        };

        // Count how many samples are getting deleted
        let size_estimate = (pts as f64 * size_of::<u128>() as f64 * 2.1) as usize;
        idx.clear();

        // Delete the file
        let path = self.path();
        if path.exists() {
            println!("Deleting segment: {:?}", self.timestamp_us);
            remove_file(&path).await?;
        }

        // Return the size of the data released
        Ok(size_estimate)
    }
}

///
/// Upon drop, take the index and current state out of this segment
/// and spawn a task to write it to disk, leaving an empty segment in place.
///
impl Drop for Segment {
    fn drop(&mut self) {
        println!("Dropping segment: {:?}", self.timestamp_us);
        if self.halt.swap(true, Ordering::SeqCst) {
            return;
        }
        let this = replace(
            self,
            Segment::new(true, &self.working_dir, self.timestamp_us),
        );
        tokio::spawn(async move {
            println!("Dehydrating segment on drop: {:?}", this.timestamp_us);
            let _ = this.dehydrate_segment().await;
        });
    }
}

#[cfg(test)]
mod tests {
    use std::{
        sync::atomic::Ordering,
        time::{Duration, Instant, SystemTime},
    };

    use futures::{stream, StreamExt};
    use rand::{seq::SliceRandom, thread_rng};
    use tempdir::TempDir;
    use tokio::{fs::read_dir, time::sleep};

    use crate::{prelude::*, stream_set::HOT_SEGMENTS};

    #[tokio::test]
    async fn can_insert_keys() {
        let sset = TkStreamSetBuilder::new().build().await.unwrap();
        let now_ts = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_micros() as i64;
        let keys = (0..10).map(|i| UserDeviceTimeKey {
            timestamp_us: now_ts + i * 1000,
            user_id: 42,
            device_id: 69,
        });
        let results = sset.insert(keys.clone()).await.unwrap();
        assert_eq!(results, vec![true; 10]);

        let results = sset.insert(keys).await.unwrap();
        assert_eq!(results, vec![false; 10]);
    }

    #[tokio::test]
    async fn can_hydrate_drop_rehydrate() {
        // Configure a temporary directory that will be deleted when the test is done
        let tmp_dir = TempDir::new("can_hydrate_drop_rehydrate").unwrap();
        let tmp_path = tmp_dir.into_path();

        let sset = TkStreamSetBuilder::new()
            .with_working_dir(tmp_path.clone())
            .build()
            .await
            .unwrap();
        let now_ts = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_micros() as i64;
        let keys = (0..50).map(|i| UserDeviceTimeKey {
            timestamp_us: now_ts + i * 1000,
            user_id: 42,
            device_id: 69,
        });
        let results = sset.insert(keys.clone()).await.unwrap();
        assert_eq!(results, vec![true; 50]);

        let results = sset.insert(keys.clone()).await.unwrap();
        assert_eq!(results, vec![false; 50]);

        drop(sset);

        let ts = Instant::now();
        while HOT_SEGMENTS.load(Ordering::Relaxed) > 0 && ts.elapsed() < Duration::from_secs(15) {
            sleep(Duration::from_millis(10)).await;
        }

        // assert that there is a file in the directory
        let mut files = read_dir(tmp_path.clone()).await.unwrap();
        let has_files = files.next_entry().await.unwrap().is_some();
        assert!(has_files);

        let sset = TkStreamSetBuilder::new()
            .with_working_dir(tmp_path.clone())
            .build()
            .await
            .unwrap();
        let results = sset.insert(keys).await.unwrap();
        assert_eq!(results, vec![false; 50]);
    }

    #[tokio::test]
    async fn can_hydrate_drop_rehydrate_across_segments() {
        // Configure a temporary directory that will be deleted when the test is done
        let tmp_dir = TempDir::new("can_hydrate_drop_rehydrate_across_segments").unwrap();
        let tmp_path = tmp_dir.into_path();

        let sset = TkStreamSetBuilder::new()
            .with_working_dir(tmp_path.clone())
            .with_segment_time_interval(Duration::from_secs(300))
            .build()
            .await
            .unwrap();
        let now_ts = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_micros() as i64;

        const NUM_BATCHES: usize = 3600; // 1 hour
        const NUM_BATCH_ROWS: usize = 200; // 1 second each

        let start = Instant::now();
        let sset_ref = &sset;
        stream::iter(0..NUM_BATCHES)
            .for_each_concurrent(num_cpus::get(), |j| async move {
                if j % 100 == 0 {
                    println!("Inserting batch {}", j);
                }
                let keys = (0..NUM_BATCH_ROWS).map(move |i| UserDeviceTimeKey {
                    timestamp_us: now_ts + (j * NUM_BATCH_ROWS + i) as i64 * 5000,
                    user_id: 42 + i as i32 % 7,
                    device_id: 69 + i as i32 % 11,
                });
                let results = sset_ref.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![true; NUM_BATCH_ROWS]);

                let results = sset_ref.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![false; NUM_BATCH_ROWS]);
            })
            .await;
        println!(
            "Inserting {} keys took {:?}",
            NUM_BATCHES * NUM_BATCH_ROWS,
            start.elapsed()
        );

        drop(sset);

        let ts = Instant::now();
        while HOT_SEGMENTS.load(Ordering::Relaxed) > 0 && ts.elapsed() < Duration::from_secs(15) {
            sleep(Duration::from_millis(10)).await;
        }

        // assert that there is a file in the directory
        let mut files = read_dir(tmp_path.clone()).await.unwrap();
        let mut num_files = 0;
        while let Ok(Some(_file)) = files.next_entry().await {
            num_files += 1;
        }
        println!("Found {} files", num_files);
        assert!(num_files > 1);

        let sset = TkStreamSetBuilder::new()
            .with_working_dir(tmp_path.clone())
            .with_segment_time_interval(Duration::from_secs(300))
            .build()
            .await
            .unwrap();

        let sset = &sset;
        stream::iter(0..NUM_BATCHES)
            .for_each_concurrent(num_cpus::get(), |j| async move {
                let keys = (0..NUM_BATCH_ROWS).map(|i| UserDeviceTimeKey {
                    timestamp_us: now_ts + (j * NUM_BATCH_ROWS + i) as i64 * 5000,
                    user_id: 42 + i as i32 % 7,
                    device_id: 69 + i as i32 % 11,
                });

                let results = sset.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![false; NUM_BATCH_ROWS]);
            })
            .await
    }

    #[tokio::test]
    async fn can_hydrate_drop_rehydrate_across_segments_random_insertion() {
        // Configure a temporary directory that will be deleted when the test is done
        let tmp_dir =
            TempDir::new("can_hydrate_drop_rehydrate_across_segments_random_insertion").unwrap();
        let tmp_path = tmp_dir.into_path();

        let sset = TkStreamSetBuilder::new()
            .with_working_dir(tmp_path.clone())
            .with_segment_time_interval(Duration::from_secs(300))
            .build()
            .await
            .unwrap();
        let now_ts = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_micros() as i64;

        const NUM_BATCHES: usize = 2 * 3600;
        const NUM_BATCH_ROWS: usize = 100;
        let mut batch_order = (0..NUM_BATCHES).collect::<Vec<_>>();
        batch_order.shuffle(&mut thread_rng());

        let start = Instant::now();
        let sset_ref = &sset;
        stream::iter(batch_order.iter())
            .for_each_concurrent(num_cpus::get(), |j| async move {
                if j % 100 == 0 {
                    println!("Inserting batch {}", j);
                }
                let keys = (0..NUM_BATCH_ROWS).map(move |i| UserDeviceTimeKey {
                    timestamp_us: now_ts + (j * NUM_BATCH_ROWS + i) as i64 * 5000,
                    user_id: 42 + i as i32 % 7,
                    device_id: 69 + i as i32 % 11,
                });
                let results = sset_ref.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![true; NUM_BATCH_ROWS]);

                let results = sset_ref.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![false; NUM_BATCH_ROWS]);
            })
            .await;
        println!(
            "Inserting {} keys took {:?}",
            NUM_BATCHES * NUM_BATCH_ROWS,
            start.elapsed()
        );

        drop(sset);

        let ts = Instant::now();
        while HOT_SEGMENTS.load(Ordering::Relaxed) > 0 && ts.elapsed() < Duration::from_secs(15) {
            sleep(Duration::from_millis(10)).await;
        }

        // assert that there is a file in the directory
        let mut files = read_dir(tmp_path.clone()).await.unwrap();
        let mut num_files = 0;
        while let Ok(Some(_file)) = files.next_entry().await {
            num_files += 1;
        }
        println!("Found {} files", num_files);
        assert!(num_files > 1);

        let sset = TkStreamSetBuilder::new()
            .with_working_dir(tmp_path.clone())
            .with_segment_time_interval(Duration::from_secs(300))
            .build()
            .await
            .unwrap();

        let sset = &sset;
        stream::iter(batch_order.iter())
            .for_each_concurrent(num_cpus::get(), |j| async move {
                let keys = (0..NUM_BATCH_ROWS).map(|i| UserDeviceTimeKey {
                    timestamp_us: now_ts + (j * NUM_BATCH_ROWS + i) as i64 * 5000,
                    user_id: 42 + i as i32 % 7,
                    device_id: 69 + i as i32 % 11,
                });

                let results = sset.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![false; NUM_BATCH_ROWS]);
            })
            .await
    }

    #[tokio::test]
    async fn can_hydrate_drop_rehydrate_new_segments_then_old_segments() {
        // Configure a temporary directory that will be deleted when the test is done
        let tmp_dir =
            TempDir::new("can_hydrate_drop_rehydrate_new_segments_then_old_segments").unwrap();
        let tmp_path = tmp_dir.into_path();

        let sset = TkStreamSetBuilder::new()
            .with_working_dir(tmp_path.clone())
            .with_segment_time_interval(Duration::from_secs(300))
            .with_memory_limit(MemoryLimit::Low)
            .build()
            .await
            .unwrap();
        let now_ts = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_micros() as i64;

        const NUM_BATCHES: usize = 24 * 3600;
        const NUM_BATCH_ROWS: usize = 100;
        let batch_order = (0..NUM_BATCHES).collect::<Vec<_>>();

        let start = Instant::now();
        let sset_ref = &sset;
        stream::iter(batch_order[(NUM_BATCHES / 2)..].iter())
            .for_each_concurrent(num_cpus::get(), |j| async move {
                if j % 100 == 0 {
                    println!("Inserting batch {}", j);
                }
                let keys = (0..NUM_BATCH_ROWS).map(move |i| UserDeviceTimeKey {
                    timestamp_us: now_ts + (j * NUM_BATCH_ROWS + i) as i64 * 5000,
                    user_id: 42 + i as i32 % 7,
                    device_id: 69 + i as i32 % 11,
                });
                let results = sset_ref.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![true; NUM_BATCH_ROWS]);

                let results = sset_ref.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![false; NUM_BATCH_ROWS]);
            })
            .await;
        println!(
            "Inserting {} keys took {:?}",
            NUM_BATCHES * NUM_BATCH_ROWS,
            start.elapsed()
        );
        stream::iter(batch_order[..(NUM_BATCHES / 2)].iter())
            .for_each_concurrent(num_cpus::get(), |j| async move {
                if j % 100 == 0 {
                    println!("Inserting batch {}", j);
                }
                let keys = (0..NUM_BATCH_ROWS).map(move |i| UserDeviceTimeKey {
                    timestamp_us: now_ts + (j * NUM_BATCH_ROWS + i) as i64 * 5000,
                    user_id: 42 + i as i32 % 7,
                    device_id: 69 + i as i32 % 11,
                });
                let results = sset_ref.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![true; NUM_BATCH_ROWS]);

                let results = sset_ref.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![false; NUM_BATCH_ROWS]);
            })
            .await;
        println!(
            "Inserting {} keys took {:?}",
            NUM_BATCHES * NUM_BATCH_ROWS,
            start.elapsed()
        );

        drop(sset);

        let ts = Instant::now();
        while HOT_SEGMENTS.load(Ordering::Relaxed) > 0 && ts.elapsed() < Duration::from_secs(15) {
            sleep(Duration::from_millis(10)).await;
        }

        // assert that there is a file in the directory
        let mut files = read_dir(tmp_path.clone()).await.unwrap();
        let mut num_files = 0;
        while let Ok(Some(_file)) = files.next_entry().await {
            num_files += 1;
        }
        println!("Found {} files", num_files);
        assert!(num_files > 1);

        let sset = TkStreamSetBuilder::new()
            .with_working_dir(tmp_path.clone())
            .with_segment_time_interval(Duration::from_secs(300))
            .build()
            .await
            .unwrap();

        let sset = &sset;
        stream::iter(batch_order.iter())
            .for_each_concurrent(num_cpus::get(), |j| async move {
                let keys = (0..NUM_BATCH_ROWS).map(|i| UserDeviceTimeKey {
                    timestamp_us: now_ts + (j * NUM_BATCH_ROWS + i) as i64 * 5000,
                    user_id: 42 + i as i32 % 7,
                    device_id: 69 + i as i32 % 11,
                });

                let results = sset.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![false; NUM_BATCH_ROWS]);
            })
            .await
    }

    // release build only
    #[cfg(not(debug_assertions))]
    #[tokio::test]
    async fn can_hydrate_drop_rehydrate_across_day_of_segments() {
        // Configure a temporary directory that will be deleted when the test is done
        let tmp_dir = TempDir::new("can_hydrate_drop_rehydrate_across_day_of_segments").unwrap();
        let tmp_path = tmp_dir.into_path();

        let sset = TkStreamSetBuilder::new()
            .with_working_dir(tmp_path.clone())
            .with_segment_time_interval(Duration::from_secs(60 * 60 * 2))
            .with_memory_limit(MemoryLimit::Low)
            .build()
            .await
            .unwrap();
        let now_ts = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_micros() as i64;

        const NUM_BATCHES: usize = 48 * 3600; // 1 day
        const NUM_BATCH_ROWS: usize = 5000; // 1 second each

        let start = Instant::now();
        let sset_ref = &sset;
        stream::iter(0..NUM_BATCHES)
            .for_each_concurrent(num_cpus::get(), |j| async move {
                if j % 100 == 0 {
                    println!("Inserting batch {}", j);
                }
                let keys = (0..NUM_BATCH_ROWS).map(move |i| UserDeviceTimeKey {
                    timestamp_us: now_ts + (j * NUM_BATCH_ROWS + i) as i64 * 1000,
                    user_id: 42 + i as i32 % 7,
                    device_id: 69 + i as i32 % 11,
                });
                let results = sset_ref.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![true; NUM_BATCH_ROWS]);

                let results = sset_ref.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![false; NUM_BATCH_ROWS]);
            })
            .await;
        println!(
            "Inserting {} keys took {:?}",
            NUM_BATCHES * NUM_BATCH_ROWS,
            start.elapsed()
        );

        drop(sset);

        let ts = Instant::now();
        while HOT_SEGMENTS.load(Ordering::Relaxed) > 0 && ts.elapsed() < Duration::from_secs(15) {
            sleep(Duration::from_millis(10)).await;
        }

        // assert that there is a file in the directory
        let mut files = read_dir(tmp_path.clone()).await.unwrap();
        let has_files = files.next_entry().await.unwrap().is_some();
        assert!(has_files);

        let sset = TkStreamSetBuilder::new()
            .with_working_dir(tmp_path.clone())
            .with_segment_time_interval(Duration::from_secs(60 * 60 * 2))
            .with_memory_limit(MemoryLimit::Low)
            .build()
            .await
            .unwrap();

        let sset = &sset;
        stream::iter(0..NUM_BATCHES)
            .for_each_concurrent(num_cpus::get(), |j| async move {
                let keys = (0..NUM_BATCH_ROWS).map(|i| UserDeviceTimeKey {
                    timestamp_us: now_ts + (j * NUM_BATCH_ROWS + i) as i64 * 1000,
                    user_id: 42 + i as i32 % 7,
                    device_id: 69 + i as i32 % 11,
                });

                let results = sset.insert(keys.clone()).await.unwrap();
                assert_eq!(results, vec![false; NUM_BATCH_ROWS]);
            })
            .await
    }
}