review: work around fd_fdstat_set_flags

In order to make progress with wasi-threads, this change temporarily
works around limitations induced by `wasi-common`'s
`fd_fdstat_set_flags` to allow `&mut self` use in the implementation.
Eventual resolution is tracked in
bytecodealliance#5643. This change
makes several related helper functions (e.g., `set_fdflags`) take `&mut
self` as well.

Co-authored-by: Alex Crichton <alex@alexcrichton.com>

crates/wasi-common/Cargo.toml

@@ -26,6 +26,7 @@ tracing = { workspace = true }
 cap-std = { workspace = true }
 cap-rand = { workspace = true }
 bitflags = { workspace = true }
+log = { workspace = true }
 
 [target.'cfg(unix)'.dependencies]
 rustix = { workspace = true, features = ["fs"] }

crates/wasi-common/cap-std-sync/src/file.rs

@@ -93,15 +93,15 @@ impl WasiFile for File {
         let fdflags = get_fd_flags(&*file)?;
         Ok(fdflags)
     }
-    async fn set_fdflags(&self, fdflags: FdFlags) -> Result<(), Error> {
+    async fn set_fdflags(&mut self, fdflags: FdFlags) -> Result<(), Error> {
         if fdflags.intersects(
             wasi_common::file::FdFlags::DSYNC
                 | wasi_common::file::FdFlags::SYNC
                 | wasi_common::file::FdFlags::RSYNC,
         ) {
             return Err(Error::invalid_argument().context("cannot set DSYNC, SYNC, or RSYNC flag"));
         }
-        let mut file = self.0.write().unwrap();
+        let file = self.0.get_mut().unwrap();
         let set_fd_flags = (*file).new_set_fd_flags(to_sysif_fdflags(fdflags))?;
         (*file).set_fd_flags(set_fd_flags)?;
         Ok(())

crates/wasi-common/cap-std-sync/src/net.rs

@@ -98,7 +98,7 @@ macro_rules! wasi_listen_write_impl {
             }
             async fn sock_accept(&self, fdflags: FdFlags) -> Result<Box<dyn WasiFile>, Error> {
                 let (stream, _) = self.0.accept()?;
-                let stream = <$stream>::from_cap_std(stream);
+                let mut stream = <$stream>::from_cap_std(stream);
                 stream.set_fdflags(fdflags).await?;
                 Ok(Box::new(stream))
             }
@@ -110,7 +110,7 @@ macro_rules! wasi_listen_write_impl {
                 let fdflags = get_fd_flags(&self.0)?;
                 Ok(fdflags)
             }
-            async fn set_fdflags(&self, fdflags: FdFlags) -> Result<(), Error> {
+            async fn set_fdflags(&mut self, fdflags: FdFlags) -> Result<(), Error> {
                 if fdflags == wasi_common::file::FdFlags::NONBLOCK {
                     self.0.set_nonblocking(true)?;
                 } else if fdflags.is_empty() {
@@ -197,7 +197,7 @@ macro_rules! wasi_stream_write_impl {
                 let fdflags = get_fd_flags(&self.0)?;
                 Ok(fdflags)
             }
-            async fn set_fdflags(&self, fdflags: FdFlags) -> Result<(), Error> {
+            async fn set_fdflags(&mut self, fdflags: FdFlags) -> Result<(), Error> {
                 if fdflags == wasi_common::file::FdFlags::NONBLOCK {
                     self.0.set_nonblocking(true)?;
                 } else if fdflags.is_empty() {

crates/wasi-common/src/ctx.rs

@@ -2,17 +2,26 @@ use crate::clocks::WasiClocks;
 use crate::dir::{DirCaps, DirEntry, WasiDir};
 use crate::file::{FileCaps, FileEntry, WasiFile};
 use crate::sched::WasiSched;
-use crate::string_array::{StringArray, StringArrayError};
+use crate::string_array::StringArray;
 use crate::table::Table;
-use crate::Error;
+use crate::{Error, StringArrayError};
 use cap_rand::RngCore;
+use std::ops::Deref;
 use std::path::{Path, PathBuf};
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
-pub struct WasiCtx {
+/// An `Arc`-wrapper around the wasi-common context to allow mutable access to
+/// the file descriptor table. This wrapper is only necessary due to the
+/// signature of `fd_fdstat_set_flags`; if that changes, there are a variety of
+/// improvements that can be made (TODO:
+/// https://github.com/bytecodealliance/wasmtime/issues/5643).
+#[derive(Clone)]
+pub struct WasiCtx(Arc<WasiCtxInner>);
+
+pub struct WasiCtxInner {
     pub args: StringArray,
     pub env: StringArray,
-    pub random: Box<dyn RngCore + Send + Sync>,
+    pub random: Mutex<Box<dyn RngCore + Send + Sync>>,
     pub clocks: WasiClocks,
     pub sched: Box<dyn WasiSched>,
     pub table: Table,
@@ -25,14 +34,14 @@ impl WasiCtx {
         sched: Box<dyn WasiSched>,
         table: Table,
     ) -> Self {
-        let s = WasiCtx {
+        let s = WasiCtx(Arc::new(WasiCtxInner {
             args: StringArray::new(),
             env: StringArray::new(),
-            random,
+            random: Mutex::new(random),
             clocks,
             sched,
             table,
-        };
+        }));
         s.set_stdin(Box::new(crate::pipe::ReadPipe::new(std::io::empty())));
         s.set_stdout(Box::new(crate::pipe::WritePipe::new(std::io::sink())));
         s.set_stderr(Box::new(crate::pipe::WritePipe::new(std::io::sink())));
@@ -77,12 +86,22 @@ impl WasiCtx {
         &self.table
     }
 
+    pub fn table_mut(&mut self) -> Option<&mut Table> {
+        Arc::get_mut(&mut self.0).map(|c| &mut c.table)
+    }
+
     pub fn push_arg(&mut self, arg: &str) -> Result<(), StringArrayError> {
-        self.args.push(arg.to_owned())
+        let s = Arc::get_mut(&mut self.0).expect(
+            "`push_arg` should only be used during initialization before the context is cloned",
+        );
+        s.args.push(arg.to_owned())
     }
 
     pub fn push_env(&mut self, var: &str, value: &str) -> Result<(), StringArrayError> {
-        self.env.push(format!("{}={}", var, value))?;
+        let s = Arc::get_mut(&mut self.0).expect(
+            "`push_env` should only be used during initialization before the context is cloned",
+        );
+        s.env.push(format!("{}={}", var, value))?;
         Ok(())
     }
 
@@ -130,3 +149,10 @@ impl WasiCtx {
         Ok(())
     }
 }
+
+impl Deref for WasiCtx {
+    type Target = WasiCtxInner;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}

crates/wasi-common/src/file.rs

@@ -64,7 +64,7 @@ pub trait WasiFile: Send + Sync {
         Ok(FdFlags::empty())
     }
 
-    async fn set_fdflags(&self, _flags: FdFlags) -> Result<(), Error> {
+    async fn set_fdflags(&mut self, _flags: FdFlags) -> Result<(), Error> {
         Err(Error::badf())
     }
 
@@ -217,11 +217,15 @@ pub struct Filestat {
 
 pub(crate) trait TableFileExt {
     fn get_file(&self, fd: u32) -> Result<Arc<FileEntry>, Error>;
+    fn get_file_mut(&mut self, fd: u32) -> Result<&mut FileEntry, Error>;
 }
 impl TableFileExt for crate::table::Table {
     fn get_file(&self, fd: u32) -> Result<Arc<FileEntry>, Error> {
         self.get(fd)
     }
+    fn get_file_mut(&mut self, fd: u32) -> Result<&mut FileEntry, Error> {
+        self.get_mut(fd)
+    }
 }
 
 pub(crate) struct FileEntry {
@@ -272,13 +276,18 @@ impl FileEntry {
 
 pub trait FileEntryExt {
     fn get_cap(&self, caps: FileCaps) -> Result<&dyn WasiFile, Error>;
+    fn get_cap_mut(&mut self, caps: FileCaps) -> Result<&mut dyn WasiFile, Error>;
 }
 
 impl FileEntryExt for FileEntry {
     fn get_cap(&self, caps: FileCaps) -> Result<&dyn WasiFile, Error> {
         self.capable_of(caps)?;
         Ok(&*self.file)
     }
+    fn get_cap_mut(&mut self, caps: FileCaps) -> Result<&mut dyn WasiFile, Error> {
+        self.capable_of(caps)?;
+        Ok(&mut *self.file)
+    }
 }
 
 bitflags! {

crates/wasi-common/src/snapshots/preview_1.rs

@@ -189,11 +189,16 @@ impl wasi_snapshot_preview1::WasiSnapshotPreview1 for WasiCtx {
         fd: types::Fd,
         flags: types::Fdflags,
     ) -> Result<(), Error> {
-        self.table()
-            .get_file(u32::from(fd))?
-            .get_cap(FileCaps::FDSTAT_SET_FLAGS)?
-            .set_fdflags(FdFlags::from(flags))
-            .await
+        if let Some(table) = self.table_mut() {
+            table
+                .get_file_mut(u32::from(fd))?
+                .get_cap_mut(FileCaps::FDSTAT_SET_FLAGS)?
+                .set_fdflags(FdFlags::from(flags))
+                .await
+        } else {
+            log::warn!("`fd_fdstat_set_flags` does not work with wasi-threads enabled; see https://github.com/bytecodealliance/wasmtime/issues/5643");
+            Err(Error::invalid_argument())
+        }
     }
 
     async fn fd_fdstat_set_rights(
@@ -1110,7 +1115,7 @@ impl wasi_snapshot_preview1::WasiSnapshotPreview1 for WasiCtx {
             while copied < buf.len() {
                 let len = (buf.len() - copied).min(MAX_SHARED_BUFFER_SIZE as u32);
                 let mut tmp = vec![0; len as usize];
-                self.random.try_fill_bytes(&mut tmp)?;
+                self.random.lock().unwrap().try_fill_bytes(&mut tmp)?;
                 let dest = buf
                     .get_range(copied..copied + len)
                     .unwrap()
@@ -1122,7 +1127,7 @@ impl wasi_snapshot_preview1::WasiSnapshotPreview1 for WasiCtx {
             // If the Wasm memory is non-shared, copy directly into the linear
             // memory.
             let mem = &mut buf.as_slice_mut()?.unwrap();
-            self.random.try_fill_bytes(mem)?;
+            self.random.lock().unwrap().try_fill_bytes(mem)?;
         }
         Ok(())
     }

crates/wasi-common/src/table.rs

@@ -76,6 +76,22 @@ impl Table {
         }
     }
 
+    /// Get a mutable reference to a resource of a given type at a given index.
+    /// Only one such reference can be borrowed at any given time.
+    pub fn get_mut<T: Any>(&mut self, key: u32) -> Result<&mut T, Error> {
+        let entry = match self.0.get_mut().unwrap().map.get_mut(&key) {
+            Some(entry) => entry,
+            None => return Err(Error::badf().context("key not in table")),
+        };
+        let entry = match Arc::get_mut(entry) {
+            Some(entry) => entry,
+            None => return Err(Error::badf().context("cannot mutably borrow shared file")),
+        };
+        entry
+            .downcast_mut::<T>()
+            .ok_or_else(|| Error::badf().context("element is a different type"))
+    }
+
     /// Remove a resource at a given index from the table. Returns the resource
     /// if it was present.
     pub fn delete<T: Any + Send + Sync>(&self, key: u32) -> Option<Arc<T>> {

crates/wasi-common/tokio/src/file.rs

@@ -116,7 +116,7 @@ macro_rules! wasi_file_impl {
             async fn get_fdflags(&self) -> Result<FdFlags, Error> {
                 block_on_dummy_executor(|| self.0.get_fdflags())
             }
-            async fn set_fdflags(&self, fdflags: FdFlags) -> Result<(), Error> {
+            async fn set_fdflags(&mut self, fdflags: FdFlags) -> Result<(), Error> {
                 block_on_dummy_executor(|| self.0.set_fdflags(fdflags))
             }
             async fn get_filestat(&self) -> Result<Filestat, Error> {