feat: reload config on theme change (#5135)

* feat: reload config on theme change

* fix: small TOCTOU window and some logging

* add pr

---------

Co-authored-by: Aram Drevekenin <aram@poor.dev>

1 files changed, 127 insertions, 29 deletions

diff --git a/zellij-utils/src/input/config.rs b/zellij-utils/src/input/config.rs
index 6599024c..4d8d9e1d 100644
--- a/zellij-utils/src/input/config.rs
+++ b/zellij-utils/src/input/config.rs
@@ -8,7 +8,7 @@ use serde::{Deserialize, Serialize};
 use std::collections::HashSet;
 use std::fs::File;
 use std::io::{self, Read};
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use thiserror::Error;
 
 use std::convert::TryFrom;
@@ -439,13 +439,16 @@ impl Config {
 }
 
 #[cfg(not(target_family = "wasm"))]
-pub async fn watch_config_file_changes<F, Fut>(config_file_path: PathBuf, on_config_change: F)
-where
+pub async fn watch_config_file_changes<F, Fut>(
+    config_file_path: PathBuf,
+    config_dir: Option<&Path>,
+    on_config_change: F,
+) where
     F: Fn(Config) -> Fut + Send + 'static,
     Fut: std::future::Future<Output = ()> + Send,
 {
     // in a gist, what we do here is fire the `on_config_change` function whenever there is a
-    // change in the config file, we do this by:
+    // change in the config file or the configured theme directory, we do this by:
     // 1. Trying to watch the provided config file for changes
     // 2. If the file is deleted or does not exist, we periodically poll for it (manually, not
     //    through filesystem events)
@@ -458,8 +461,75 @@ where
     use notify::{self, Config as WatcherConfig, Event, PollWatcher, RecursiveMode, Watcher};
     use std::time::Duration;
     use tokio::sync::mpsc;
+
+    fn cli_args_for_config(config_file_path: &Path, config_dir: Option<&Path>) -> CliArgs {
+        let mut cli_args_for_config = CliArgs::default();
+        cli_args_for_config.config = Some(config_file_path.to_path_buf());
+        cli_args_for_config.config_dir = config_dir.map(Path::to_path_buf);
+        cli_args_for_config
+    }
+
+    fn load_config_and_theme_dir(
+        config_file_path: &Path,
+        config_dir: Option<&Path>,
+    ) -> Option<(Config, Option<PathBuf>)> {
+        let cli_args_for_config = cli_args_for_config(config_file_path, config_dir);
+        Setup::from_cli_args(&cli_args_for_config)
+            .map(|(config, _, config_options, _, _)| {
+                let theme_dir = config_options.theme_dir.or_else(|| {
+                    let config_dir = config_dir
+                        .map(Path::to_path_buf)
+                        .or_else(home::find_default_config_dir);
+                    home::get_theme_dir(config_dir).filter(|dir| dir.exists())
+                });
+                (config, theme_dir)
+            })
+            .ok()
+    }
+
+    fn event_is_for_config_file(event: &Event, config_file_path: &Path) -> bool {
+        event.paths.iter().any(|path| path == config_file_path)
+    }
+
+    fn event_is_in_theme_dir(event: &Event, theme_dir: Option<&Path>) -> bool {
+        theme_dir.map_or(false, |theme_dir| {
+            event.paths.iter().any(|path| path.starts_with(theme_dir))
+        })
+    }
+
+    async fn reload_config_after_change<F, Fut>(
+        config_file_path: &Path,
+        config_dir: Option<&Path>,
+        watched_theme_dir: Option<&Path>,
+        on_config_change: &F,
+    ) -> Option<bool>
+    where
+        F: Fn(Config) -> Fut + Send + 'static,
+        Fut: std::future::Future<Output = ()> + Send,
+    {
+        tokio::time::sleep(Duration::from_millis(100)).await;
+
+        if !config_file_path.exists() {
+            return None;
+        }
+
+        let (new_config, new_theme_dir) =
+            match load_config_and_theme_dir(config_file_path, config_dir) {
+                Some(loaded) => loaded,
+                None => {
+                    log::error!("Failed to reload config from {:?}", config_file_path);
+                    return None;
+                },
+            };
+        on_config_change(new_config).await;
+        Some(new_theme_dir.as_deref() != watched_theme_dir)
+    }
+
     loop {
         if config_file_path.exists() {
+            let watched_theme_dir =
+                load_config_and_theme_dir(config_file_path.as_path(), config_dir)
+                    .and_then(|(_, theme_dir)| theme_dir);
             let (tx, mut rx) = mpsc::unbounded_channel();
 
             let mut watcher = match PollWatcher::new(
@@ -469,40 +539,68 @@ where
                 WatcherConfig::default().with_poll_interval(Duration::from_secs(1)),
             ) {
                 Ok(watcher) => watcher,
-                Err(_) => break,
+                Err(e) => {
+                    log::error!("Failed to create config watcher: {}", e);
+                    break;
+                },
             };
 
-            if watcher
-                .watch(&config_file_path, RecursiveMode::NonRecursive)
-                .is_err()
-            {
+            if let Err(e) = watcher.watch(&config_file_path, RecursiveMode::NonRecursive) {
+                log::error!("Failed to watch config file {:?}: {}", config_file_path, e);
                 break;
             }
 
+            if let Some(watched_theme_dir) = &watched_theme_dir {
+                if let Err(e) = watcher.watch(watched_theme_dir, RecursiveMode::NonRecursive) {
+                    log::error!(
+                        "Failed to watch theme dir {:?}, continuing without it: {}",
+                        watched_theme_dir,
+                        e,
+                    );
+                }
+            }
+
             while let Some(event_result) = rx.recv().await {
-                match event_result {
-                    Ok(event) => {
-                        if event.paths.contains(&config_file_path) {
-                            if event.kind.is_remove() {
-                                break;
-                            } else if event.kind.is_create() || event.kind.is_modify() {
-                                tokio::time::sleep(Duration::from_millis(100)).await;
+                let event = match event_result {
+                    Ok(event) => event,
+                    Err(e) => {
+                        log::error!("Config watcher event error: {}", e);
+                        break;
+                    },
+                };
 
-                                if !config_file_path.exists() {
-                                    continue;
-                                }
+                if event_is_for_config_file(&event, config_file_path.as_path()) {
+                    if event.kind.is_remove() {
+                        break;
+                    }
 
-                                let mut cli_args_for_config = CliArgs::default();
-                                cli_args_for_config.config = Some(PathBuf::from(&config_file_path));
-                                if let Ok(new_config) = Setup::from_cli_args(&cli_args_for_config)
-                                    .map_err(|e| e.to_string())
-                                {
-                                    on_config_change(new_config.0).await;
-                                }
-                            }
+                    if event.kind.is_create() || event.kind.is_modify() {
+                        if reload_config_after_change(
+                            config_file_path.as_path(),
+                            config_dir,
+                            watched_theme_dir.as_deref(),
+                            &on_config_change,
+                        )
+                        .await
+                        .unwrap_or(false)
+                        {
+                            break;
                         }
-                    },
-                    Err(_) => break,
+                    }
+                } else if event_is_in_theme_dir(&event, watched_theme_dir.as_deref())
+                    && (event.kind.is_remove() || event.kind.is_create() || event.kind.is_modify())
+                {
+                    let should_restart_watcher = reload_config_after_change(
+                        config_file_path.as_path(),
+                        config_dir,
+                        watched_theme_dir.as_deref(),
+                        &on_config_change,
+                    )
+                    .await
+                    .unwrap_or(true);
+                    if should_restart_watcher {
+                        break;
+                    }
                 }
             }
         }