feat(plugins): update and render plugins asynchronously (#2410)

* working-ish minus a few race conditions

* relax atomicity

* various refactoringz

* remove commented out code

* clarify some stuffs

* refactor(plugins): move PluginMap and friends to a different file

* refactor(plugins): move zellij_exports and friends to a different file

* style(fmt): rustfmt

* fix(e2e): adjust tests for flakiness async loading

1 files changed, 167 insertions, 77 deletions

diff --git a/zellij-server/src/plugins/plugin_loader.rs b/zellij-server/src/plugins/plugin_loader.rs
index c5413388..86afd8ca 100644
--- a/zellij-server/src/plugins/plugin_loader.rs
+++ b/zellij-server/src/plugins/plugin_loader.rs
@@ -1,4 +1,5 @@
-use crate::plugins::wasm_bridge::{wasi_read_string, zellij_exports, PluginEnv, PluginMap};
+use crate::plugins::plugin_map::{PluginEnv, PluginMap, RunningPlugin, Subscriptions};
+use crate::plugins::zellij_exports::{wasi_read_string, zellij_exports};
 use highway::{HighwayHash, PortableHash};
 use log::info;
 use semver::Version;
@@ -146,21 +147,8 @@ fn assert_plugin_version(instance: &Instance, plugin_env: &PluginEnv) -> Result<
     Ok(())
 }
 
-fn load_plugin_instance(instance: &mut Instance) -> Result<()> {
-    let err_context = || format!("failed to load plugin from instance {instance:#?}");
-
-    let load_function = instance
-        .exports
-        .get_function("_start")
-        .with_context(err_context)?;
-    // This eventually calls the `.load()` method
-    load_function.call(&[]).with_context(err_context)?;
-    Ok(())
-}
-
 pub struct PluginLoader<'a> {
     plugin_cache: Arc<Mutex<HashMap<PathBuf, Module>>>,
-    plugin_map: Arc<Mutex<PluginMap>>,
     plugin_path: PathBuf,
     loading_indication: &'a mut LoadingIndication,
     senders: ThreadSenders,
@@ -206,15 +194,20 @@ impl<'a> PluginLoader<'a> {
         )?;
         plugin_loader
             .load_module_from_memory()
-            .and_then(|module| plugin_loader.create_plugin_instance_and_environment(module))
-            .and_then(|(instance, plugin_env)| {
-                plugin_loader.load_plugin_instance(&instance, &plugin_env)?;
-                plugin_loader.clone_instance_for_other_clients(
+            .and_then(|module| {
+                plugin_loader.create_plugin_instance_environment_and_subscriptions(module)
+            })
+            .and_then(|(instance, plugin_env, subscriptions)| {
+                plugin_loader.load_plugin_instance(
                     &instance,
                     &plugin_env,
-                    &connected_clients,
+                    &plugin_map,
+                    &subscriptions,
                 )
             })
+            .and_then(|_| {
+                plugin_loader.clone_instance_for_other_clients(&connected_clients, &plugin_map)
+            })
             .with_context(err_context)?;
         display_loading_stage!(end, loading_indication, senders, plugin_id);
         Ok(())
@@ -237,7 +230,6 @@ impl<'a> PluginLoader<'a> {
         let err_context = || format!("failed to start plugin {plugin:#?} for client {client_id}");
         let mut plugin_loader = PluginLoader::new(
             &plugin_cache,
-            &plugin_map,
             loading_indication,
             &senders,
             plugin_id,
@@ -252,19 +244,69 @@ impl<'a> PluginLoader<'a> {
             .load_module_from_memory()
             .or_else(|_e| plugin_loader.load_module_from_hd_cache())
             .or_else(|_e| plugin_loader.compile_module())
-            .and_then(|module| plugin_loader.create_plugin_instance_and_environment(module))
-            .and_then(|(instance, plugin_env)| {
-                plugin_loader.load_plugin_instance(&instance, &plugin_env)?;
-                plugin_loader.clone_instance_for_other_clients(
+            .and_then(|module| {
+                plugin_loader.create_plugin_instance_environment_and_subscriptions(module)
+            })
+            .and_then(|(instance, plugin_env, subscriptions)| {
+                plugin_loader.load_plugin_instance(
                     &instance,
                     &plugin_env,
+                    &plugin_map,
+                    &subscriptions,
+                )
+            })
+            .and_then(|_| {
+                plugin_loader.clone_instance_for_other_clients(
                     &connected_clients.lock().unwrap(),
+                    &plugin_map,
                 )
             })
             .with_context(err_context)?;
         display_loading_stage!(end, loading_indication, senders, plugin_id);
         Ok(())
     }
+    pub fn add_client(
+        client_id: ClientId,
+        plugin_dir: PathBuf,
+        plugin_cache: Arc<Mutex<HashMap<PathBuf, Module>>>,
+        senders: ThreadSenders,
+        store: Store,
+        plugin_map: Arc<Mutex<PluginMap>>,
+        connected_clients: Arc<Mutex<Vec<ClientId>>>,
+        loading_indication: &mut LoadingIndication,
+    ) -> Result<()> {
+        let mut new_plugins = HashSet::new();
+        for (&(plugin_id, _), _) in &*plugin_map.lock().unwrap() {
+            new_plugins.insert((plugin_id, client_id));
+        }
+        for (plugin_id, existing_client_id) in new_plugins {
+            let mut plugin_loader = PluginLoader::new_from_different_client_id(
+                &plugin_cache,
+                &plugin_map,
+                loading_indication,
+                &senders,
+                plugin_id,
+                existing_client_id,
+                &store,
+                &plugin_dir,
+            )?;
+            plugin_loader
+                .load_module_from_memory()
+                .and_then(|module| {
+                    plugin_loader.create_plugin_instance_environment_and_subscriptions(module)
+                })
+                .and_then(|(instance, plugin_env, subscriptions)| {
+                    plugin_loader.load_plugin_instance(
+                        &instance,
+                        &plugin_env,
+                        &plugin_map,
+                        &subscriptions,
+                    )
+                })?
+        }
+        connected_clients.lock().unwrap().push(client_id);
+        Ok(())
+    }
 
     pub fn reload_plugin(
         plugin_id: u32,
@@ -297,22 +339,26 @@ impl<'a> PluginLoader<'a> {
         )?;
         plugin_loader
             .compile_module()
-            .and_then(|module| plugin_loader.create_plugin_instance_and_environment(module))
-            .and_then(|(instance, plugin_env)| {
-                plugin_loader.load_plugin_instance(&instance, &plugin_env)?;
-                plugin_loader.clone_instance_for_other_clients(
+            .and_then(|module| {
+                plugin_loader.create_plugin_instance_environment_and_subscriptions(module)
+            })
+            .and_then(|(instance, plugin_env, subscriptions)| {
+                plugin_loader.load_plugin_instance(
                     &instance,
                     &plugin_env,
-                    &connected_clients,
+                    &plugin_map,
+                    &subscriptions,
                 )
             })
+            .and_then(|_| {
+                plugin_loader.clone_instance_for_other_clients(&connected_clients, &plugin_map)
+            })
             .with_context(err_context)?;
         display_loading_stage!(end, loading_indication, senders, plugin_id);
         Ok(())
     }
     pub fn new(
         plugin_cache: &Arc<Mutex<HashMap<PathBuf, Module>>>,
-        plugin_map: &Arc<Mutex<PluginMap>>,
         loading_indication: &'a mut LoadingIndication,
         senders: &ThreadSenders,
         plugin_id: u32,
@@ -328,7 +374,6 @@ impl<'a> PluginLoader<'a> {
         let plugin_path = plugin.path.clone();
         Ok(PluginLoader {
             plugin_cache: plugin_cache.clone(),
-            plugin_map: plugin_map.clone(),
             plugin_path,
             loading_indication,
             senders: senders.clone(),
@@ -354,19 +399,63 @@ impl<'a> PluginLoader<'a> {
         plugin_dir: &'a PathBuf,
     ) -> Result<Self> {
         let err_context = || "Failed to find existing plugin";
-        let (_old_instance, old_user_env, (rows, cols)) = {
+        let (running_plugin, _subscriptions) = {
             let mut plugin_map = plugin_map.lock().unwrap();
             plugin_map
                 .remove(&(plugin_id, client_id))
                 .with_context(err_context)?
         };
-        let tab_index = old_user_env.tab_index;
-        let size = Size { rows, cols };
-        let plugin_config = old_user_env.plugin.clone();
-        loading_indication.set_name(old_user_env.name());
+        let running_plugin = running_plugin.lock().unwrap();
+        let tab_index = running_plugin.plugin_env.tab_index;
+        let size = Size {
+            rows: running_plugin.rows,
+            cols: running_plugin.columns,
+        };
+        let plugin_config = running_plugin.plugin_env.plugin.clone();
+        loading_indication.set_name(running_plugin.plugin_env.name());
+        PluginLoader::new(
+            plugin_cache,
+            loading_indication,
+            senders,
+            plugin_id,
+            client_id,
+            store,
+            plugin_config,
+            plugin_dir,
+            tab_index,
+            size,
+        )
+    }
+    pub fn new_from_different_client_id(
+        plugin_cache: &Arc<Mutex<HashMap<PathBuf, Module>>>,
+        plugin_map: &Arc<Mutex<PluginMap>>,
+        loading_indication: &'a mut LoadingIndication,
+        senders: &ThreadSenders,
+        plugin_id: u32,
+        client_id: ClientId,
+        store: &Store,
+        plugin_dir: &'a PathBuf,
+    ) -> Result<Self> {
+        let err_context = || "Failed to find existing plugin";
+        let (running_plugin, _subscriptions) = {
+            let plugin_map = plugin_map.lock().unwrap();
+            plugin_map
+                .iter()
+                .find(|((p_id, _c_id), _)| p_id == &plugin_id)
+                .with_context(err_context)?
+                .1
+                .clone()
+        };
+        let running_plugin = running_plugin.lock().unwrap();
+        let tab_index = running_plugin.plugin_env.tab_index;
+        let size = Size {
+            rows: running_plugin.rows,
+            cols: running_plugin.columns,
+        };
+        let plugin_config = running_plugin.plugin_env.plugin.clone();
+        loading_indication.set_name(running_plugin.plugin_env.name());
         PluginLoader::new(
             plugin_cache,
-            plugin_map,
             loading_indication,
             senders,
             plugin_id,
@@ -464,10 +553,10 @@ impl<'a> PluginLoader<'a> {
             .with_context(err_context)?;
         Ok(module)
     }
-    pub fn create_plugin_instance_and_environment(
+    pub fn create_plugin_instance_environment_and_subscriptions(
         &mut self,
         module: Module,
-    ) -> Result<(Instance, PluginEnv)> {
+    ) -> Result<(Instance, PluginEnv, Arc<Mutex<Subscriptions>>)> {
         let err_context = || {
             format!(
                 "Failed to create instance and plugin env for plugin {}",
@@ -499,12 +588,12 @@ impl<'a> PluginLoader<'a> {
             plugin: mut_plugin,
             senders: self.senders.clone(),
             wasi_env,
-            subscriptions: Arc::new(Mutex::new(HashSet::new())),
             plugin_own_data_dir: self.plugin_own_data_dir.clone(),
             tab_index: self.tab_index,
         };
 
-        let zellij = zellij_exports(&self.store, &plugin_env);
+        let subscriptions = Arc::new(Mutex::new(HashSet::new()));
+        let zellij = zellij_exports(&self.store, &plugin_env, &subscriptions);
         let instance =
             Instance::new(&module, &zellij.chain_back(wasi)).with_context(err_context)?;
         assert_plugin_version(&instance, &plugin_env).with_context(err_context)?;
@@ -514,12 +603,14 @@ impl<'a> PluginLoader<'a> {
             .lock()
             .unwrap()
             .insert(cloned_plugin.path, module);
-        Ok((instance, plugin_env))
+        Ok((instance, plugin_env, subscriptions))
     }
     pub fn load_plugin_instance(
         &mut self,
         instance: &Instance,
         plugin_env: &PluginEnv,
+        plugin_map: &Arc<Mutex<PluginMap>>,
+        subscriptions: &Arc<Mutex<Subscriptions>>,
     ) -> Result<()> {
         let err_context = || format!("failed to load plugin from instance {instance:#?}");
         let main_user_instance = instance.clone();
@@ -548,13 +639,16 @@ impl<'a> PluginLoader<'a> {
             self.senders,
             self.plugin_id
         );
-        let mut plugin_map = self.plugin_map.lock().unwrap();
-        plugin_map.insert(
+        plugin_map.lock().unwrap().insert(
             (self.plugin_id, self.client_id),
             (
-                main_user_instance,
-                main_user_env,
-                (self.size.rows, self.size.cols),
+                Arc::new(Mutex::new(RunningPlugin::new(
+                    main_user_instance,
+                    main_user_env,
+                    self.size.rows,
+                    self.size.cols,
+                ))),
+                subscriptions.clone(),
             ),
         );
         display_loading_stage!(
@@ -567,9 +661,8 @@ impl<'a> PluginLoader<'a> {
     }
     pub fn clone_instance_for_other_clients(
         &mut self,
-        instance: &Instance,
-        plugin_env: &PluginEnv,
         connected_clients: &[ClientId],
+        plugin_map: &Arc<Mutex<PluginMap>>,
     ) -> Result<()> {
         if !connected_clients.is_empty() {
             display_loading_stage!(
@@ -578,14 +671,32 @@ impl<'a> PluginLoader<'a> {
                 self.senders,
                 self.plugin_id
             );
-            let mut plugin_map = self.plugin_map.lock().unwrap();
             for client_id in connected_clients {
-                let (instance, new_plugin_env) =
-                    clone_plugin_for_client(&plugin_env, *client_id, &instance, &self.store)?;
-                plugin_map.insert(
-                    (self.plugin_id, *client_id),
-                    (instance, new_plugin_env, (self.size.rows, self.size.cols)),
-                );
+                let mut loading_indication = LoadingIndication::new("".into());
+                let mut plugin_loader_for_client = PluginLoader::new_from_different_client_id(
+                    &self.plugin_cache.clone(),
+                    &plugin_map,
+                    &mut loading_indication,
+                    &self.senders.clone(),
+                    self.plugin_id,
+                    *client_id,
+                    &self.store,
+                    &self.plugin_dir,
+                )?;
+                plugin_loader_for_client
+                    .load_module_from_memory()
+                    .and_then(|module| {
+                        plugin_loader_for_client
+                            .create_plugin_instance_environment_and_subscriptions(module)
+                    })
+                    .and_then(|(instance, plugin_env, subscriptions)| {
+                        plugin_loader_for_client.load_plugin_instance(
+                            &instance,
+                            &plugin_env,
+                            plugin_map,
+                            &subscriptions,
+                        )
+                    })?
             }
             display_loading_stage!(
                 indicate_cloning_plugin_for_other_clients_success,
@@ -633,24 +744,3 @@ fn create_plugin_fs_entries(plugin_own_data_dir: &PathBuf) -> Result<()> {
         .with_context(err_context)?;
     Ok(())
 }
-
-fn clone_plugin_for_client(
-    plugin_env: &PluginEnv,
-    client_id: ClientId,
-    instance: &Instance,
-    store: &Store,
-) -> Result<(Instance, PluginEnv)> {
-    let err_context = || format!("Failed to clone plugin for client {client_id}");
-    let mut new_plugin_env = plugin_env.clone();
-    new_plugin_env.client_id = client_id;
-    let module = instance.module().clone();
-    let wasi = new_plugin_env
-        .wasi_env
-        .import_object(&module)
-        .with_context(err_context)?;
-    let zellij = zellij_exports(store, &new_plugin_env);
-    let mut instance =
-        Instance::new(&module, &zellij.chain_back(wasi)).with_context(err_context)?;
-    load_plugin_instance(&mut instance).with_context(err_context)?;
-    Ok((instance, new_plugin_env))
-}