use futures::FutureExt; use futures::StreamExt; pub trait DemoStorage: Send + Sync { fn duplicate(&self) -> Box; fn upload<'f, 's, 'own>( &'own self, user_id: String, demo_id: String, stream: futures_util::stream::BoxStream<'s, axum::body::Bytes>, ) -> futures::future::BoxFuture<'f, Result<(), String>> where 's: 'f, 'own: 'f; fn load<'f, 'own>( &'own self, user_id: String, demo_id: String, ) -> futures::future::BoxFuture<'f, Result> where 'own: 'f; } pub struct FileStorage { folder: std::sync::Arc, } impl FileStorage { pub fn new

(folder: P) -> Self where P: Into, { Self { folder: std::sync::Arc::new(folder.into()), } } } impl DemoStorage for FileStorage { fn duplicate(&self) -> Box { Box::new(Self { folder: self.folder.clone(), }) } fn upload<'f, 's, 'own>( &'own self, user_id: String, demo_id: String, stream: futures_util::stream::BoxStream<'s, axum::body::Bytes>, ) -> futures::future::BoxFuture<'f, Result<(), String>> where 's: 'f, 'own: 'f, { let path = self.folder.clone(); async move { let user_folder = std::path::Path::new(path.as_path()).join(format!("{}/", user_id)); if !tokio::fs::try_exists(&user_folder).await.unwrap_or(false) { tokio::fs::create_dir_all(&user_folder).await.unwrap(); } let demo_file_path = user_folder.join(format!("{}.dem", demo_id)); async { // Convert the stream into an `AsyncRead`. let body_with_io_error = stream.map(|b| Ok::<_, std::io::Error>(b)); let body_reader = tokio_util::io::StreamReader::new(body_with_io_error); futures::pin_mut!(body_reader); // Create the file. `File` implements `AsyncWrite`. let mut file = tokio::io::BufWriter::new(tokio::fs::File::create(demo_file_path).await?); // Copy the body into the file. tokio::io::copy(&mut body_reader, &mut file).await?; Ok::<_, std::io::Error>(()) } .await .map_err(|err| err.to_string()) } .boxed() } fn load<'f, 'own>( &'own self, user_id: String, demo_id: String, ) -> futures::future::BoxFuture<'f, Result> where 'own: 'f, { async move { let user_folder = std::path::Path::new(self.folder.as_path()).join(format!("{}/", user_id)); let demo_file_path = user_folder.join(format!("{}.dem", demo_id)); let file = std::fs::File::open(demo_file_path.as_path()).unwrap(); let mmap = unsafe { memmap2::MmapOptions::new().map(&file).unwrap() }; Ok(crate::analysis::AnalysisData::MemMapped( std::sync::Arc::new(mmap), )) } .boxed() } } pub struct S3Storage { bucket: std::sync::Arc, } impl S3Storage { pub fn new( bucket_name: &str, region: s3::region::Region, credentials: s3::creds::Credentials, ) -> Self { let mut bucket = s3::bucket::Bucket::new(bucket_name, region, credentials).unwrap(); bucket.set_path_style(); Self { bucket: bucket.into(), } } } impl DemoStorage for S3Storage { fn duplicate(&self) -> Box { Box::new(Self { bucket: self.bucket.clone(), }) } fn upload<'f, 's, 'own>( &'own self, user_id: String, demo_id: String, stream: futures_util::stream::BoxStream<'s, axum::body::Bytes>, ) -> futures::future::BoxFuture<'f, Result<(), String>> where 's: 'f, 'own: 'f, { async move { let path = std::path::PathBuf::new().join(user_id).join(demo_id); let path = path.to_str().unwrap(); // Convert the stream into an `AsyncRead`. let body_with_io_error = stream.map(|b| Ok::<_, std::io::Error>(b)); let body_reader = tokio_util::io::StreamReader::new(body_with_io_error); futures::pin_mut!(body_reader); self.bucket.list(String::new(), None).await.unwrap(); self.bucket .put_object_stream(&mut body_reader, path) .await .unwrap(); Ok(()) } .boxed() } fn load<'f, 'own>( &'own self, user_id: String, demo_id: String, ) -> futures::future::BoxFuture<'f, Result> where 'own: 'f, { async move { let path = std::path::PathBuf::new().join(user_id).join(demo_id); let path = path.to_str().unwrap(); let resp = self.bucket.get_object(path).await.unwrap(); Ok(crate::analysis::AnalysisData::Preloaded( resp.to_vec().into(), )) } .boxed() } }