[wasi-common] add custom FdPool container for managing fd allocs/deallocs (#1329)
* Rename FdEntry to Entry
* Add custom FdSet container for managing fd allocs/deallocs
This commit adds a custom `FdSet` container which is intended
for use in `wasi-common` to track WASI fd allocs/deallocs. The
main aim for this container is to abstract away the current
approach of spawning new handles
```rust
fd = fd.checked_add(1).ok_or(...)?;
```
and to make it possible to reuse unused/reclaimed handles
which currently is not done.
The struct offers 3 methods to manage its functionality:
* `FdSet::new` initialises the internal data structures,
and most notably, it preallocates an `FdSet::BATCH_SIZE`
worth of handles in such a way that we always start popping
from the "smallest" handle (think of it as of reversed stack,
I guess; it's not a binary heap since we don't really care
whether internally the handles are sorted in some way, just that
the "largets" handle is at the bottom. Why will become clear
when describing `allocate` method.)
* `FdSet::allocate` pops the next available handle if one is available.
The tricky bit here is that, if we run out of handles, we preallocate
the next `FdSet::BATCH_SIZE` worth of handles starting from the
latest popped handle (i.e., the "largest" handle). This
works only because we make sure to only ever pop and push already
existing handles from the back, and push _new_ handles (from the
preallocation step) from the front. When we ultimately run out
of _all_ available handles, we then return `None` for the client
to handle in some way (e.g., throwing an error such as `WasiError::EMFILE`
or whatnot).
* `FdSet::deallocate` returns the already allocated handle back to
the pool for further reuse.
When figuring out the internals, I've tried to optimise for both
alloc and dealloc performance, and I believe we've got an amortised
`O(1)~*` performance for both (if my maths is right, and it may very
well not be, so please verify!).
In order to keep `FdSet` fairly generic, I've made sure not to hard-code
it for the current type system generated by `wig` (i.e., `wasi::__wasi_fd_t`
representing WASI handle), but rather, any type which wants to be managed
by `FdSet` needs to conform to `Fd` trait. This trait is quite simple as
it only requires a couple of rudimentary traits (although `std:#️⃣:Hash`
is quite a powerful assumption here!), and a custom method
```rust
Fd::next(&self) -> Option<Self>;
```
which is there to encapsulate creating another handle from the given one.
In the current state of the code, that'd be simply `u32::checked_add(1)`.
When `wiggle` makes it way into the `wasi-common`, I'd imagine it being
similar to
```rust
fn next(&self) -> Option<Self> {
self.0.checked_add(1).map(Self::from)
}
```
Anyhow, I'd be happy to learn your thoughts about this design!
* Fix compilation on other targets
* Rename FdSet to FdPool
* Fix FdPool unit tests
* Skip preallocation step in FdPool
* Replace 'replace' calls with direct assignment
* Reuse FdPool from snapshot1 in snapshot0
* Refactor FdPool::allocate
* Remove entry before deallocating the fd
* Refactor the design to accommodate `u32` as underlying type
This commit refactors the design by ensuring that the underlying
type in `FdPool` which we use to track and represent raw file
descriptors is `u32`. As a result, the structure of `FdPool` is
simplified massively as we no longer need to track the claimed
descriptors; in a way, we trust the caller to return the handle
after it's done with it. In case the caller decides to be clever
and return a handle which was not yet legally allocated, we panic.
This should never be a problem in `wasi-common` unless we hit a
bug.
To make all of this work, `Fd` trait is modified to require two
methods: `as_raw(&self) -> u32` and `from_raw(raw_fd: u32) -> Self`
both of which are used to convert to and from the `FdPool`'s underlying
type `u32`.
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
use crate::fdentry::{Descriptor, FdEntry};
|
||||
use crate::sys::fdentry_impl::OsHandle;
|
||||
use crate::entry::{Descriptor, Entry};
|
||||
use crate::fdpool::FdPool;
|
||||
use crate::sys::entry_impl::OsHandle;
|
||||
use crate::virtfs::{VirtualDir, VirtualDirEntry};
|
||||
use crate::wasi::{self, WasiError, WasiResult};
|
||||
use std::borrow::Borrow;
|
||||
@@ -37,18 +38,18 @@ pub enum WasiCtxBuilderError {
|
||||
|
||||
type WasiCtxBuilderResult<T> = std::result::Result<T, WasiCtxBuilderError>;
|
||||
|
||||
enum PendingFdEntry {
|
||||
Thunk(fn() -> io::Result<FdEntry>),
|
||||
enum PendingEntry {
|
||||
Thunk(fn() -> io::Result<Entry>),
|
||||
File(File),
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for PendingFdEntry {
|
||||
impl std::fmt::Debug for PendingEntry {
|
||||
fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Thunk(f) => write!(
|
||||
fmt,
|
||||
"PendingFdEntry::Thunk({:p})",
|
||||
f as *const fn() -> io::Result<FdEntry>
|
||||
f as *const fn() -> io::Result<Entry>
|
||||
),
|
||||
Self::File(f) => write!(fmt, "PendingFdEntry::File({:?})", f),
|
||||
}
|
||||
@@ -102,7 +103,9 @@ impl PendingCString {
|
||||
|
||||
/// A builder allowing customizable construction of `WasiCtx` instances.
|
||||
pub struct WasiCtxBuilder {
|
||||
fds: Option<HashMap<wasi::__wasi_fd_t, PendingFdEntry>>,
|
||||
stdin: Option<PendingEntry>,
|
||||
stdout: Option<PendingEntry>,
|
||||
stderr: Option<PendingEntry>,
|
||||
preopens: Option<Vec<(PathBuf, Descriptor)>>,
|
||||
args: Option<Vec<PendingCString>>,
|
||||
env: Option<HashMap<PendingCString, PendingCString>>,
|
||||
@@ -111,14 +114,14 @@ pub struct WasiCtxBuilder {
|
||||
impl WasiCtxBuilder {
|
||||
/// Builder for a new `WasiCtx`.
|
||||
pub fn new() -> Self {
|
||||
let mut fds = HashMap::new();
|
||||
|
||||
fds.insert(0, PendingFdEntry::Thunk(FdEntry::null));
|
||||
fds.insert(1, PendingFdEntry::Thunk(FdEntry::null));
|
||||
fds.insert(2, PendingFdEntry::Thunk(FdEntry::null));
|
||||
let stdin = Some(PendingEntry::Thunk(Entry::null));
|
||||
let stdout = Some(PendingEntry::Thunk(Entry::null));
|
||||
let stderr = Some(PendingEntry::Thunk(Entry::null));
|
||||
|
||||
Self {
|
||||
fds: Some(fds),
|
||||
stdin,
|
||||
stdout,
|
||||
stderr,
|
||||
preopens: Some(Vec::new()),
|
||||
args: Some(Vec::new()),
|
||||
env: Some(HashMap::new()),
|
||||
@@ -160,37 +163,27 @@ impl WasiCtxBuilder {
|
||||
|
||||
/// Inherit stdin from the host process.
|
||||
pub fn inherit_stdin(&mut self) -> &mut Self {
|
||||
self.fds
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.insert(0, PendingFdEntry::Thunk(FdEntry::duplicate_stdin));
|
||||
self.stdin = Some(PendingEntry::Thunk(Entry::duplicate_stdin));
|
||||
self
|
||||
}
|
||||
|
||||
/// Inherit stdout from the host process.
|
||||
pub fn inherit_stdout(&mut self) -> &mut Self {
|
||||
self.fds
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.insert(1, PendingFdEntry::Thunk(FdEntry::duplicate_stdout));
|
||||
self.stdout = Some(PendingEntry::Thunk(Entry::duplicate_stdout));
|
||||
self
|
||||
}
|
||||
|
||||
/// Inherit stdout from the host process.
|
||||
pub fn inherit_stderr(&mut self) -> &mut Self {
|
||||
self.fds
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.insert(2, PendingFdEntry::Thunk(FdEntry::duplicate_stderr));
|
||||
self.stderr = Some(PendingEntry::Thunk(Entry::duplicate_stderr));
|
||||
self
|
||||
}
|
||||
|
||||
/// Inherit the stdin, stdout, and stderr streams from the host process.
|
||||
pub fn inherit_stdio(&mut self) -> &mut Self {
|
||||
let fds = self.fds.as_mut().unwrap();
|
||||
fds.insert(0, PendingFdEntry::Thunk(FdEntry::duplicate_stdin));
|
||||
fds.insert(1, PendingFdEntry::Thunk(FdEntry::duplicate_stdout));
|
||||
fds.insert(2, PendingFdEntry::Thunk(FdEntry::duplicate_stderr));
|
||||
self.stdin = Some(PendingEntry::Thunk(Entry::duplicate_stdin));
|
||||
self.stdout = Some(PendingEntry::Thunk(Entry::duplicate_stdout));
|
||||
self.stderr = Some(PendingEntry::Thunk(Entry::duplicate_stderr));
|
||||
self
|
||||
}
|
||||
|
||||
@@ -234,28 +227,19 @@ impl WasiCtxBuilder {
|
||||
|
||||
/// Provide a File to use as stdin
|
||||
pub fn stdin(&mut self, file: File) -> &mut Self {
|
||||
self.fds
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.insert(0, PendingFdEntry::File(file));
|
||||
self.stdin = Some(PendingEntry::File(file));
|
||||
self
|
||||
}
|
||||
|
||||
/// Provide a File to use as stdout
|
||||
pub fn stdout(&mut self, file: File) -> &mut Self {
|
||||
self.fds
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.insert(1, PendingFdEntry::File(file));
|
||||
self.stdout = Some(PendingEntry::File(file));
|
||||
self
|
||||
}
|
||||
|
||||
/// Provide a File to use as stderr
|
||||
pub fn stderr(&mut self, file: File) -> &mut Self {
|
||||
self.fds
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.insert(2, PendingFdEntry::File(file));
|
||||
self.stderr = Some(PendingEntry::File(file));
|
||||
self
|
||||
}
|
||||
|
||||
@@ -338,28 +322,33 @@ impl WasiCtxBuilder {
|
||||
})
|
||||
.collect::<WasiCtxBuilderResult<Vec<CString>>>()?;
|
||||
|
||||
let mut fds: HashMap<wasi::__wasi_fd_t, FdEntry> = HashMap::new();
|
||||
// Populate the non-preopen fds.
|
||||
for (fd, pending) in self.fds.take().unwrap() {
|
||||
let mut fd_pool = FdPool::new();
|
||||
let mut entries: HashMap<wasi::__wasi_fd_t, Entry> = HashMap::new();
|
||||
// Populate the non-preopen entries.
|
||||
for pending in vec![
|
||||
self.stdin.take().unwrap(),
|
||||
self.stdout.take().unwrap(),
|
||||
self.stderr.take().unwrap(),
|
||||
] {
|
||||
let fd = fd_pool
|
||||
.allocate()
|
||||
.ok_or(WasiCtxBuilderError::TooManyFilesOpen)?;
|
||||
log::debug!("WasiCtx inserting ({:?}, {:?})", fd, pending);
|
||||
match pending {
|
||||
PendingFdEntry::Thunk(f) => {
|
||||
fds.insert(fd, f()?);
|
||||
PendingEntry::Thunk(f) => {
|
||||
entries.insert(fd, f()?);
|
||||
}
|
||||
PendingFdEntry::File(f) => {
|
||||
fds.insert(fd, FdEntry::from(Descriptor::OsHandle(OsHandle::from(f)))?);
|
||||
PendingEntry::File(f) => {
|
||||
entries.insert(fd, Entry::from(Descriptor::OsHandle(OsHandle::from(f)))?);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Then add the preopen fds. Startup code in the guest starts looking at fd 3 for preopens,
|
||||
// so we start from there. This variable is initially 2, though, because the loop
|
||||
// immediately does the increment and check for overflow.
|
||||
let mut preopen_fd: wasi::__wasi_fd_t = 2;
|
||||
// Then add the preopen entries.
|
||||
for (guest_path, dir) in self.preopens.take().unwrap() {
|
||||
// We do the increment at the beginning of the loop body, so that we don't overflow
|
||||
// unnecessarily if we have exactly the maximum number of file descriptors.
|
||||
preopen_fd = preopen_fd
|
||||
.checked_add(1)
|
||||
let preopen_fd = fd_pool
|
||||
.allocate()
|
||||
.ok_or(WasiCtxBuilderError::TooManyFilesOpen)?;
|
||||
|
||||
match &dir {
|
||||
@@ -378,27 +367,26 @@ impl WasiCtxBuilder {
|
||||
}
|
||||
}
|
||||
|
||||
// We don't currently allow setting file descriptors other than 0-2, but this will avoid
|
||||
// collisions if we restore that functionality in the future.
|
||||
while fds.contains_key(&preopen_fd) {
|
||||
preopen_fd = preopen_fd
|
||||
.checked_add(1)
|
||||
.ok_or(WasiCtxBuilderError::TooManyFilesOpen)?;
|
||||
}
|
||||
let mut fe = FdEntry::from(dir)?;
|
||||
let mut fe = Entry::from(dir)?;
|
||||
fe.preopen_path = Some(guest_path);
|
||||
log::debug!("WasiCtx inserting ({:?}, {:?})", preopen_fd, fe);
|
||||
fds.insert(preopen_fd, fe);
|
||||
log::debug!("WasiCtx fds = {:?}", fds);
|
||||
entries.insert(preopen_fd, fe);
|
||||
log::debug!("WasiCtx entries = {:?}", entries);
|
||||
}
|
||||
|
||||
Ok(WasiCtx { args, env, fds })
|
||||
Ok(WasiCtx {
|
||||
args,
|
||||
env,
|
||||
entries,
|
||||
fd_pool,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct WasiCtx {
|
||||
fds: HashMap<wasi::__wasi_fd_t, FdEntry>,
|
||||
fd_pool: FdPool,
|
||||
entries: HashMap<wasi::__wasi_fd_t, Entry>,
|
||||
pub(crate) args: Vec<CString>,
|
||||
pub(crate) env: Vec<CString>,
|
||||
}
|
||||
@@ -420,53 +408,42 @@ impl WasiCtx {
|
||||
}
|
||||
|
||||
/// Check if `WasiCtx` contains the specified raw WASI `fd`.
|
||||
pub(crate) unsafe fn contains_fd_entry(&self, fd: wasi::__wasi_fd_t) -> bool {
|
||||
self.fds.contains_key(&fd)
|
||||
pub(crate) unsafe fn contains_entry(&self, fd: wasi::__wasi_fd_t) -> bool {
|
||||
self.entries.contains_key(&fd)
|
||||
}
|
||||
|
||||
/// Get an immutable `FdEntry` corresponding to the specified raw WASI `fd`.
|
||||
pub(crate) unsafe fn get_fd_entry(&self, fd: wasi::__wasi_fd_t) -> WasiResult<&FdEntry> {
|
||||
self.fds.get(&fd).ok_or(WasiError::EBADF)
|
||||
/// Get an immutable `Entry` corresponding to the specified raw WASI `fd`.
|
||||
pub(crate) unsafe fn get_entry(&self, fd: wasi::__wasi_fd_t) -> WasiResult<&Entry> {
|
||||
self.entries.get(&fd).ok_or(WasiError::EBADF)
|
||||
}
|
||||
|
||||
/// Get a mutable `FdEntry` corresponding to the specified raw WASI `fd`.
|
||||
pub(crate) unsafe fn get_fd_entry_mut(
|
||||
&mut self,
|
||||
fd: wasi::__wasi_fd_t,
|
||||
) -> WasiResult<&mut FdEntry> {
|
||||
self.fds.get_mut(&fd).ok_or(WasiError::EBADF)
|
||||
/// Get a mutable `Entry` corresponding to the specified raw WASI `fd`.
|
||||
pub(crate) unsafe fn get_entry_mut(&mut self, fd: wasi::__wasi_fd_t) -> WasiResult<&mut Entry> {
|
||||
self.entries.get_mut(&fd).ok_or(WasiError::EBADF)
|
||||
}
|
||||
|
||||
/// Insert the specified `FdEntry` into the `WasiCtx` object.
|
||||
/// Insert the specified `Entry` into the `WasiCtx` object.
|
||||
///
|
||||
/// The `FdEntry` will automatically get another free raw WASI `fd` assigned. Note that
|
||||
/// the two subsequent free raw WASI `fd`s do not have to be stored contiguously.
|
||||
pub(crate) fn insert_fd_entry(&mut self, fe: FdEntry) -> WasiResult<wasi::__wasi_fd_t> {
|
||||
// Never insert where stdio handles are expected to be.
|
||||
let mut fd = 3;
|
||||
while self.fds.contains_key(&fd) {
|
||||
if let Some(next_fd) = fd.checked_add(1) {
|
||||
fd = next_fd;
|
||||
} else {
|
||||
return Err(WasiError::EMFILE);
|
||||
}
|
||||
}
|
||||
self.fds.insert(fd, fe);
|
||||
pub(crate) fn insert_entry(&mut self, fe: Entry) -> WasiResult<wasi::__wasi_fd_t> {
|
||||
let fd = self.fd_pool.allocate().ok_or(WasiError::EMFILE)?;
|
||||
self.entries.insert(fd, fe);
|
||||
Ok(fd)
|
||||
}
|
||||
|
||||
/// Insert the specified `FdEntry` with the specified raw WASI `fd` key into the `WasiCtx`
|
||||
/// Insert the specified `Entry` with the specified raw WASI `fd` key into the `WasiCtx`
|
||||
/// object.
|
||||
pub(crate) fn insert_fd_entry_at(
|
||||
&mut self,
|
||||
fd: wasi::__wasi_fd_t,
|
||||
fe: FdEntry,
|
||||
) -> Option<FdEntry> {
|
||||
self.fds.insert(fd, fe)
|
||||
pub(crate) fn insert_entry_at(&mut self, fd: wasi::__wasi_fd_t, fe: Entry) -> Option<Entry> {
|
||||
self.entries.insert(fd, fe)
|
||||
}
|
||||
|
||||
/// Remove `FdEntry` corresponding to the specified raw WASI `fd` from the `WasiCtx` object.
|
||||
pub(crate) fn remove_fd_entry(&mut self, fd: wasi::__wasi_fd_t) -> WasiResult<FdEntry> {
|
||||
self.fds.remove(&fd).ok_or(WasiError::EBADF)
|
||||
/// Remove `Entry` corresponding to the specified raw WASI `fd` from the `WasiCtx` object.
|
||||
pub(crate) fn remove_entry(&mut self, fd: wasi::__wasi_fd_t) -> WasiResult<Entry> {
|
||||
// Remove the `fd` from valid entries.
|
||||
let entry = self.entries.remove(&fd).ok_or(WasiError::EBADF)?;
|
||||
// Next, deallocate the `fd`.
|
||||
self.fd_pool.deallocate(fd);
|
||||
Ok(entry)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user