feat: make wasi-nn immutable

Signed-off-by: Harald Hoyer <harald@profian.com>

Author: haraldh

crates/wasi-nn/src/ctx.rs

@@ -6,14 +6,15 @@ use crate::r#impl::UsageError;
 use crate::witx::types::{Graph, GraphEncoding, GraphExecutionContext};
 use std::collections::HashMap;
 use std::hash::Hash;
+use std::sync::RwLock;
 use thiserror::Error;
 use wiggle::GuestError;
 
 /// Capture the state necessary for calling into the backend ML libraries.
 pub struct WasiNnCtx {
-    pub(crate) backends: HashMap<u8, Box<dyn Backend>>,
-    pub(crate) graphs: Table<Graph, Box<dyn BackendGraph>>,
-    pub(crate) executions: Table<GraphExecutionContext, Box<dyn BackendExecutionContext>>,
+    pub(crate) backends: RwLock<HashMap<u8, Box<dyn Backend>>>,
+    pub(crate) graphs: RwLock<Table<Graph, Box<dyn BackendGraph>>>,
+    pub(crate) executions: RwLock<Table<GraphExecutionContext, Box<dyn BackendExecutionContext>>>,
 }
 
 impl WasiNnCtx {
@@ -27,9 +28,9 @@ impl WasiNnCtx {
             Box::new(OpenvinoBackend::default()) as Box<dyn Backend>,
         );
         Ok(Self {
-            backends,
-            graphs: Table::default(),
-            executions: Table::default(),
+            backends: RwLock::new(backends),
+            graphs: RwLock::new(Table::default()),
+            executions: RwLock::new(Table::default()),
         })
     }
 }

crates/wasi-nn/src/impl.rs

@@ -26,61 +26,61 @@ pub enum UsageError {
 
 impl<'a> WasiEphemeralNn for WasiNnCtx {
     fn load<'b>(
-        &mut self,
+        &self,
         builders: &GraphBuilderArray<'_>,
         encoding: GraphEncoding,
         target: ExecutionTarget,
     ) -> Result<Graph> {
         let encoding_id: u8 = encoding.into();
-        let graph = if let Some(backend) = self.backends.get_mut(&encoding_id) {
+        let graph = if let Some(backend) = self.backends.write().unwrap().get_mut(&encoding_id) {
             backend.load(builders, target)?
         } else {
             return Err(UsageError::InvalidEncoding(encoding).into());
         };
-        let graph_id = self.graphs.insert(graph);
+        let graph_id = self.graphs.write().unwrap().insert(graph);
         Ok(graph_id)
     }
 
-    fn init_execution_context(&mut self, graph_id: Graph) -> Result<GraphExecutionContext> {
-        let exec_context = if let Some(graph) = self.graphs.get_mut(graph_id) {
+    fn init_execution_context(&self, graph_id: Graph) -> Result<GraphExecutionContext> {
+        let exec_context = if let Some(graph) = self.graphs.write().unwrap().get_mut(graph_id) {
             graph.init_execution_context()?
         } else {
             return Err(UsageError::InvalidGraphHandle.into());
         };
 
-        let exec_context_id = self.executions.insert(exec_context);
+        let exec_context_id = self.executions.write().unwrap().insert(exec_context);
         Ok(exec_context_id)
     }
 
     fn set_input<'b>(
-        &mut self,
+        &self,
         exec_context_id: GraphExecutionContext,
         index: u32,
         tensor: &Tensor<'b>,
     ) -> Result<()> {
-        if let Some(exec_context) = self.executions.get_mut(exec_context_id) {
+        if let Some(exec_context) = self.executions.write().unwrap().get_mut(exec_context_id) {
             Ok(exec_context.set_input(index, tensor)?)
         } else {
             Err(UsageError::InvalidGraphHandle.into())
         }
     }
 
-    fn compute(&mut self, exec_context_id: GraphExecutionContext) -> Result<()> {
-        if let Some(exec_context) = self.executions.get_mut(exec_context_id) {
+    fn compute(&self, exec_context_id: GraphExecutionContext) -> Result<()> {
+        if let Some(exec_context) = self.executions.write().unwrap().get_mut(exec_context_id) {
             Ok(exec_context.compute()?)
         } else {
             Err(UsageError::InvalidExecutionContextHandle.into())
         }
     }
 
     fn get_output<'b>(
-        &mut self,
+        &self,
         exec_context_id: GraphExecutionContext,
         index: u32,
         out_buffer: &GuestPtr<'_, u8>,
         out_buffer_max_size: u32,
     ) -> Result<u32> {
-        if let Some(exec_context) = self.executions.get_mut(exec_context_id) {
+        if let Some(exec_context) = self.executions.write().unwrap().get_mut(exec_context_id) {
             let mut destination = out_buffer
                 .as_array(out_buffer_max_size)
                 .as_slice_mut()?