diff --git a/lib/src/async_db.rs b/lib/src/async_db.rs new file mode 100644 index 0000000..3883160 --- /dev/null +++ b/lib/src/async_db.rs @@ -0,0 +1,188 @@ +use std::{marker::PhantomData, path::Path, sync::Arc}; + +use async_stream::stream; +use tokio::pin; + +use async_compression::tokio::bufread::ZstdDecoder; +use async_compression::tokio::write::ZstdEncoder; +use async_compression::Level; +use futures::stream::{self, StreamExt}; +use futures_core::stream::Stream; +use futures_util::pin_mut; + +use tokio::{ + fs, + io::{self, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}, +}; + +use fmmap::tokio::{AsyncMmapFile, AsyncOptions}; + +type LSize = u32; +const LEN_SIZE: usize = std::mem::size_of::(); +const BINCODE_CFG: bincode::config::Configuration = bincode::config::standard(); + +trait ErrorToString { + type Output; + fn str_err(self) -> std::result::Result; +} + +impl ErrorToString for std::result::Result +where + E: std::error::Error, +{ + type Output = T; + fn str_err(self) -> std::result::Result { + self.map_err(|e| e.to_string()) + } +} + +pub struct WriterOpts { + pub compress_lvl: Level, + pub data_buf_size: usize, + pub out_buf_size: usize, +} + +impl Default for WriterOpts { + fn default() -> Self { + Self { + compress_lvl: Level::Default, + data_buf_size: 500 * 1024 * 1024, + out_buf_size: 200 * 1024 * 1024, + } + } +} + +pub struct Writer +where + T: bincode::Encode, +{ + out: io::BufWriter, + data_buf: Vec, + table: Vec, + compress_lvl: Level, + _t: PhantomData<*const T>, +} + +impl Writer +where + T: bincode::Encode, +{ + pub async fn new>(path: P, opts: WriterOpts) -> Result { + let out = fs::File::create(path).await.str_err()?; + let out = io::BufWriter::with_capacity(opts.out_buf_size, out); + let data_buf: Vec = Vec::with_capacity(opts.data_buf_size); + + let compress_lvl = opts.compress_lvl; + + let table: Vec = vec![]; + + Ok(Self { + out, + data_buf, + table, + compress_lvl, + _t: PhantomData, + }) + } + + pub async fn push(&mut self, item: T) -> Result<(), String> { + let pos: LSize = self.data_buf.len() as LSize; + + let item_data = bincode::encode_to_vec(item, BINCODE_CFG).str_err()?; + + let mut zencoder = ZstdEncoder::with_quality(&mut self.data_buf, self.compress_lvl); + zencoder.write_all(&item_data).await.str_err()?; + zencoder.flush().await.str_err()?; + + self.table.push(pos); + + Ok(()) + } + + pub async fn load(&mut self, source: S) -> Result<(), String> + where + S: Stream + std::marker::Unpin, + { + let hint = source.size_hint(); + let hint = std::cmp::max(hint.0, hint.1.unwrap_or(0)); + if hint > 0 { + self.table.reserve(hint); + } + + pin_mut!(source); + while let Some(item) = source.next().await { + self.push(item).await?; + } + + Ok(()) + } + + pub async fn finish(mut self) -> Result<(), String> { + // finish tab + let pos: LSize = self.data_buf.len() as LSize; + self.table.push(pos); + + // write tab + let tab_size = (self.table.len() * LEN_SIZE) as LSize; + for pos in self.table { + let pos_data = (pos + tab_size).to_le_bytes(); + self.out.write_all(&pos_data).await.str_err()?; + } + + // copy data + self.out.write_all(&self.data_buf[..]).await.str_err()?; + + self.out.flush().await.str_err()?; + Ok(()) + } +} + +#[cfg(test)] +mod test { + use super::*; + use tempfile::tempdir; + + #[derive(bincode::Encode, bincode::Decode, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] + struct TestData { + num: u64, + test: String, + } + + fn gen_data(count: usize) -> impl Iterator { + (0..count).into_iter().map(|i| TestData { + num: i as u64, + test: "test".repeat(i), + }) + } + + fn stream_iter(it: I) -> impl Stream + where + I: Iterator, + { + stream! { + for i in it { + yield i; + } + } + } + + #[tokio::test] + async fn test_write() { + let dir = tempdir().expect("tempdir"); + let tmpfile = dir.path().join("test.tmp"); + let opts = WriterOpts { + compress_lvl: Level::Default, + data_buf_size: 10 * 1024 * 1024, + out_buf_size: 10 * 1024 * 1024, + }; + let mut writer: Writer = Writer::new(&tmpfile, opts).await.expect("new writer"); + + let items_iter = gen_data(5); + let items: Vec = items_iter.collect(); + + let src = stream_iter(items.clone().into_iter()); + pin_mut!(src); + writer.load(src).await.expect("load"); + writer.finish().await.expect("finish write"); + } +} diff --git a/lib/src/lib.rs b/lib/src/lib.rs index b1ff844..5987921 100644 --- a/lib/src/lib.rs +++ b/lib/src/lib.rs @@ -1,3 +1,6 @@ +#[cfg(feature = "async")] +pub mod async_db; +#[cfg(feature = "sync")] pub mod db; pub mod questions; pub mod source;