feat(plugins): reload plugin at runtime (#2372)

* fix(plugins): proper error when wasm file does not exist

* reload working

* race condition handling

* refactor(plugins): start plugin

* refactor(plugins): plugin-loader

* refactor(plugins): load/reload plugin

* refactor(plugins): apply cached events

* fix(plugins): gittery loading

* chore(plugins): rename reload-plugin to start-or-reload-plugin

* chore(styling): small cleanups

* style(fmt): rustfmt

* style(fmt): cleanups

* style(fmt): cleanups

* test(e2e): update snapshots

* test(e2e): update snapshots

* chore(repo): comment plugin optimization because it doubles the CI time

11 files changed, 1030 insertions, 590 deletions

diff --git a/zellij-server/src/lib.rs b/zellij-server/src/lib.rs
index bf0f1a82f..f82c9d682 100644
--- a/zellij-server/src/lib.rs
+++ b/zellij-server/src/lib.rs
@@ -764,7 +764,7 @@ fn init_session(
                 Some(&to_screen),
                 Some(&to_pty),
                 Some(&to_plugin),
-                None,
+                Some(&to_server),
                 Some(&to_pty_writer),
                 Some(&to_background_jobs),
                 None,
diff --git a/zellij-server/src/panes/grid.rs b/zellij-server/src/panes/grid.rs
index b76d21ecb..6f44f6eb4 100644
--- a/zellij-server/src/panes/grid.rs
+++ b/zellij-server/src/panes/grid.rs
@@ -1254,7 +1254,7 @@ impl Grid {
             let new_row = Row::new(self.width).canonical();
             self.viewport.push(new_row);
         }
-        if self.cursor.y == self.height - 1 {
+        if self.cursor.y == self.height.saturating_sub(1) {
             if self.scroll_region.is_none() {
                 if self.alternate_screen_state.is_none() {
                     self.transfer_rows_to_lines_above(1);
@@ -1406,7 +1406,7 @@ impl Grid {
     }
     fn line_wrap(&mut self) {
         self.cursor.x = 0;
-        if self.cursor.y == self.height - 1 {
+        if self.cursor.y == self.height.saturating_sub(1) {
             if self.alternate_screen_state.is_none() {
                 self.transfer_rows_to_lines_above(1);
             } else {
diff --git a/zellij-server/src/panes/plugin_pane.rs b/zellij-server/src/panes/plugin_pane.rs
index 8a45c3c10..7915069c5 100644
--- a/zellij-server/src/panes/plugin_pane.rs
+++ b/zellij-server/src/panes/plugin_pane.rs
@@ -545,6 +545,12 @@ impl Pane for PluginPane {
             self.loading_indication.to_string().as_bytes().to_vec(),
         );
     }
+    fn start_loading_indication(&mut self, loading_indication: LoadingIndication) {
+        self.loading_indication.merge(loading_indication);
+        self.handle_plugin_bytes_for_all_clients(
+            self.loading_indication.to_string().as_bytes().to_vec(),
+        );
+    }
     fn progress_animation_offset(&mut self) {
         if self.loading_indication.ended {
             return;
diff --git a/zellij-server/src/plugins/mod.rs b/zellij-server/src/plugins/mod.rs
index 2aca679a3..eeb19c201 100644
--- a/zellij-server/src/plugins/mod.rs
+++ b/zellij-server/src/plugins/mod.rs
@@ -1,11 +1,11 @@
-mod start_plugin;
+mod plugin_loader;
 mod wasm_bridge;
 use log::info;
 use std::{collections::HashMap, fs, path::PathBuf};
 use wasmer::Store;
 
 use crate::screen::ScreenInstruction;
-use crate::{pty::PtyInstruction, thread_bus::Bus, ClientId};
+use crate::{pty::PtyInstruction, thread_bus::Bus, ClientId, ServerInstruction};
 
 use wasm_bridge::WasmBridge;
 
@@ -32,7 +32,15 @@ pub enum PluginInstruction {
     ),
     Update(Vec<(Option<u32>, Option<ClientId>, Event)>), // Focused plugin / broadcast, client_id, event data
     Unload(u32),                                         // plugin_id
-    Resize(u32, usize, usize),                           // plugin_id, columns, rows
+    Reload(
+        Option<bool>,   // should float
+        Option<String>, // pane title
+        RunPlugin,
+        usize, // tab index
+        ClientId,
+        Size,
+    ),
+    Resize(u32, usize, usize), // plugin_id, columns, rows
     AddClient(ClientId),
     RemoveClient(ClientId),
     NewTab(
@@ -43,7 +51,7 @@ pub enum PluginInstruction {
         usize, // tab_index
         ClientId,
     ),
-    ApplyCachedEvents(u32), // u32 is the plugin id
+    ApplyCachedEvents(Vec<u32>), // a list of plugin id
     Exit,
 }
 
@@ -53,6 +61,7 @@ impl From<&PluginInstruction> for PluginContext {
             PluginInstruction::Load(..) => PluginContext::Load,
             PluginInstruction::Update(..) => PluginContext::Update,
             PluginInstruction::Unload(..) => PluginContext::Unload,
+            PluginInstruction::Reload(..) => PluginContext::Reload,
             PluginInstruction::Resize(..) => PluginContext::Resize,
             PluginInstruction::Exit => PluginContext::Exit,
             PluginInstruction::AddClient(_) => PluginContext::AddClient,
@@ -103,6 +112,42 @@ pub(crate) fn plugin_thread_main(
             PluginInstruction::Unload(pid) => {
                 wasm_bridge.unload_plugin(pid)?;
             },
+            PluginInstruction::Reload(
+                should_float,
+                pane_title,
+                run,
+                tab_index,
+                client_id,
+                size,
+            ) => match wasm_bridge.reload_plugin(&run) {
+                Ok(_) => {
+                    let _ = bus
+                        .senders
+                        .send_to_server(ServerInstruction::UnblockInputThread);
+                },
+                Err(err) => match err.downcast_ref::<ZellijError>() {
+                    Some(ZellijError::PluginDoesNotExist) => {
+                        log::warn!("Plugin {} not found, starting it instead", run.location);
+                        match wasm_bridge.load_plugin(&run, tab_index, size, client_id) {
+                            Ok(plugin_id) => {
+                                drop(bus.senders.send_to_screen(ScreenInstruction::AddPlugin(
+                                    should_float,
+                                    run,
+                                    pane_title,
+                                    tab_index,
+                                    plugin_id,
+                                )));
+                            },
+                            Err(e) => {
+                                log::error!("Failed to load plugin: {e}");
+                            },
+                        };
+                    },
+                    _ => {
+                        return Err(err);
+                    },
+                },
+            },
             PluginInstruction::Resize(pid, new_columns, new_rows) => {
                 wasm_bridge.resize_plugin(pid, new_columns, new_rows)?;
             },
diff --git a/zellij-server/src/plugins/plugin_loader.rs b/zellij-server/src/plugins/plugin_loader.rs
new file mode 100644
index 000000000..c54133889
--- /dev/null
+++ b/zellij-server/src/plugins/plugin_loader.rs
@@ -0,0 +1,656 @@
+use crate::plugins::wasm_bridge::{wasi_read_string, zellij_exports, PluginEnv, PluginMap};
+use highway::{HighwayHash, PortableHash};
+use log::info;
+use semver::Version;
+use std::{
+    collections::{HashMap, HashSet},
+    fmt, fs,
+    path::PathBuf,
+    sync::{Arc, Mutex},
+};
+use url::Url;
+use wasmer::{ChainableNamedResolver, Instance, Module, Store};
+use wasmer_wasi::{Pipe, WasiState};
+
+use crate::{
+    logging_pipe::LoggingPipe, screen::ScreenInstruction, thread_bus::ThreadSenders,
+    ui::loading_indication::LoadingIndication, ClientId,
+};
+
+use zellij_utils::{
+    consts::{VERSION, ZELLIJ_CACHE_DIR, ZELLIJ_TMP_DIR},
+    errors::prelude::*,
+    input::plugins::PluginConfig,
+    pane_size::Size,
+};
+
+macro_rules! display_loading_stage {
+    ($loading_stage:ident, $loading_indication:expr, $senders:expr, $plugin_id:expr) => {{
+        $loading_indication.$loading_stage();
+        drop(
+            $senders.send_to_screen(ScreenInstruction::UpdatePluginLoadingStage(
+                $plugin_id,
+                $loading_indication.clone(),
+            )),
+        );
+    }};
+}
+
+/// Custom error for plugin version mismatch.
+///
+/// This is thrown when, during starting a plugin, it is detected that the plugin version doesn't
+/// match the zellij version. This is treated as a fatal error and leads to instantaneous
+/// termination.
+#[derive(Debug)]
+pub struct VersionMismatchError {
+    zellij_version: String,
+    plugin_version: String,
+    plugin_path: PathBuf,
+    // true for builtin plugins
+    builtin: bool,
+}
+
+impl std::error::Error for VersionMismatchError {}
+
+impl VersionMismatchError {
+    pub fn new(
+        zellij_version: &str,
+        plugin_version: &str,
+        plugin_path: &PathBuf,
+        builtin: bool,
+    ) -> Self {
+        VersionMismatchError {
+            zellij_version: zellij_version.to_owned(),
+            plugin_version: plugin_version.to_owned(),
+            plugin_path: plugin_path.to_owned(),
+            builtin,
+        }
+    }
+}
+
+impl fmt::Display for VersionMismatchError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let first_line = if self.builtin {
+            "It seems your version of zellij was built with outdated core plugins."
+        } else {
+            "If you're seeing this error a plugin version doesn't match the current
+zellij version."
+        };
+
+        write!(
+            f,
+            "{}
+Detected versions:
+
+- Plugin version: {}
+- Zellij version: {}
+- Offending plugin: {}
+
+If you're a user:
+    Please contact the distributor of your zellij version and report this error
+    to them.
+
+If you're a developer:
+    Please run zellij with updated plugins. The easiest way to achieve this
+    is to build zellij with `cargo xtask install`. Also refer to the docs:
+    https://github.com/zellij-org/zellij/blob/main/CONTRIBUTING.md#building
+",
+            first_line,
+            self.plugin_version.trim_end(),
+            self.zellij_version.trim_end(),
+            self.plugin_path.display()
+        )
+    }
+}
+
+// Returns `Ok` if the plugin version matches the zellij version.
+// Returns an `Err` otherwise.
+fn assert_plugin_version(instance: &Instance, plugin_env: &PluginEnv) -> Result<()> {
+    let err_context = || {
+        format!(
+            "failed to determine plugin version for plugin {}",
+            plugin_env.plugin.path.display()
+        )
+    };
+
+    let plugin_version_func = match instance.exports.get_function("plugin_version") {
+        Ok(val) => val,
+        Err(_) => {
+            return Err(anyError::new(VersionMismatchError::new(
+                VERSION,
+                "Unavailable",
+                &plugin_env.plugin.path,
+                plugin_env.plugin.is_builtin(),
+            )))
+        },
+    };
+
+    let plugin_version = plugin_version_func
+        .call(&[])
+        .map_err(anyError::new)
+        .and_then(|_| wasi_read_string(&plugin_env.wasi_env))
+        .and_then(|string| Version::parse(&string).context("failed to parse plugin version"))
+        .with_context(err_context)?;
+    let zellij_version = Version::parse(VERSION)
+        .context("failed to parse zellij version")
+        .with_context(err_context)?;
+    if plugin_version != zellij_version {
+        return Err(anyError::new(VersionMismatchError::new(
+            VERSION,
+            &plugin_version.to_string(),
+            &plugin_env.plugin.path,
+            plugin_env.plugin.is_builtin(),
+        )));
+    }
+
+    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,
+    plugin_id: u32,
+    client_id: ClientId,
+    store: Store,
+    plugin: PluginConfig,
+    plugin_dir: &'a PathBuf,
+    tab_index: usize,
+    plugin_own_data_dir: PathBuf,
+    size: Size,
+    wasm_blob_on_hd: Option<(Vec<u8>, PathBuf)>,
+}
+
+impl<'a> PluginLoader<'a> {
+    pub fn reload_plugin_from_memory(
+        plugin_id: u32,
+        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 err_context = || format!("failed to reload plugin {plugin_id} from memory");
+        let mut connected_clients: Vec<ClientId> =
+            connected_clients.lock().unwrap().iter().copied().collect();
+        if connected_clients.is_empty() {
+            return Err(anyhow!("No connected clients, cannot reload plugin"));
+        }
+        let first_client_id = connected_clients.remove(0);
+
+        let mut plugin_loader = PluginLoader::new_from_existing_plugin_attributes(
+            &plugin_cache,
+            &plugin_map,
+            loading_indication,
+            &senders,
+            plugin_id,
+            first_client_id,
+            &store,
+            &plugin_dir,
+        )?;
+        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(
+                    &instance,
+                    &plugin_env,
+                    &connected_clients,
+                )
+            })
+            .with_context(err_context)?;
+        display_loading_stage!(end, loading_indication, senders, plugin_id);
+        Ok(())
+    }
+
+    pub fn start_plugin(
+        plugin_id: u32,
+        client_id: ClientId,
+        plugin: &PluginConfig,
+        tab_index: usize,
+        plugin_dir: PathBuf,
+        plugin_cache: Arc<Mutex<HashMap<PathBuf, Module>>>,
+        senders: ThreadSenders,
+        store: Store,
+        plugin_map: Arc<Mutex<PluginMap>>,
+        size: Size,
+        connected_clients: Arc<Mutex<Vec<ClientId>>>,
+        loading_indication: &mut LoadingIndication,
+    ) -> Result<()> {
+        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,
+            client_id,
+            &store,
+            plugin.clone(),
+            &plugin_dir,
+            tab_index,
+            size,
+        )?;
+        plugin_loader
+            .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(
+                    &instance,
+                    &plugin_env,
+                    &connected_clients.lock().unwrap(),
+                )
+            })
+            .with_context(err_context)?;
+        display_loading_stage!(end, loading_indication, senders, plugin_id);
+        Ok(())
+    }
+
+    pub fn reload_plugin(
+        plugin_id: u32,
+        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 err_context = || format!("failed to reload plugin id {plugin_id}");
+
+        let mut connected_clients: Vec<ClientId> =
+            connected_clients.lock().unwrap().iter().copied().collect();
+        if connected_clients.is_empty() {
+            return Err(anyhow!("No connected clients, cannot reload plugin"));
+        }
+        let first_client_id = connected_clients.remove(0);
+
+        let mut plugin_loader = PluginLoader::new_from_existing_plugin_attributes(
+            &plugin_cache,
+            &plugin_map,
+            loading_indication,
+            &senders,
+            plugin_id,
+            first_client_id,
+            &store,
+            &plugin_dir,
+        )?;
+        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(
+                    &instance,
+                    &plugin_env,
+                    &connected_clients,
+                )
+            })
+            .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,
+        client_id: ClientId,
+        store: &Store,
+        plugin: PluginConfig,
+        plugin_dir: &'a PathBuf,
+        tab_index: usize,
+        size: Size,
+    ) -> Result<Self> {
+        let plugin_own_data_dir = ZELLIJ_CACHE_DIR.join(Url::from(&plugin.location).to_string());
+        create_plugin_fs_entries(&plugin_own_data_dir)?;
+        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(),
+            plugin_id,
+            client_id,
+            store: store.clone(),
+            plugin,
+            plugin_dir,
+            tab_index,
+            plugin_own_data_dir,
+            size,
+            wasm_blob_on_hd: None,
+        })
+    }
+    pub fn new_from_existing_plugin_attributes(
+        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 (_old_instance, old_user_env, (rows, cols)) = {
+            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());
+        PluginLoader::new(
+            plugin_cache,
+            plugin_map,
+            loading_indication,
+            senders,
+            plugin_id,
+            client_id,
+            store,
+            plugin_config,
+            plugin_dir,
+            tab_index,
+            size,
+        )
+    }
+    pub fn load_module_from_memory(&mut self) -> Result<Module> {
+        display_loading_stage!(
+            indicate_loading_plugin_from_memory,
+            self.loading_indication,
+            self.senders,
+            self.plugin_id
+        );
+        let module = self
+            .plugin_cache
+            .lock()
+            .unwrap()
+            .remove(&self.plugin_path)
+            .ok_or(anyhow!("Plugin is not stored in memory"))?;
+        display_loading_stage!(
+            indicate_loading_plugin_from_memory_success,
+            self.loading_indication,
+            self.senders,
+            self.plugin_id
+        );
+        Ok(module)
+    }
+    pub fn load_module_from_hd_cache(&mut self) -> Result<Module> {
+        display_loading_stage!(
+            indicate_loading_plugin_from_memory_notfound,
+            self.loading_indication,
+            self.senders,
+            self.plugin_id
+        );
+        display_loading_stage!(
+            indicate_loading_plugin_from_hd_cache,
+            self.loading_indication,
+            self.senders,
+            self.plugin_id
+        );
+        let (_wasm_bytes, cached_path) = self.plugin_bytes_and_cache_path()?;
+        let timer = std::time::Instant::now();
+        let module = unsafe { Module::deserialize_from_file(&self.store, &cached_path)? };
+        log::info!(
+            "Loaded plugin '{}' from cache folder at '{}' in {:?}",
+            self.plugin_path.display(),
+            ZELLIJ_CACHE_DIR.display(),
+            timer.elapsed(),
+        );
+        display_loading_stage!(
+            indicate_loading_plugin_from_hd_cache_success,
+            self.loading_indication,
+            self.senders,
+            self.plugin_id
+        );
+        Ok(module)
+    }
+    pub fn compile_module(&mut self) -> Result<Module> {
+        display_loading_stage!(
+            indicate_loading_plugin_from_hd_cache_notfound,
+            self.loading_indication,
+            self.senders,
+            self.plugin_id
+        );
+        display_loading_stage!(
+            indicate_compiling_plugin,
+            self.loading_indication,
+            self.senders,
+            self.plugin_id
+        );
+        let (wasm_bytes, cached_path) = self.plugin_bytes_and_cache_path()?;
+        let timer = std::time::Instant::now();
+        let err_context = || "failed to recover cache dir";
+        let module = fs::create_dir_all(ZELLIJ_CACHE_DIR.to_owned())
+            .map_err(anyError::new)
+            .and_then(|_| {
+                // compile module
+                Module::new(&self.store, &wasm_bytes).map_err(anyError::new)
+            })
+            .and_then(|m| {
+                // serialize module to HD cache for faster loading in the future
+                m.serialize_to_file(&cached_path).map_err(anyError::new)?;
+                log::info!(
+                    "Compiled plugin '{}' in {:?}",
+                    self.plugin_path.display(),
+                    timer.elapsed()
+                );
+                Ok(m)
+            })
+            .with_context(err_context)?;
+        Ok(module)
+    }
+    pub fn create_plugin_instance_and_environment(
+        &mut self,
+        module: Module,
+    ) -> Result<(Instance, PluginEnv)> {
+        let err_context = || {
+            format!(
+                "Failed to create instance and plugin env for plugin {}",
+                self.plugin_id
+            )
+        };
+        let mut wasi_env = WasiState::new("Zellij")
+            .env("CLICOLOR_FORCE", "1")
+            .map_dir("/host", ".")
+            .and_then(|wasi| wasi.map_dir("/data", &self.plugin_own_data_dir))
+            .and_then(|wasi| wasi.map_dir("/tmp", ZELLIJ_TMP_DIR.as_path()))
+            .and_then(|wasi| {
+                wasi.stdin(Box::new(Pipe::new()))
+