diff --git a/programs/fileio.c b/programs/fileio.c index 9f79b5c57..4af2c0942 100644 --- a/programs/fileio.c +++ b/programs/fileio.c @@ -343,6 +343,7 @@ int FIO_rust_setDictBufferMalloc(const char* fileName, void** buffer, size_t* loadedSize); int FIO_rust_removeFile(const char* path); +int FIO_rust_passThrough(ReadPoolCtx_t* readCtx, WritePoolCtx_t* writeCtx); #ifdef ZSTD_LZ4COMPRESS int FIO_rust_LZ4_GetBlockSize_FromBlockId(int id); #endif @@ -2123,24 +2124,7 @@ static void FIO_freeDResources(dRess_t ress) * @return : 0 (no error) */ static int FIO_passThrough(dRess_t *ress) { - size_t const blockSize = MIN(MIN(64 KB, ZSTD_DStreamInSize()), ZSTD_DStreamOutSize()); - IOJob_t *writeJob = AIO_WritePool_acquireJob(ress->writeCtx); - AIO_ReadPool_fillBuffer(ress->readCtx, blockSize); - - while(ress->readCtx->srcBufferLoaded) { - size_t writeSize; - writeSize = MIN(blockSize, ress->readCtx->srcBufferLoaded); - assert(writeSize <= writeJob->bufferSize); - memcpy(writeJob->buffer, ress->readCtx->srcBuffer, writeSize); - writeJob->usedBufferSize = writeSize; - AIO_WritePool_enqueueAndReacquireWriteJob(&writeJob); - AIO_ReadPool_consumeBytes(ress->readCtx, writeSize); - AIO_ReadPool_fillBuffer(ress->readCtx, blockSize); - } - assert(ress->readCtx->reachedEof); - AIO_WritePool_releaseIoJob(writeJob); - AIO_WritePool_sparseWriteEnd(ress->writeCtx); - return 0; + return FIO_rust_passThrough(ress->readCtx, ress->writeCtx); } /* FIO_zstdErrorHelp() : diff --git a/rust/src/fileio_asyncio.rs b/rust/src/fileio_asyncio.rs index e642b1ad4..5929a3cc8 100644 --- a/rust/src/fileio_asyncio.rs +++ b/rust/src/fileio_asyncio.rs @@ -29,6 +29,7 @@ const MAX_IO_JOBS: usize = 10; const IO_QUEUE_SIZE: usize = MAX_IO_JOBS - 2; const SPARSE_SEGMENT_SIZE: usize = 32 * 1024; const SPARSE_SKIP_CHUNK: u64 = 1 << 30; +const PASS_THROUGH_MAX_BLOCK_SIZE: usize = 64 * 1024; /// C's `FIO_prefs_t` from `programs/fileio_types.h`. /// @@ -1350,6 +1351,48 @@ pub unsafe extern "C" fn AIO_ReadPool_closeFile(ctx: *mut ReadPoolCtx_t) -> c_in unsafe { libc::fclose(file) } } +#[no_mangle] +pub unsafe extern "C" fn FIO_rust_passThrough( + read_ctx: *mut ReadPoolCtx_t, + write_ctx: *mut WritePoolCtx_t, +) -> c_int { + assert!(!read_ctx.is_null()); + assert!(!write_ctx.is_null()); + + let read_context = read_ctx.cast::(); + let write_context = write_ctx.cast::(); + let read_inner = unsafe { base_inner(read_context) }; + let write_inner = unsafe { base_inner(write_context) }; + let block_size = PASS_THROUGH_MAX_BLOCK_SIZE + .min(unsafe { read_at::(read_context, read_inner.layout.job_buffer_size) }) + .min(unsafe { read_at::(write_context, write_inner.layout.job_buffer_size) }); + + let mut write_job = unsafe { AIO_WritePool_acquireJob(write_ctx) }; + unsafe { AIO_ReadPool_fillBuffer(read_ctx, block_size) }; + + while unsafe { read_at::(read_context, read_inner.layout.read_src_buffer_loaded) } != 0 { + let loaded = + unsafe { read_at::(read_context, read_inner.layout.read_src_buffer_loaded) }; + let write_size = block_size.min(loaded); + assert!(write_size <= unsafe { (*write_job).bufferSize }); + let source = + unsafe { read_at::<*const u8>(read_context, read_inner.layout.read_src_buffer) }; + unsafe { + ptr::copy_nonoverlapping(source, (*write_job).buffer.cast::(), write_size); + (*write_job).usedBufferSize = write_size; + AIO_WritePool_enqueueAndReacquireWriteJob(&mut write_job); + AIO_ReadPool_consumeBytes(read_ctx, write_size); + AIO_ReadPool_fillBuffer(read_ctx, block_size); + } + } + assert!(unsafe { read_at::(read_context, read_inner.layout.read_reached_eof) } != 0); + unsafe { + AIO_WritePool_releaseIoJob(write_job); + AIO_WritePool_sparseWriteEnd(write_ctx); + } + 0 +} + impl PoolInner { unsafe fn all_jobs_available(&self) -> bool { let state = self.jobs.lock().unwrap_or_else(|error| error.into_inner()); @@ -1739,4 +1782,73 @@ mod tests { ); unsafe { AIO_ReadPool_free(ctx) }; } + + #[cfg(unix)] + fn run_pass_through( + input: &[u8], + read_buffer_size: usize, + write_buffer_size: usize, + async_io: c_int, + ) -> Vec { + let input_file = unsafe { libc::tmpfile() }; + assert!(!input_file.is_null()); + if !input.is_empty() { + assert_eq!( + unsafe { libc::fwrite(input.as_ptr().cast(), 1, input.len(), input_file) }, + input.len() + ); + } + assert_eq!(unsafe { libc::fflush(input_file) }, 0); + assert_eq!(unsafe { libc::fseek(input_file, 0, libc::SEEK_SET) }, 0); + + let output_file = unsafe { libc::tmpfile() }; + assert!(!output_file.is_null()); + let mut prefs = test_prefs(async_io); + prefs.testMode = 0; + prefs.sparseFileSupport = 0; + let read_ctx = unsafe { AIO_ReadPool_create(&prefs, read_buffer_size) }; + let write_ctx = unsafe { AIO_WritePool_create(&prefs, write_buffer_size) }; + unsafe { + AIO_ReadPool_setFile(read_ctx, input_file); + AIO_WritePool_setFile(write_ctx, output_file); + } + + assert_eq!(unsafe { FIO_rust_passThrough(read_ctx, write_ctx) }, 0); + assert_eq!(unsafe { libc::fseek(output_file, 0, libc::SEEK_END) }, 0); + let output_size = unsafe { libc::ftell(output_file) }; + assert!(output_size >= 0); + let output_size = usize::try_from(output_size).unwrap(); + assert_eq!(output_size, input.len()); + assert_eq!(unsafe { libc::fseek(output_file, 0, libc::SEEK_SET) }, 0); + let mut output = vec![0_u8; output_size]; + if output_size != 0 { + assert_eq!( + unsafe { libc::fread(output.as_mut_ptr().cast(), 1, output_size, output_file) }, + output_size + ); + } + + unsafe { + AIO_WritePool_free(write_ctx); + AIO_ReadPool_free(read_ctx); + } + output + } + + #[cfg(unix)] + #[test] + fn pass_through_copies_across_bounded_job_buffers() { + let input = b"pass-through data spans several blocks"; + for async_io in [0, 1] { + assert_eq!(run_pass_through(input, 5, 3, async_io).as_slice(), input); + } + } + + #[cfg(unix)] + #[test] + fn pass_through_handles_empty_input() { + for async_io in [0, 1] { + assert!(run_pass_through(&[], 5, 3, async_io).is_empty()); + } + } }