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;
const PART: usize = 5 * MIB;
const DATA: usize = 12 * MIB;
fn prefix_of(uri: &str) -> String {
uri.trim_start_matches("s3://").splitn(2, '/').nth(1).unwrap().to_string()
}
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
}
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()
}
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
}
#[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(())
})
}
}
#[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)
}
#[derive(Debug, Default)]
struct PartsInFlight {
now: AtomicUsize,
peak: AtomicUsize,
}
#[derive(Debug, Clone)]
struct CountedPart;
impl aws_smithy_types::config_bag::Storable for CountedPart {
type Storer = aws_smithy_types::config_bag::StoreReplace<Self>;
}
#[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)
}
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 }
}
#[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);
}
#[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");
}
#[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;
}
#[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;
}
#[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;
}
#[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;
}