mod settings;
use std::io::{Result, ErrorKind};
use std::io::SeekFrom;
use std::time::Instant;
use rand::Rng;
use aws_sdk_s3::Client;
use human_bytes::human_bytes;
use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
use hyperfile::file::hyper::Hyper;
use hyperfile::file::tokio_wrapper::HyperFileTokio;
use hyperfile::file::flags::FileFlags;
use hyperfile::file::mode::FileMode;
use hyperfile::config::{
HyperFileConfigBuilder, HyperFileConfig,
HyperFileRuntimeConfig, HyperFileMetaConfig,
};
use hyperfile::staging::config::StagingConfig;
#[cfg(feature = "wal")]
use hyperfile::wal::config::HyperFileWalConfig;
use settings::*;
async fn random_write(client: &Client, config: &HyperFileConfig, data: &mut Vec<u8>, write_zero: bool) -> Result<()> {
let rand_data_offset = rand::random_range(DEFAULT_RANDOM_WRITE_OFFSET);
let mut rand_data_bytes = rand::random_range(DEFAULT_RANDOM_WRITE_BYTES);
if rand_data_offset + rand_data_bytes >= DEFAULT_MAX_FILE_SIZE {
rand_data_bytes = DEFAULT_MAX_FILE_SIZE - rand_data_offset;
}
let total_len = rand_data_offset + rand_data_bytes;
if total_len > data.len() {
data.resize(total_len, 0);
}
if write_zero {
print!("random write zero offset {} bytes {} ..", rand_data_offset, rand_data_bytes);
} else {
print!("random write offset {} bytes {} ..", rand_data_offset, rand_data_bytes);
}
let flags = FileFlags::rdwr();
let mode = FileMode::default_file();
let mut file = HyperFileTokio::open_or_create_with_config(&client, config.clone(), flags, mode).await?;
let filled_buf = &mut data[rand_data_offset..total_len];
if write_zero {
filled_buf.fill(0);
} else {
rand::fill(filled_buf);
}
file.seek(SeekFrom::Start(rand_data_offset as u64)).await?;
let write_bytes = if write_zero {
file.write_zero(rand_data_bytes).await?
} else {
file.write(filled_buf).await?
};
assert!(write_bytes == rand_data_bytes);
file.flush().await?;
let last_cno = file.last_cno().await?;
file.shutdown().await?;
println!(". Done with last cno {}", last_cno);
Ok(())
}
async fn random_truncate(client: &Client, config: &HyperFileConfig, data: &mut Vec<u8>) -> Result<()> {
let curr_len = data.len();
let new_file_len = if curr_len == 0 {
rand::random_range(0..DEFAULT_MAX_FILE_SIZE)
} else if rand::rng().random_ratio(2, 10) {
curr_len - 1
} else if rand::rng().random_ratio(2, 10) {
curr_len + 1
} else {
rand::random_range(0..DEFAULT_MAX_FILE_SIZE)
};
print!("random truncate file to new size {} ..", new_file_len);
let flags = FileFlags::rdwr();
let mode = FileMode::default_file();
let mut file = HyperFileTokio::open_or_create_with_config(&client, config.clone(), flags, mode).await?;
let stat = file.metadata().await?;
assert!(stat.st_size == data.len() as i64);
if new_file_len > stat.st_size as usize {
print!(".[EXPEND]..");
} else if new_file_len < stat.st_size as usize {
print!(".[SHRINK]..");
} else {
print!(".[NO CHANGE]..");
}
file.set_len(new_file_len as u64).await?;
file.flush().await?;
let last_cno = file.last_cno().await?;
data.resize(new_file_len, 0);
file.shutdown().await?;
println!(". Done with last cno {}", last_cno);
Ok(())
}
async fn read_check(client: &Client, config: &HyperFileConfig, data: &mut Vec<u8>) -> Result<()> {
let total_bytes = data.len();
let flags = FileFlags::rdonly();
let mode = FileMode::default_file();
let mut file = HyperFileTokio::open_or_create_with_config(&client, config.clone(), flags, mode).await?;
let stat = file.metadata().await?;
assert!(stat.st_size == total_bytes as i64);
let mut buf = Vec::new();
buf.resize(total_bytes, 0);
file.seek(SeekFrom::Start(0)).await?;
let start = Instant::now();
let read_bytes = file.read(&mut buf).await?;
let nanos = start.elapsed().as_nanos();
let md5_data = md5::compute(&data);
let md5_file = md5::compute(&buf);
println!("real data in memory md5: {:?} - read data from file md5: {:?}", md5_data, md5_file);
assert!(read_bytes == total_bytes);
let throughput = (read_bytes as u128 * 1000_000_000) / nanos;
println!("bytes read {} throughput {}/s", human_bytes(read_bytes as f64), human_bytes(throughput as f64));
if md5_data != md5_file {
panic!("md5 check failed");
}
Ok(())
}
#[tokio::main]
async fn main() -> Result<()> {
env_logger::init();
let uri = std::env::var("HYPERFILE_STAGING_ROOT_URI").unwrap();
#[cfg(feature = "wal")]
let wal_uri = std::env::var("HYPERFILE_WAL_ROOT_URI").unwrap();
let config = aws_config::load_from_env().await;
let client = aws_sdk_s3::Client::new(&config);
let _ = Hyper::fs_unlink(&client, &uri).await
.or_else(|e| {
if e.kind() == ErrorKind::NotFound {
return Ok(());
}
Err(e)
})?;
let mut data: Vec<u8> = Vec::new();
let meta_config = HyperFileMetaConfig::default();
let staging_config = StagingConfig::new_s3_uri(&uri, None);
let runtime_config = HyperFileRuntimeConfig::default_large();
#[cfg(not(feature = "wal"))]
let file_config = HyperFileConfigBuilder::new()
.with_meta_config(&meta_config)
.with_staging_config(&staging_config)
.with_runtime_config(&runtime_config)
.build();
#[cfg(feature = "wal")]
let wal_config = HyperFileWalConfig::new(&wal_uri);
#[cfg(feature = "wal")]
let file_config = HyperFileConfigBuilder::new()
.with_meta_config(&meta_config)
.with_staging_config(&staging_config)
.with_runtime_config(&runtime_config)
.with_wal_config(&wal_config)
.build();
for _ in 0..NUM_ITER {
if rand::rng().random_ratio(5, 10) {
random_write(&client, &file_config, &mut data, false).await?;
} else if rand::rng().random_ratio(2, 10) {
random_write(&client, &file_config, &mut data, true).await?;
} else {
random_truncate(&client, &file_config, &mut data).await?;
}
}
println!("Final read check...");
read_check(&client, &file_config, &mut data).await?;
Ok(())
}