From 7f11a3236412e9e44cfa5c34df6ffeda0d29ba69 Mon Sep 17 00:00:00 2001
From: Isaac Clayton <slightknack@gmail.com>
Date: Tue, 12 Jul 2022 17:07:33 +0200
Subject: [PATCH] Implement periodic yielding using
 epoch_deadline_async_yield_and_update

---
 crates/plugin_runtime/src/lib.rs            |  4 +-
 crates/plugin_runtime/src/plugin.rs         | 46 +++++++++++++++------
 crates/zed/src/languages/language_plugin.rs |  4 +-
 3 files changed, 39 insertions(+), 15 deletions(-)

diff --git a/crates/plugin_runtime/src/lib.rs b/crates/plugin_runtime/src/lib.rs
index 555d925874..c4c145d801 100644
--- a/crates/plugin_runtime/src/lib.rs
+++ b/crates/plugin_runtime/src/lib.rs
@@ -23,7 +23,7 @@ mod tests {
         }
 
         async {
-            let mut runtime = PluginBuilder::new_with_default_ctx()
+            let (mut runtime, incrementer) = PluginBuilder::new_with_default_ctx()
                 .unwrap()
                 .host_function("mystery_number", |input: u32| input + 7)
                 .unwrap()
@@ -53,6 +53,8 @@ mod tests {
                 .await
                 .unwrap();
 
+            std::thread::spawn(move || incrementer.block_on());
+
             let plugin = TestPlugin {
                 noop: runtime.function("noop").unwrap(),
                 constant: runtime.function("constant").unwrap(),
diff --git a/crates/plugin_runtime/src/plugin.rs b/crates/plugin_runtime/src/plugin.rs
index 6e0764d1a7..dbc6a539df 100644
--- a/crates/plugin_runtime/src/plugin.rs
+++ b/crates/plugin_runtime/src/plugin.rs
@@ -62,6 +62,8 @@ pub struct PluginBuilder {
     wasi_ctx: WasiCtx,
     engine: Engine,
     linker: Linker<WasiCtxAlloc>,
+    delta: u64,
+    epoch: std::time::Duration,
 }
 
 /// Creates a default engine for compiling Wasm.
@@ -83,10 +85,11 @@ impl PluginBuilder {
         let linker = Linker::new(&engine);
 
         Ok(PluginBuilder {
-            // host_functions: HashMap::new(),
             wasi_ctx,
             engine,
             linker,
+            delta: 1,
+            epoch: std::time::Duration::from_millis(100),
         })
     }
 
@@ -242,7 +245,17 @@ impl PluginBuilder {
 
     /// Initializes a [`Plugin`] from a given compiled Wasm module.
     /// Both binary (`.wasm`) and text (`.wat`) module formats are supported.
-    pub async fn init<T: AsRef<[u8]>>(self, precompiled: bool, module: T) -> Result<Plugin, Error> {
+    pub async fn init<T: AsRef<[u8]>>(
+        self,
+        precompiled: bool,
+        module: T,
+    ) -> Result<
+        (
+            Plugin,
+            std::pin::Pin<Box<dyn Future<Output = ()> + Send + 'static>>,
+        ),
+        Error,
+    > {
         Plugin::init(precompiled, module.as_ref().to_vec(), self).await
     }
 }
@@ -303,12 +316,17 @@ impl Plugin {
         println!();
     }
 
-    async fn init<T, F: Future<Output = ()> + Send + 'static>(
+    async fn init(
         precompiled: bool,
         module: Vec<u8>,
         plugin: PluginBuilder,
-        spawn_incrementer: impl Fn(F) -> T,
-    ) -> Result<(Self, T), Error> {
+    ) -> Result<
+        (
+            Self,
+            std::pin::Pin<Box<dyn Future<Output = ()> + Send + 'static>>,
+        ),
+        Error,
+    > {
         // initialize the WebAssembly System Interface context
         let engine = plugin.engine;
         let mut linker = plugin.linker;
@@ -323,7 +341,6 @@ impl Plugin {
                 alloc: None,
             },
         );
-        store.epoch_deadline_async_yield_and_update(1);
 
         let module = if precompiled {
             unsafe { Module::deserialize(&engine, module)? }
@@ -331,6 +348,16 @@ impl Plugin {
             Module::new(&engine, module)?
         };
 
+        // set up automatic yielding after given duration
+        store.epoch_deadline_async_yield_and_update(plugin.delta);
+        let epoch = plugin.epoch;
+        let incrementer = Box::pin(async move {
+            loop {
+                smol::Timer::after(epoch).await;
+                engine.increment_epoch();
+            }
+        });
+
         // load the provided module into the asynchronous runtime
         linker.module_async(&mut store, "", &module).await?;
         let instance = linker.instantiate_async(&mut store, &module).await?;
@@ -345,13 +372,6 @@ impl Plugin {
         });
 
         let plugin = Plugin { store, instance };
-        let incrementer = spawn_incrementer(async move {
-            loop {
-                smol::Timer::after(std::time::Duration::from_millis(100)).await;
-
-                engine.increment_epoch();
-            }
-        });
 
         Ok((plugin, incrementer))
     }
diff --git a/crates/zed/src/languages/language_plugin.rs b/crates/zed/src/languages/language_plugin.rs
index ac649c425d..cf812fd2ad 100644
--- a/crates/zed/src/languages/language_plugin.rs
+++ b/crates/zed/src/languages/language_plugin.rs
@@ -9,7 +9,7 @@ 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()?
+    let (plugin, incrementer) = PluginBuilder::new_with_default_ctx()?
         .host_function_async("command", |command: String| async move {
             let mut args = command.split(' ');
             let command = args.next().unwrap();
@@ -25,6 +25,8 @@ pub async fn new_json(executor: Arc<Background>) -> Result<PluginLspAdapter> {
             include_bytes!("../../../../plugins/bin/json_language.wasm.pre"),
         )
         .await?;
+
+    executor.spawn(incrementer).detach();
     PluginLspAdapter::new(plugin, executor).await
 }