Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 4366971

Browse files
George-MiaoGeorge-Miao
committed
feat: add test
1 parent 58561d0 commit 4366971

8 files changed

Lines changed: 194 additions & 28 deletions

File tree

‎compio-buf/src/io_buf.rs‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -855,7 +855,7 @@ where
855855
#[cfg(feature = "memmap2")]
856856
impl SetLen for memmap2::MmapMut {
857857
unsafe fn set_len(&mut self, len: usize) {
858-
assert!(len <= self.len())
858+
debug_assert!(len <= self.len())
859859
}
860860
}
861861

@@ -949,6 +949,39 @@ mod test {
949949
assert!(buf.capacity() >= 10);
950950
}
951951

952+
#[test]
953+
#[cfg(feature = "memmap2")]
954+
fn tests_memmap2() {
955+
use std::{
956+
fs::{OpenOptions, remove_file},
957+
io::{Seek, SeekFrom, Write},
958+
};
959+
960+
use memmap2::MmapOptions;
961+
962+
use super::*;
963+
964+
let path = std::env::temp_dir().join("compio_buf_mmap_mut_test");
965+
966+
let mut file = OpenOptions::new()
967+
.read(true)
968+
.write(true)
969+
.create(true)
970+
.truncate(true)
971+
.open(&path)
972+
.unwrap();
973+
let data = b"hello memmap2";
974+
file.write_all(data).unwrap();
975+
file.flush().unwrap();
976+
file.seek(SeekFrom::Start(0)).unwrap();
977+
let mmap = unsafe { MmapOptions::new().map(&file).unwrap() };
978+
979+
let uninit_slice = mmap.as_init();
980+
assert_eq!(uninit_slice, data);
981+
982+
remove_file(path).unwrap();
983+
}
984+
952985
#[test]
953986
fn test_other_reserve() {
954987
let mut buf = [1, 1, 4, 5, 1, 4];

‎compio-driver/Cargo.toml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ pin-project-lite = { workspace = true }
2828
thin-cell = { version = "0.1.1" }
2929
smallvec = { workspace = true, optional = true, features = ["union"] }
3030
synchrony = { workspace = true, features = ["waker_slot"] }
31+
memmap2 = { version = "0.9.9", optional = true }
3132

3233
# Windows specific dependencies
3334
[target.'cfg(windows)'.dependencies]
@@ -88,6 +89,8 @@ io-uring-cqe32 = []
8889
iocp-global = []
8990
iocp-wait-packet = []
9091

92+
buffer-mmap = ["memmap2"]
93+
9194
# Nightly features
9295
once_cell_try = []
9396
proc_macro_diagnostic = ["dep:compile_warning"]

‎compio-driver/src/buffer_pool/fallback.rs‎

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
//! The fallback buffer pool. It is backed by a [`VecDeque`] of [`Vec<u8>`].
1+
//! The fallback buffer pool. It is backed by a [`VecDeque`] of [`BufferBlock`].
22
//! An [`OwnedBuffer`] is selected when the op is created. It keeps a strong
33
//! reference to the buffer pool. The [`BorrowedBuffer`] is created after the op
44
//! returns successfully.
@@ -16,12 +16,14 @@ use std::{
1616

1717
use compio_buf::{IntoInner, IoBuf, IoBufMut, SetLen, Slice};
1818

19+
use super::BufferBlock;
20+
1921
struct BufferPoolInner {
20-
buffers: RefCell<VecDeque<Vec<u8>>>,
22+
buffers: RefCell<VecDeque<BufferBlock>>,
2123
}
2224

2325
impl BufferPoolInner {
24-
pub(crate) fn add_buffer(&self, mut buffer: Vec<u8>) {
26+
pub(crate) fn add_buffer(&self, mut buffer: BufferBlock) {
2527
buffer.clear();
2628
self.buffers.borrow_mut().push_back(buffer)
2729
}
@@ -42,17 +44,17 @@ impl Debug for BufferPool {
4244
}
4345

4446
impl BufferPool {
45-
pub(crate) fn new(buffer_len: u16, buffer_size: usize) -> Self {
47+
pub(crate) fn new(buffer_len: u16, buffer_size: usize) -> io::Result<Self> {
4648
// To match the behavior of io-uring, extend the number of buffers.
4749
let buffers = (0..buffer_len.next_power_of_two())
48-
.map(|_| Vec::with_capacity(buffer_size))
49-
.collect();
50+
.map(|_| BufferBlock::with_capacity(buffer_size))
51+
.collect::<io::Result<_>>()?;
5052

51-
Self {
53+
Ok(Self {
5254
inner: Rc::new(BufferPoolInner {
5355
buffers: RefCell::new(buffers),
5456
}),
55-
}
57+
})
5658
}
5759

5860
/// Select an [`OwnedBuffer`] when the op creates.
@@ -73,7 +75,7 @@ impl BufferPool {
7375
}
7476

7577
/// Return the buffer to the pool.
76-
pub(crate) fn add_buffer(&self, buffer: Vec<u8>) {
78+
pub(crate) fn add_buffer(&self, buffer: BufferBlock) {
7779
self.inner.add_buffer(buffer);
7880
}
7981

@@ -90,12 +92,12 @@ impl BufferPool {
9092

9193
#[doc(hidden)]
9294
pub struct OwnedBuffer {
93-
buffer: ManuallyDrop<Slice<Vec<u8>>>,
95+
buffer: ManuallyDrop<Slice<BufferBlock>>,
9496
pool: ManuallyDrop<Rc<BufferPoolInner>>,
9597
}
9698

9799
impl OwnedBuffer {
98-
fn new(buffer: Slice<Vec<u8>>, pool: Rc<BufferPoolInner>) -> Self {
100+
fn new(buffer: Slice<BufferBlock>, pool: Rc<BufferPoolInner>) -> Self {
99101
Self {
100102
buffer: ManuallyDrop::new(buffer),
101103
pool: ManuallyDrop::new(pool),
@@ -132,7 +134,7 @@ impl Drop for OwnedBuffer {
132134
}
133135

134136
impl IntoInner for OwnedBuffer {
135-
type Inner = Slice<Vec<u8>>;
137+
type Inner = Slice<BufferBlock>;
136138

137139
fn into_inner(mut self) -> Self::Inner {
138140
// SAFETY: `self` is forgotten in this method.
@@ -150,12 +152,12 @@ impl IntoInner for OwnedBuffer {
150152
/// When IO operation finish, user will obtain a `BorrowedBuffer` to access the
151153
/// filled data
152154
pub struct BorrowedBuffer<'a> {
153-
buffer: ManuallyDrop<Slice<Vec<u8>>>,
155+
buffer: ManuallyDrop<Slice<BufferBlock>>,
154156
pool: &'a BufferPool,
155157
}
156158

157159
impl<'a> BorrowedBuffer<'a> {
158-
pub(crate) fn new(buffer: Slice<Vec<u8>>, pool: &'a BufferPool) -> Self {
160+
pub(crate) fn new(buffer: Slice<BufferBlock>, pool: &'a BufferPool) -> Self {
159161
Self {
160162
buffer: ManuallyDrop::new(buffer),
161163
pool,

‎compio-driver/src/buffer_pool/iour.rs‎

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
//! The io-uring buffer pool. It is backed by a [`Vec`] of [`Vec<u8>`].
1+
//! The io-uring buffer pool. It is backed by a [`Vec`] of [`BufferBlock`].
22
//! The kernel selects the buffer and returns `flags`. The crate
33
//! [`io_uring_buf_ring`] handles the returning of buffer on drop.
44
@@ -9,14 +9,28 @@ use std::{
99
ops::{Deref, DerefMut},
1010
};
1111

12-
use io_uring_buf_ring::IoUringBufRing;
12+
use io_uring_buf_ring::{Buffer, IoUringBufRing};
13+
14+
use crate::buffer_pool::BufferBlock;
1315

1416
/// Buffer pool
1517
///
1618
/// A buffer pool to allow user no need to specify a specific buffer to do the
1719
/// IO operation
1820
pub struct BufferPool {
19-
buf_ring: IoUringBufRing<Vec<u8>>,
21+
buf_ring: IoUringBufRing<BufferBlock>,
22+
}
23+
24+
unsafe impl Buffer for BufferBlock {
25+
fn ptr(&self) -> *mut std::mem::MaybeUninit<u8> {
26+
self.as_ptr() as _
27+
}
28+
29+
fn len(&self) -> usize {
30+
self.capacity()
31+
}
32+
33+
fn drop(self) {}
2034
}
2135

2236
impl Debug for BufferPool {
@@ -26,15 +40,15 @@ impl Debug for BufferPool {
2640
}
2741

2842
impl BufferPool {
29-
pub(crate) fn new(buf_ring: IoUringBufRing<Vec<u8>>) -> Self {
43+
pub(crate) fn new(buf_ring: IoUringBufRing<BufferBlock>) -> Self {
3044
Self { buf_ring }
3145
}
3246

3347
pub(crate) fn buffer_group(&self) -> u16 {
3448
self.buf_ring.buffer_group()
3549
}
3650

37-
pub(crate) fn into_inner(self) -> IoUringBufRing<Vec<u8>> {
51+
pub(crate) fn into_inner(self) -> IoUringBufRing<BufferBlock> {
3852
self.buf_ring
3953
}
4054

@@ -61,7 +75,7 @@ impl BufferPool {
6175
///
6276
/// When IO operation finish, user will obtain a `BorrowedBuffer` to access the
6377
/// filled data
64-
pub struct BorrowedBuffer<'a>(io_uring_buf_ring::BorrowedBuffer<'a, Vec<u8>>);
78+
pub struct BorrowedBuffer<'a>(io_uring_buf_ring::BorrowedBuffer<'a, BufferBlock>);
6579

6680
impl Debug for BorrowedBuffer<'_> {
6781
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {

‎compio-driver/src/buffer_pool/mod.rs‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,117 @@ cfg_if::cfg_if! {
1313
}
1414
}
1515

16+
use std::ops::{Deref, DerefMut};
17+
18+
/// Buffer type for the buffer pool.
19+
///
20+
/// Backed by [`Vec`].
21+
#[cfg(not(feature = "buffer-mmap"))]
22+
pub struct BufferBlock(Vec<u8>);
23+
24+
/// Buffer type for the buffer pool.
25+
///
26+
/// Backed by [`memmap2`].
27+
#[cfg(feature = "buffer-mmap")]
28+
pub struct BufferBlock(memmap2::MmapMut);
29+
30+
impl Deref for BufferBlock {
31+
type Target = [u8];
32+
33+
fn deref(&self) -> &Self::Target {
34+
self.0.deref()
35+
}
36+
}
37+
38+
impl DerefMut for BufferBlock {
39+
fn deref_mut(&mut self) -> &mut Self::Target {
40+
self.0.deref_mut()
41+
}
42+
}
43+
44+
#[cfg(not(feature = "buffer-mmap"))]
45+
mod vec_buffer {
46+
use std::io;
47+
48+
use compio_buf::{IoBuf, IoBufMut, ReserveError, ReserveExactError, SetLen};
49+
50+
use super::BufferBlock;
51+
52+
impl BufferBlock {
53+
pub(crate) fn with_capacity(cap: usize) -> io::Result<Self> {
54+
Ok(Self(Vec::with_capacity(cap)))
55+
}
56+
57+
pub(crate) fn capacity(&self) -> usize {
58+
self.0.capacity()
59+
}
60+
}
61+
62+
impl IoBuf for BufferBlock {
63+
fn as_init(&self) -> &[u8] {
64+
self
65+
}
66+
}
67+
68+
impl SetLen for BufferBlock {
69+
unsafe fn set_len(&mut self, len: usize) {
70+
unsafe { self.0.set_len(len) }
71+
}
72+
}
73+
74+
impl IoBufMut for BufferBlock {
75+
fn as_uninit(&mut self) -> &mut [std::mem::MaybeUninit<u8>] {
76+
self.0.as_uninit()
77+
}
78+
79+
fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
80+
IoBufMut::reserve(&mut self.0, len)
81+
}
82+
83+
fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
84+
IoBufMut::reserve_exact(&mut self.0, len)
85+
}
86+
}
87+
}
88+
89+
#[cfg(feature = "buffer-mmap")]
90+
mod mmap_buffer {
91+
use std::io;
92+
93+
use compio_buf::{IoBuf, IoBufMut, SetLen};
94+
use memmap2::MmapMut;
95+
96+
use super::BufferBlock;
97+
98+
impl BufferBlock {
99+
pub(crate) fn with_capacity(cap: usize) -> io::Result<Self> {
100+
MmapMut::map_anon(cap).map(BufferBlock)
101+
}
102+
103+
pub(crate) fn capacity(&self) -> usize {
104+
self.0.len()
105+
}
106+
}
107+
108+
impl IoBuf for BufferBlock {
109+
fn as_init(&self) -> &[u8] {
110+
self
111+
}
112+
}
113+
114+
impl SetLen for BufferBlock {
115+
unsafe fn set_len(&mut self, len: usize) {
116+
assert!(len <= self.len())
117+
}
118+
}
119+
120+
impl IoBufMut for BufferBlock {
121+
fn as_uninit(&mut self) -> &mut [std::mem::MaybeUninit<u8>] {
122+
unsafe { std::mem::transmute(self.as_mut_slice()) }
123+
}
124+
}
125+
}
126+
16127
/// Trait to get the selected buffer of an io operation.
17128
pub trait TakeBuffer {
18129
/// Selected buffer type. It keeps the reference to the buffer pool and

‎compio-driver/src/sys/iocp/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -473,7 +473,7 @@ impl Driver {
473473
buffer_len: u16,
474474
buffer_size: usize,
475475
) -> io::Result<BufferPool> {
476-
Ok(BufferPool::new(buffer_len, buffer_size))
476+
BufferPool::new(buffer_len, buffer_size)
477477
}
478478

479479
/// # Safety

‎compio-driver/src/sys/iour/mod.rs‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ use io_uring::{
3535
use slab::Slab;
3636

3737
use crate::{
38-
AsyncifyPool, BufferPool, DriverType, Entry, ProactorBuilder,
38+
AsyncifyPool, BufferBlock, BufferPool, DriverType, Entry, ProactorBuilder,
3939
key::{ErasedKey, RefExt},
4040
syscall,
4141
};
@@ -422,11 +422,14 @@ impl Driver {
422422
));
423423
}
424424

425-
let buf_ring = io_uring_buf_ring::IoUringBufRing::new_with_flags(
425+
let bufs = (0..buffer_len)
426+
.map(|_| BufferBlock::with_capacity(buffer_size))
427+
.collect::<io::Result<Vec<_>>>()?;
428+
429+
let buf_ring = io_uring_buf_ring::IoUringBufRing::new_with_buffers_and_flags(
426430
&self.inner,
427-
buffer_len,
431+
bufs,
428432
buffer_group as _,
429-
buffer_size,
430433
0,
431434
)?;
432435

‎compio-driver/src/sys/poll/mod.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -672,11 +672,11 @@ impl Driver {
672672
Ok(BufferPool::new_poll(crate::FallbackBufferPool::new(
673673
buffer_len,
674674
buffer_size,
675-
)))
675+
)?))
676676
}
677677
#[cfg(not(fusion))]
678678
{
679-
Ok(BufferPool::new(buffer_len, buffer_size))
679+
BufferPool::new(buffer_len, buffer_size)
680680
}
681681
}
682682

0 commit comments

Comments
 (0)