hyperfile 0.11.2

#![cfg(not(any(feature = "blocking", feature = "wal")))]
//! Blocks written by reference are sent ahead under the stream format, as dirty
//! blocks are: their bytes are read from the block source when their piece is
//! uploaded, and the flush reads only what was not sent. On memory staging.
//!
//! ```bash
//! cargo test --release --test functional_stream_write_ref
//! ```

use btree_ondisk::{MemoryBlockLoader, NullNodeCache};
use hyperfile::block_source::BlockSource;
use hyperfile::config::{HyperFileConfigBuilder, HyperFileMetaConfig, HyperFileRuntimeConfig};
use hyperfile::file::file::HyperFile;
use hyperfile::file::flags::{FileFlags, HyperFileFlags};
use hyperfile::file::mode::{FileMode, HyperFileMode};
use hyperfile::file::HyperTrait;
use hyperfile::inode::FlushInodeFlag;
use hyperfile::meta_format::BlockPtrFormat;
use hyperfile::staging::config::StagingConfig;
use hyperfile::staging::memory::MemoryStaging;
use hyperfile::staging::{Staging, StagingIntercept};
use hyperfile::SegmentId;
use std::io::{Error, ErrorKind, Result};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};

type MemFile<'a> = HyperFile<'a, MemoryStaging, MemoryBlockLoader<hyperfile::BlockPtr>, NullNodeCache>;

const BS: usize = 4096;
const MIB: usize = 1 << 20;
const FMT: BlockPtrFormat = BlockPtrFormat::StreamSegment;

/// Sending ahead at `target` in 5 MiB pieces. The dirty-data limits stay at
/// their defaults unless `raised`: a reference is not held back by them.
fn runtime(target: usize, raised: bool) -> HyperFileRuntimeConfig {
    let mut rt = HyperFileRuntimeConfig::default();
    if raised {
        rt.data_cache_dirty_max_bytes_threshold = 1 << 40;
        rt.data_cache_dirty_max_blocks_threshold = (1 << 40) / BS;
        rt.segment_buffer_size = 1 << 40;
    }
    rt.data_cache_dirty_max_flush_interval = u64::MAX / 2;
    rt.data_cache_blocks = 0;
    rt.stream_segment_bytes = target;
    rt.segment_mpu_chunk_size = 5 * MIB;
    rt
}

fn config(s: &MemoryStaging, rt: &HyperFileRuntimeConfig) -> hyperfile::config::HyperFileConfig {
    let mut meta = HyperFileMetaConfig::default();
    meta.block_ptr_format = FMT;
    HyperFileConfigBuilder::new().with_staging_config(s.config()).with_meta_config(&meta).with_runtime_config(rt).build()
}

async fn create(name: &str, rt: &HyperFileRuntimeConfig) -> (MemoryStaging, MemFile<'static>) {
    let s = MemoryStaging::new(StagingConfig::new_memory(name), rt.clone());
    let f = MemFile::new(s.clone(), s.to_block_loader(), NullNodeCache, config(&s, rt),
        HyperFileFlags::from_flags(FileFlags::rdwr()), HyperFileMode::from_mode(FileMode::default_file()))
        .await.expect("create");
    (s, f)
}

async fn reopen(s: &MemoryStaging, rt: &HyperFileRuntimeConfig) -> MemFile<'static> {
    MemFile::open(s.clone(), s.to_block_loader(), NullNodeCache, config(s, rt), HyperFileFlags::from_flags(FileFlags::rdwr()))
        .await.expect("open")
}

fn block(b: usize, g: u8) -> Vec<u8> {
    (0..BS).map(|i| (b as u8).wrapping_mul(29) ^ g ^ (i as u8)).collect()
}

/// A source over a buffer that can change, counting the bytes read from it.
#[derive(Clone)]
struct Source {
    data: Arc<Mutex<Vec<u8>>>,
    read: Arc<AtomicUsize>,
}

impl Source {
    fn new(blocks: usize, g: u8) -> Self {
        Self { data: Arc::new(Mutex::new((0..blocks).flat_map(|b| block(b, g)).collect())), read: Arc::new(AtomicUsize::new(0)) }
    }
    fn set(&self, b: usize, g: u8) {
        self.data.lock().unwrap()[b * BS..(b + 1) * BS].copy_from_slice(&block(b, g));
    }
}

impl BlockSource for Source {
    fn read_at(&self, off: u64, buf: &mut [u8]) -> Result<()> {
        let d = self.data.lock().unwrap();
        buf.copy_from_slice(&d[off as usize..off as usize + buf.len()]);
        self.read.fetch_add(buf.len(), Ordering::Relaxed);
        Ok(())
    }
}

async fn check_cold(s: &MemoryStaging, rt: &HyperFileRuntimeConfig, name: &str, n: usize, want: impl Fn(usize) -> Vec<u8>) {
    let mut g = reopen(s, rt).await;
    let mut v = vec![0u8; BS];
    for b in 0..n {
        g.read(b * BS, &mut v).await.unwrap();
        assert!(v == want(b), "{}: block {} reads back wrong", name, b);
    }
}

/// Written by reference only, with the dirty-data limits at their defaults:
/// objects are sent ahead, whole and in pieces, their bytes read once from the
/// source, and the flush reads only what was not sent.
#[tokio::test]
async fn references_are_sent_ahead_and_the_flush_reads_only_the_rest() {
    for (name, target) in [("wr-whole", 5 * MIB), ("wr-pieces", 16 * MIB)] {
        let rt = runtime(target, false);
        let (s, mut f) = create(name, &rt).await;
        let n = 64 * MIB / BS;
        let src = Source::new(n, 1);
        f.set_block_source(Some(Arc::new(src.clone()))).unwrap();
        for k in 0..n / 256 {
            f.write_ref(k * 256 * BS, 256 * BS).await.unwrap();
            tokio::task::yield_now().await;
        }
        // Let what is in flight land.
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        let before = s.segment_ids().len();
        let read_ahead = src.read.load(Ordering::Relaxed);
        f.flush_timing().reset();
        f.flush().await.expect("flush");
        let t = f.flush_timing().snapshot();
        let read_by_flush = src.read.load(Ordering::Relaxed) - read_ahead;
        assert!(before >= 3, "{}: nothing was sent ahead before the flush ({} objects)", name, before);
        assert!(read_ahead >= 32 * MIB, "{}: {} bytes read ahead", name, read_ahead);
        assert!(read_by_flush <= n * BS - read_ahead + 16 * MIB,
            "{}: the flush read {} bytes, {} were sent ahead", name, read_by_flush, read_ahead);
        assert!(src.read.load(Ordering::Relaxed) <= n * BS + 16 * MIB, "{}: read more than once", name);
        assert_eq!(t.stream_dead_bytes, 0, "{}", name);
        drop(f);
        check_cold(&s, &rt, name, n, |b| block(b, 1)).await;
    }
}

/// Holds part 1's completion for `ms`, then refuses it when told; every other
/// object passes at once. Part 1 by name: objects can arrive in either order.
#[derive(Clone)]
struct Slow { ms: u64, fail_first: bool }

impl StagingIntercept<MemoryStaging> for Slow {
    fn after_flush_inode(&self, _: &MemoryStaging, _: &[u8], _: FlushInodeFlag)
        -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + '_ + Send>> {
        Box::pin(async { Ok(()) })
    }
    fn after_remove_inode(&self, _: &MemoryStaging)
        -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + '_ + Send>> {
        Box::pin(async { Ok(()) })
    }
    fn before_segment_done(&self, _: &MemoryStaging, segid: SegmentId, _: &[u8], _: usize)
        -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + '_ + Send>> {
        let first = segid.part_id() == Some(1);
        let (ms, fail) = (if first { self.ms } else { 0 }, first && self.fail_first);
        Box::pin(async move {
            tokio::time::sleep(std::time::Duration::from_millis(ms)).await;
            if fail { Err(Error::new(ErrorKind::ConnectionReset, "object refused")) } else { Ok(()) }
        })
    }
}

/// A block sent by reference and written again while its object is in flight —
/// by reference with new bytes, by value, zeroed, and in part (which reads the
/// rest of it from the source and keeps its pointer) — leaves a dead slot and
/// keeps the newer write; the rest of the object is stored.
#[tokio::test]
async fn a_reference_written_again_in_flight_keeps_the_newer_write() {
    let rt = runtime(5 * MIB, false);
    let name = "wr-again";
    let (s, mut f) = create(name, &rt).await;
    let n = 5 * MIB / BS;
    let src = Source::new(2 * n, 1);
    f.set_block_source(Some(Arc::new(src.clone()))).unwrap();
    f.staging_interceptor(Slow { ms: 300, fail_first: false });
    f.write_ref(0, n * BS).await.unwrap();
    // In flight now. Again, three ways.
    for b in 0..8 { src.set(b, 2); }
    f.write_ref(0, 8 * BS).await.unwrap();
    f.write(8 * BS, &[block(8, 3), block(9, 3)].concat()).await.unwrap();
    f.write_zero(10 * BS, BS).await.unwrap();
    f.write(11 * BS + 100, &[0xa5u8; 50]).await.unwrap();
    tokio::time::sleep(std::time::Duration::from_millis(400)).await;
    f.flush_timing().reset();
    f.flush().await.expect("flush");
    assert_eq!(f.flush_timing().snapshot().stream_dead_bytes, (12 * BS) as u64, "{}: dead slots", name);
    drop(f);
    check_cold(&s, &rt, name, n, |b| match b {
        0..=7 => block(b, 2),
        8 | 9 => block(b, 3),
        10 => vec![0u8; BS],
        11 => { let mut v = block(b, 1); v[100..150].fill(0xa5); v },
        _ => block(b, 1),
    }).await;
}

/// An object holding references that fails: they are references still, and the
/// flush reads them from the source.
#[tokio::test]
async fn references_in_an_object_that_fails_are_written_by_the_flush() {
    let rt = runtime(5 * MIB, false);
    let name = "wr-fail";
    let (s, mut f) = create(name, &rt).await;
    let n = 10 * MIB / BS;
    let src = Source::new(n, 1);
    f.set_block_source(Some(Arc::new(src.clone()))).unwrap();
    f.staging_interceptor(Slow { ms: 50, fail_first: true });
    f.write_ref(0, n * BS).await.unwrap();
    tokio::time::sleep(std::time::Duration::from_millis(200)).await;
    f.flush_timing().reset();
    f.flush().await.expect("flush");
    assert!(f.flush_timing().snapshot().stream_dead_bytes >= (5 * MIB) as u64, "{}: the refused object counts dead", name);
    drop(f);
    check_cold(&s, &rt, name, n, |b| block(b, 1)).await;
}

/// A leaf over blocks in two objects: references in one still in flight, and
/// later blocks in another that lands first. The leaf stays dirty while the
/// references are not stored, so when their object fails, the flush writes the
/// leaf with their new places.
#[tokio::test]
async fn a_leaf_over_references_not_yet_stored_stays_dirty() {
    let name = "wr-leaf";
    let mut rt = runtime(16 * 1024, false);
    rt.segment_mpu_chunk_size = 5 * MIB;
    let (s, mut f) = create(name, &rt).await;
    // A published file with leaves of the block map, about half full.
    let n = 8192;
    for b in (0..n).step_by(2) {
        f.write(b * BS, &block(b, 1)).await.unwrap();
    }
    f.flush().await.expect("publish the base");
    let src = Source::new(n, 2);
    f.set_block_source(Some(Arc::new(src.clone()))).unwrap();
    f.staging_interceptor(Slow { ms: 300, fail_first: true });
    // The first object: references, sent and slow to fail.
    f.write_ref(0, 4 * BS).await.unwrap();
    // The second: blocks under the same leaf, by value, stored at once.
    for b in [5usize, 7, 9, 11] {
        f.write(b * BS, &block(b, 3)).await.unwrap();
    }
    tokio::time::sleep(std::time::Duration::from_millis(50)).await;
    // Takes the second in.
    f.write(4001 * BS, &block(4001, 3)).await.unwrap();
    tokio::time::sleep(std::time::Duration::from_millis(400)).await;
    f.flush().await.expect("flush");
    drop(f);
    // The file ends at the last block written, n - 2.
    check_cold(&s, &rt, name, n - 1, |b| match b {
        0..=3 => block(b, 2),
        5 | 7 | 9 | 11 | 4001 => block(b, 3),
        b if b % 2 == 0 => block(b, 1),
        _ => vec![0u8; BS],
    }).await;
}

/// References and dirty blocks mixed in one leaf, sent ahead together.
#[tokio::test]
async fn references_and_dirty_blocks_together() {
    for (name, target) in [("wr-mix-whole", 5 * MIB), ("wr-mix-pieces", 16 * MIB)] {
        let rt = runtime(target, true);
        let (s, mut f) = create(name, &rt).await;
        let n = 48 * MIB / BS;
        let src = Source::new(n, 1);
        f.set_block_source(Some(Arc::new(src.clone()))).unwrap();
        for k in 0..n / 64 {
            if k % 3 == 0 {
                for b in k * 64..(k + 1) * 64 { f.write(b * BS, &block(b, 4)).await.unwrap(); }
            } else {
                f.write_ref(k * 64 * BS, 64 * BS).await.unwrap();
            }
            tokio::task::yield_now().await;
        }
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        let before = s.segment_ids().len();
        f.flush().await.expect("flush");
        // Two objects at least; filled in pieces, the last is still open.
        assert!(before >= 2, "{}: nothing was sent ahead ({} objects)", name, before);
        drop(f);
        check_cold(&s, &rt, name, n, |b| block(b, if (b / 64) % 3 == 0 { 4 } else { 1 })).await;
    }
}