Work on supporting both epoch and fuel

crates/plugin_runtime/src/lib.rs

@@ -23,7 +23,7 @@ mod tests {
         }
 
         async {
-            let mut runtime = PluginBuilder::new_with_default_ctx(PluginYield::default_fuel())
+            let mut runtime = PluginBuilder::new_fuel_with_default_ctx(PluginYield::default_fuel())
                 .unwrap()
                 .host_function("mystery_number", |input: u32| input + 7)
                 .unwrap()

crates/plugin_runtime/src/plugin.rs

@@ -55,27 +55,43 @@ impl<A: Serialize, R: DeserializeOwned> Clone for WasiFn<A, R> {
     }
 }
 
+pub struct PluginYieldEpoch {
+    delta: u64,
+    epoch: std::time::Duration,
+}
+
+impl Into<PluginYield> for PluginYieldEpoch {
+    fn into(self) -> PluginYield {
+        PluginYield::Epoch(self)
+    }
+}
+
+pub struct PluginYieldFuel {
+    initial: u64,
+    refill: u64,
+}
+
+impl Into<PluginYield> for PluginYieldFuel {
+    fn into(self) -> PluginYield {
+        PluginYield::Fuel(self)
+    }
+}
+
 pub enum PluginYield {
-    Epoch {
+    Epoch(PluginYieldEpoch),
-        delta: u64,
+    Fuel(PluginYieldFuel),
-        epoch: std::time::Duration,
-    },
-    Fuel {
-        initial: u64,
-        refill: u64,
-    },
 }
 
 impl PluginYield {
-    pub fn default_epoch() -> Self {
+    pub fn default_epoch() -> PluginYieldEpoch {
-        PluginYield::Epoch {
+        PluginYieldEpoch {
             delta: 1,
             epoch: Duration::from_millis(1),
         }
     }
 
-    pub fn default_fuel() -> Self {
+    pub fn default_fuel() -> PluginYieldFuel {
-        PluginYield::Fuel {
+        PluginYieldFuel {
             initial: 1000,
             refill: 1000,
         }
@@ -91,65 +107,94 @@ pub struct PluginBuilder {
     engine: Engine,
     linker: Linker<WasiCtxAlloc>,
     yield_when: PluginYield,
-    created_epoch_incrementer: bool,
 }
 
 impl PluginBuilder {
-    /// Create a new [`PluginBuilder`] with the given WASI context.
+    fn create_engine(yield_when: &PluginYield) -> Result<(Engine, Linker<WasiCtxAlloc>), Error> {
-    /// Using the default context is a safe bet, see [`new_with_default_context`].
-    pub fn new(wasi_ctx: WasiCtx, yield_when: PluginYield) -> Result<Self, Error> {
         let mut config = Config::default();
         config.async_support(true);
+        let engine = Engine::new(&config)?;
+        let linker = Linker::new(&engine);
 
         match yield_when {
-            PluginYield::Epoch { .. } => {
+            PluginYield::Epoch(_) => {
                 config.epoch_interruption(true);
             }
-            PluginYield::Fuel { .. } => {
+            PluginYield::Fuel(_) => {
                 config.consume_fuel(true);
             }
         }
 
-        let engine = Engine::new(&config)?;
+        Ok((engine, linker))
-        let linker = Linker::new(&engine);
+    }
+
+    /// Create a new [`PluginBuilder`] with the given WASI context.
+    /// Using the default context is a safe bet, see [`new_with_default_context`].
+    pub fn new_epoch<C>(
+        wasi_ctx: WasiCtx,
+        yield_epoch: PluginYieldEpoch,
+        callback: C,
+    ) -> Result<Self, Error>
+    where
+        C: FnOnce(std::pin::Pin<Box<dyn Future<Output = ()> + Send + 'static>>) -> (),
+    {
+        let epoch = yield_epoch.epoch;
+        let yield_when = PluginYield::Epoch(yield_epoch);
+        let (engine, linker) = Self::create_engine(&yield_when)?;
+
+        let engine_ref = &engine;
+
+        // we can't create the future until after initializing
+        // because we need the engine to load the plugin
+        // we could use an Arc, but that'd suck
+        callback(Box::pin(async move {
+            loop {
+                smol::Timer::after(epoch).await;
+                engine_ref.increment_epoch();
+            }
+        }));
 
         Ok(PluginBuilder {
             wasi_ctx,
             engine,
             linker,
             yield_when,
-            created_epoch_incrementer: false,
         })
     }
 
-    /// Create a new `PluginBuilder` that inherits the
+    pub fn new_fuel(wasi_ctx: WasiCtx, yield_fuel: PluginYieldFuel) -> Result<Self, Error> {
-    /// host processes' access to `stdout` and `stderr`.
+        let yield_when = PluginYield::Fuel(yield_fuel);
-    pub fn new_with_default_ctx(yield_when: PluginYield) -> Result<Self, Error> {
+        let (engine, linker) = Self::create_engine(&yield_when)?;
-        let wasi_ctx = WasiCtxBuilder::new()
+
-            .inherit_stdout()
+        Ok(PluginBuilder {
-            .inherit_stderr()
+            wasi_ctx,
-            .build();
+            engine,
-        Self::new(wasi_ctx, yield_when)
+            linker,
+            yield_when,
+        })
     }
 
-    /// Creates a epoch incrementer if this plugin is configured to increment epochs.
+    /// Create a new `WasiCtx` that inherits the
-    /// Will panic if this plugin is not configured to increment epochs.
+    /// host processes' access to `stdout` and `stderr`.
-    pub fn create_epoch_incrementer(
+    fn default_ctx() -> WasiCtx {
-        &mut self,
+        WasiCtxBuilder::new()
-    ) -> std::pin::Pin<Box<dyn Future<Output = ()> + Send + 'static>> {
+            .inherit_stdout()
-        if let PluginYield::Epoch { epoch, .. } = self.yield_when {
+            .inherit_stderr()
-            self.created_epoch_incrementer = true;
+            .build()
-            let epoch = epoch.clone();
+    }
-            let engine = &self.engine;
+
-            Box::pin(async move {
+    pub fn new_epoch_with_default_ctx<C>(
-                loop {
+        yield_epoch: PluginYieldEpoch,
-                    smol::Timer::after(epoch);
+        callback: C,
-                    engine.increment_epoch();
+    ) -> Result<Self, Error>
-                }
+    where
-            })
+        C: FnOnce(std::pin::Pin<Box<dyn Future<Output = ()> + Send + 'static>>) -> (),
-        } else {
+    {
-            panic!("Tried creating an epoch incrementer, but one does not yet exist.")
+        Self::new_epoch(Self::default_ctx(), yield_epoch, callback)
-        }
+    }
+
+    pub fn new_fuel_with_default_ctx(yield_fuel: PluginYieldFuel) -> Result<Self, Error> {
+        Self::new_fuel(Self::default_ctx(), yield_fuel)
     }
 
     /// Add an `async` host function. See [`host_function`] for details.
@@ -297,11 +342,6 @@ impl PluginBuilder {
     /// Will panic if this is plugin uses `PluginYield::Epoch`,
     /// but an epoch incrementer has not yet been created.
     pub async fn init<T: AsRef<[u8]>>(self, precompiled: bool, module: T) -> Result<Plugin, Error> {
-        if let PluginYield::Epoch { .. } = self.yield_when {
-            if !self.created_epoch_incrementer {
-                panic!("Must create epoch incrementer to run epoch-based plugin");
-            }
-        }
         Plugin::init(precompiled, module.as_ref().to_vec(), self).await
     }
 }
@@ -390,10 +430,10 @@ impl Plugin {
 
         // set up automatic yielding based on configuration
         match plugin.yield_when {
-            PluginYield::Epoch { delta, .. } => {
+            PluginYield::Epoch(PluginYieldEpoch { delta, .. }) => {
                 store.epoch_deadline_async_yield_and_update(delta);
             }
-            PluginYield::Fuel { initial, refill } => {
+            PluginYield::Fuel(PluginYieldFuel { initial, refill }) => {
                 store.add_fuel(initial).unwrap();
                 store.out_of_fuel_async_yield(u64::MAX, refill);
             }

crates/zed/src/languages/language_plugin.rs

@@ -9,7 +9,10 @@ use std::{any::Any, path::PathBuf, sync::Arc};
 use util::ResultExt;
 
 pub async fn new_json(executor: Arc<Background>) -> Result<PluginLspAdapter> {
-    let plugin = PluginBuilder::new_with_default_ctx(PluginYield::default_epoch())?
+    let plugin =
+        PluginBuilder::new_epoch_with_default_ctx(PluginYield::default_epoch(), |future| {
+            executor.spawn(future).detach()
+        })?
         .host_function_async("command", |command: String| async move {
             let mut args = command.split(' ');
             let command = args.next().unwrap();