hyperfile 0.11.2

//! Segment uploads large enough to go multipart, on a real bucket.
//!
//! A segment longer than `segment_mpu_chunk_size` is written as a multipart upload
//! instead of a single put. These pin what that path has to hold to the same terms
//! as a single put: the complete carries a condition, so it cannot replace an object
//! already under its name, and an upload a failure leaves open is aborted rather than
//! left to be billed until a lifecycle rule removes it. And no more parts are in flight
//! at once than `segment_mpu_max_concurrency` allows.
//!
//! Run in both builds: the default one uploads through the scatter path, `wal`
//! through the contiguous one.
//!
//! ```bash
//! cargo test --test integration_s3_multipart -- --ignored --test-threads=1
//! cargo test --features wal --test integration_s3_multipart -- --ignored --test-threads=1
//! ```

mod common;
use common::*;

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

use hyperfile::config::{FlushConflictPolicy, HyperFileRuntimeConfig};
use hyperfile::file::hyper::Hyper;
use hyperfile::file::flags::FileFlags;
use hyperfile::file::mode::FileMode;
use hyperfile::inode::FlushInodeFlag;
use hyperfile::segment::Segment;
use hyperfile::staging::StagingIntercept;
use hyperfile::staging::s3::S3Staging;
use hyperfile::SegmentId;

const MIB: usize = 1024 * 1024;
/// The smallest part S3 accepts, so a few MiB of data is a few parts.
const PART: usize = 5 * MIB;
/// Three parts' worth, so a flush of it goes multipart.
const DATA: usize = 12 * MIB;

fn prefix_of(uri: &str) -> String {
    uri.trim_start_matches("s3://").splitn(2, '/').nth(1).unwrap().to_string()
}

/// Settings under which one flush of `DATA` is one multipart segment: small parts,
/// and every threshold that would flush or write a partial along the way out of reach.
fn multipart_runtime() -> HyperFileRuntimeConfig {
    let mut r = HyperFileRuntimeConfig::default();
    r.segment_mpu_chunk_size = PART;
    r.data_cache_dirty_max_bytes_threshold = 64 * MIB;
    r.data_cache_dirty_max_blocks_threshold = 64 * MIB / 4096;
    r.data_cache_dirty_max_flush_interval = u64::MAX / 2;
    r
}

fn payload() -> Vec<u8> {
    payload_of(DATA)
}

fn payload_of(len: usize) -> Vec<u8> {
    (0..len).map(|i| (i.wrapping_mul(31) >> 7) as u8).collect()
}

async fn seeded(client: &aws_sdk_s3::Client) -> TestFile {
    let tf = TestFile::new(client).await;
    let mut h = Hyper::fs_open_or_create_with_default_opt(client, tf.uri(), FileFlags::rdwr(),
        FileMode::default_file()).await.expect("create");
    h.fs_write(0, b"before").await.unwrap();
    h.fs_flush().await.unwrap();
    let _ = h.fs_release().await;
    tf
}

/// Keys with a multipart upload still open under the container's prefix.
async fn open_uploads(client: &aws_sdk_s3::Client, prefix: &str) -> Vec<(String, String)> {
    let out = client.list_multipart_uploads().bucket(test_bucket()).prefix(format!("{}/", prefix))
        .send().await.expect("list multipart uploads");
    out.uploads().iter()
        .map(|u| (u.key().unwrap_or("").to_string(), u.upload_id().unwrap_or("").to_string()))
        .collect()
}

/// What was left open under the container, aborted before it is returned.
///
/// Called before anything is asserted, so a test that fails -- as every one does in a
/// revert experiment -- still leaves no parts behind to be billed.
async fn left_open(client: &aws_sdk_s3::Client, prefix: &str) -> Vec<(String, String)> {
    let open = open_uploads(client, prefix).await;
    for (key, id) in &open {
        let _ = client.abort_multipart_upload().bucket(test_bucket()).key(key).upload_id(id).send().await;
    }
    open
}

async fn etag(client: &aws_sdk_s3::Client, key: &str) -> String {
    client.head_object().bucket(test_bucket()).key(key).send().await.expect("head")
        .e_tag().unwrap_or_default().to_string()
}

async fn inode_counter(client: &aws_sdk_s3::Client, prefix: &str) -> u64 {
    let b = client.get_object().bucket(test_bucket()).key(format!("{}/inode", prefix)).send().await
        .expect("get inode").body.collect().await.unwrap().into_bytes();
    u64::from_le_bytes(b[88..96].try_into().unwrap())
}

async fn read_back(client: &aws_sdk_s3::Client, uri: &str, len: usize) -> Vec<u8> {
    let mut h = Hyper::fs_open(client, uri, FileFlags::rdonly()).await.expect("open");
    let mut buf = vec![0u8; len];
    let n = h.fs_read(0, &mut buf).await.expect("read");
    buf.truncate(n);
    let _ = h.fs_release().await;
    buf
}

/// Before the first segment write only, puts an object of its own under that
/// segment's name: `junk` if given, else the segment's own bytes -- which is what a
/// complete that landed and lost its answer leaves behind.
#[derive(Clone, Default)]
struct NameTakenFirst {
    junk: Option<Arc<Vec<u8>>>,
    done: Arc<AtomicBool>,
    at: Arc<std::sync::Mutex<Option<String>>>,
}

impl NameTakenFirst {
    fn key(&self) -> String {
        self.at.lock().unwrap().clone().expect("no segment was written")
    }
}

impl StagingIntercept<S3Staging> for NameTakenFirst {
    fn after_flush_inode(&self, _: &S3Staging, _: &[u8], _: FlushInodeFlag)
        -> std::pin::Pin<Box<dyn std::future::Future<Output = std::io::Result<()>> + '_ + Send>> {
        Box::pin(async { Ok(()) })
    }
    fn after_remove_inode(&self, _: &S3Staging)
        -> std::pin::Pin<Box<dyn std::future::Future<Output = std::io::Result<()>> + '_ + Send>> {
        Box::pin(async { Ok(()) })
    }
    fn before_segment_done(&self, staging: &S3Staging, segid: SegmentId, buf: &[u8], len: usize)
        -> std::pin::Pin<Box<dyn std::future::Future<Output = std::io::Result<()>> + '_ + Send>> {
        let first = !self.done.swap(true, Ordering::SeqCst);
        let key = format!("{}/{}", staging.root_path, Segment::segid_to_staging_file_id(segid));
        if first {
            *self.at.lock().unwrap() = Some(key.clone());
        }
        let client = staging.client.clone();
        let bucket = staging.bucket.clone();
        let body = match &self.junk {
            Some(j) => j.to_vec(),
            None => buf[..len].to_vec(),
        };
        Box::pin(async move {
            if first {
                client.put_object().bucket(bucket).key(key).body(body.into()).send().await
                    .map_err(|e| std::io::Error::other(format!("taking the name first: {}", e)))?;
            }
            Ok(())
        })
    }
}

/// Refuses every upload of one part number, before the request is sent.
#[derive(Debug)]
struct RefusePart(i32);

impl aws_sdk_s3::config::Intercept for RefusePart {
    fn name(&self) -> &'static str {
        "RefusePart"
    }
    fn read_before_execution(
        &self,
        ctx: &aws_sdk_s3::config::interceptors::BeforeSerializationInterceptorContextRef<'_>,
        _cfg: &mut aws_sdk_s3::config::ConfigBag,
    ) -> Result<(), aws_sdk_s3::error::BoxError> {
        if let Some(input) = ctx.input().downcast_ref::<aws_sdk_s3::operation::upload_part::UploadPartInput>() {
            if input.part_number() == Some(self.0) {
                return Err(format!("part {} refused by the test", self.0).into());
            }
        }
        Ok(())
    }
}

async fn client_refusing_part(n: i32) -> aws_sdk_s3::Client {
    let sdk = aws_config::from_env().region(aws_config::Region::new(test_region())).load().await;
    let conf = aws_sdk_s3::config::Builder::from(&sdk).interceptor(RefusePart(n)).build();
    aws_sdk_s3::Client::from_conf(conf)
}

/// One client's part uploads in flight, and the most there were at once.
#[derive(Debug, Default)]
struct PartsInFlight {
    now: AtomicUsize,
    peak: AtomicUsize,
}

/// Left in an operation's config bag when its start was counted, so that its end is.
#[derive(Debug, Clone)]
struct CountedPart;

impl aws_smithy_types::config_bag::Storable for CountedPart {
    type Storer = aws_smithy_types::config_bag::StoreReplace<Self>;
}

/// Counts part uploads from when each starts to when it has finished, retries and all.
#[derive(Debug)]
struct CountParts(Arc<PartsInFlight>);

impl aws_sdk_s3::config::Intercept for CountParts {
    fn name(&self) -> &'static str {
        "CountParts"
    }
    fn read_before_execution(
        &self,
        ctx: &aws_sdk_s3::config::interceptors::BeforeSerializationInterceptorContextRef<'_>,
        cfg: &mut aws_sdk_s3::config::ConfigBag,
    ) -> Result<(), aws_sdk_s3::error::BoxError> {
        if ctx.input().downcast_ref::<aws_sdk_s3::operation::upload_part::UploadPartInput>().is_some() {
            cfg.interceptor_state().store_put(CountedPart);
            let now = self.0.now.fetch_add(1, Ordering::SeqCst) + 1;
            self.0.peak.fetch_max(now, Ordering::SeqCst);
        }
        Ok(())
    }
    fn read_after_execution(
        &self,
        _ctx: &aws_sdk_s3::config::interceptors::FinalizerInterceptorContextRef<'_>,
        _rc: &aws_sdk_s3::config::RuntimeComponents,
        cfg: &mut aws_sdk_s3::config::ConfigBag,
    ) -> Result<(), aws_sdk_s3::error::BoxError> {
        if cfg.load::<CountedPart>().is_some() {
            self.0.now.fetch_sub(1, Ordering::SeqCst);
        }
        Ok(())
    }
}

async fn client_counting_parts() -> (aws_sdk_s3::Client, Arc<PartsInFlight>) {
    let counts = Arc::new(PartsInFlight::default());
    let sdk = aws_config::from_env().region(aws_config::Region::new(test_region())).load().await;
    let conf = aws_sdk_s3::config::Builder::from(&sdk).interceptor(CountParts(counts.clone())).build();
    (aws_sdk_s3::Client::from_conf(conf), counts)
}

/// What one flush of `len` bytes under a part bound did: the flush's result, the
/// segment's ETag if it was published, and the most parts that were in flight at once.
struct Bounded {
    res: std::io::Result<u64>,
    tag: Option<String>,
    peak: usize,
    open: Vec<(String, String)>,
    read: Vec<u8>,
}

async fn flush_under_bound(client: &aws_sdk_s3::Client, tf: &TestFile, data: &[u8], bound: usize) -> Bounded {
    let p = prefix_of(tf.uri());
    let (counting, counts) = client_counting_parts().await;
    let mut r = multipart_runtime();
    r.segment_mpu_max_concurrency = bound;
    let res = {
        let mut h = Hyper::fs_open_opt(&counting, tf.uri(), FileFlags::rdwr(), &r).await.unwrap();
        h.fs_write(0, data).await.unwrap();
        let res = h.fs_flush().await;
        let _ = h.fs_release().await;
        res
    };
    let open = left_open(client, &p).await;
    let tag = match &res {
        Ok(cno) => {
            let key = format!("{}/{}", p, Segment::segid_to_staging_file_id(SegmentId::new_from_cno(*cno).at_part(0)));
            Some(etag(client, &key).await)
        }
        Err(_) => None,
    };
    let read = read_back(client, tf.uri(), data.len()).await;
    Bounded { res, tag, peak: counts.peak.load(Ordering::SeqCst), open, read }
}

/// No more parts are in flight at once than the bound, and the upload still completes.
///
/// Six parts under a bound of two: every part used to start at once, so the count
/// seen was six; bounded, the third starts only when one of the first two finishes.
#[tokio::test]
#[ignore]
async fn no_more_parts_are_in_flight_than_the_bound() {
    let _ = env_logger::try_init();
    let client = make_client().await;
    let tf = seeded(&client).await;

    let data = payload_of(28 * MIB);
    let b = flush_under_bound(&client, &tf, &data, 2).await;
    tf.cleanup(&client).await;
    b.res.expect("flush");
    let tag = b.tag.unwrap();
    assert!(tag.trim_matches('"').ends_with("-6"),
        "the segment was not a six-part upload, so the bound was not what limited it: {}", tag);
    assert_eq!(b.peak, 2, "the most part uploads in flight at once was {}, under a bound of 2", b.peak);
    assert!(b.read == data, "the data came back wrong");
    assert!(b.open.is_empty(), "an upload was left open: {:?}", b.open);
}

/// A bound of zero is taken as one: the parts go one at a time, rather than none
/// going at all and the complete being sent with no parts.
#[tokio::test]
#[ignore]
async fn a_bound_of_zero_is_taken_as_one() {
    let _ = env_logger::try_init();
    let client = make_client().await;
    let tf = seeded(&client).await;

    let data = payload();
    let b = flush_under_bound(&client, &tf, &data, 0).await;
    tf.cleanup(&client).await;
    assert!(b.open.is_empty(), "an upload was left open: {:?}", b.open);
    b.res.expect("flush under a bound of zero");
    assert!(b.tag.unwrap().trim_matches('"').ends_with("-3"), "the segment was not a three-part upload");
    assert_eq!(b.peak, 1, "the most part uploads in flight at once was {}, under a bound of 0", b.peak);
    assert!(b.read == data, "the data came back wrong");
}

/// The baseline: a segment larger than a part is published, through multipart, and
/// reads back. Nothing else in the suites writes one.
#[tokio::test]
#[ignore]
async fn a_segment_larger_than_a_part_is_published_through_multipart() {
    let _ = env_logger::try_init();
    let client = make_client().await;
    let tf = seeded(&client).await;
    let p = prefix_of(tf.uri());

    let data = payload();
    let cno = {
        let mut h = Hyper::fs_open_opt(&client, tf.uri(), FileFlags::rdwr(), &multipart_runtime()).await.unwrap();
        h.fs_write(0, &data).await.unwrap();
        let cno = h.fs_flush().await.expect("flush");
        let _ = h.fs_release().await;
        cno
    };
    let open = left_open(&client, &p).await;
    let key = format!("{}/{}", p, Segment::segid_to_staging_file_id(SegmentId::new_from_cno(cno).at_part(0)));
    let tag = etag(&client, &key).await;
    assert!(tag.trim_matches('"').ends_with("-3"),
        "the segment was not a three-part upload, so nothing here tested multipart: {}", tag);
    assert!(read_back(&client, tf.uri(), DATA).await == data, "the data came back wrong");
    assert!(open.is_empty(), "an upload was left open: {:?}", open);
    tf.cleanup(&client).await;
}

/// The complete does not replace an object already under its name.
///
/// Every other create-only write carries `If-None-Match: *`; a multipart complete
/// that does not can overwrite a segment an inode already points at. Here something
/// takes the name while the upload is under way. The complete is refused, the upload
/// that was refused is aborted, the object that was there is untouched, and the flush,
/// told the name was taken, goes again at the next number.
#[tokio::test]
#[ignore]
async fn the_complete_does_not_replace_an_object_already_under_its_name() {
    let _ = env_logger::try_init();
    let client = make_client().await;
    let tf = seeded(&client).await;
    let p = prefix_of(tf.uri());

    let taken = NameTakenFirst { junk: Some(Arc::new(vec![0xEE; 4096])), ..Default::default() };
    let data = payload();
    let res = {
        let mut h = Hyper::fs_open_opt(&client, tf.uri(), FileFlags::rdwr(), &multipart_runtime()).await.unwrap();
        h.with_staging_interceptor(taken.clone());
        h.fs_write(0, &data).await.unwrap();
        let res = h.fs_flush().await;
        let _ = h.fs_release().await;
        res
    };
    let open = left_open(&client, &p).await;
    let cno = res.expect("the flush goes again at the next number and succeeds");
    let key = taken.key();
    let tag = etag(&client, &key).await;
    assert!(!tag.contains('-'), "the complete replaced what was under its name: now {}", tag);
    let body = client.get_object().bucket(test_bucket()).key(&key).send().await.unwrap()
        .body.collect().await.unwrap().into_bytes();
    assert!(body.iter().all(|b| *b == 0xEE), "the object under the taken name changed");
    assert!(open.is_empty(), "the refused upload was left open: {:?}", open);
    assert!(key.ends_with(&Segment::segid_to_staging_file_id(SegmentId::new_from_cno(cno - 1).at_part(0))),
        "the retry did not land on the next number: {} and {}", key, cno);
    assert!(read_back(&client, tf.uri(), DATA).await == data, "the data came back wrong");
    tf.cleanup(&client).await;
}

/// A part that fails aborts the upload.
#[tokio::test]
#[ignore]
async fn a_failed_part_aborts_the_upload() {
    let _ = env_logger::try_init();
    let client = make_client().await;
    let tf = seeded(&client).await;
    let p = prefix_of(tf.uri());

    let refusing = client_refusing_part(2).await;
    let res = {
        let mut h = Hyper::fs_open_opt(&refusing, tf.uri(), FileFlags::rdwr(), &multipart_runtime()).await.unwrap();
        h.fs_write(0, &payload()).await.unwrap();
        let res = h.fs_flush().await;
        let _ = h.fs_release().await;
        res
    };
    let open = left_open(&client, &p).await;
    let e = res.expect_err("part 2 is refused");
    assert!(open.is_empty(), "an upload with a failed part was left open: {:?}", open);
    assert!(e.to_string().contains("aborted"), "the error does not say what became of the upload: {}", e);
    assert_eq!(read_back(&client, tf.uri(), 6).await, b"before");
    tf.cleanup(&client).await;
}

/// The name already holds this flush's own segment, under `FailFast`.
///
/// That object is a commit to anyone who finds it: a whole segment, built on the
/// current inode. It is what a single put tried again after a lost answer leaves; a
/// multipart upload does not reach it that way -- S3 answers a complete repeated on the
/// same upload with success -- but the name can be taken so by other means, and what
/// is under test is what follows. The complete is refused with 200 and an error in the
/// body, which has to be read as the name being taken and not as some other failure;
/// the flush fails; and its rollback has to move the counter past the number, or the
/// next open takes up a flush its caller was told failed.
#[tokio::test]
#[ignore]
async fn a_name_holding_this_flushs_own_segment_is_abandoned_under_fail_fast() {
    let _ = env_logger::try_init();
    let client = make_client().await;
    let tf = seeded(&client).await;
    let p = prefix_of(tf.uri());

    let mut runtime = multipart_runtime();
    runtime.flush_conflict_policy = FlushConflictPolicy::FailFast;
    let landed = NameTakenFirst::default();
    let res = {
        let mut h = Hyper::fs_open_opt(&client, tf.uri(), FileFlags::rdwr(), &runtime).await.unwrap();
        h.with_staging_interceptor(landed.clone());
        h.fs_write(0, &payload()).await.unwrap();
        let res = h.fs_flush().await;
        let _ = h.fs_release().await;
        res
    };
    let open = left_open(&client, &p).await;
    let e = res.expect_err("the complete finds the name taken");
    assert_eq!(e.kind(), std::io::ErrorKind::AlreadyExists, "not reported as a taken name: {}", e);
    let key = landed.key();
    let tag = etag(&client, &key).await;
    assert!(!tag.contains('-'), "the complete replaced the object that was there: now {}", tag);
    assert!(open.is_empty(), "the refused upload was left open: {:?}", open);
    let n: u64 = key.rsplit('/').next().unwrap().split('.').next().unwrap().parse().unwrap();
    assert!(inode_counter(&client, &p).await >= n, "the counter was not moved past {}", n);
    assert_eq!(read_back(&client, tf.uri(), 6).await, b"before",
        "a flush its caller was told failed came back");
    tf.cleanup(&client).await;
}