Simplify the code, let the runtime shutdown the UICommand tasks

src/bridge/mod.rs

@@ -12,15 +12,10 @@ use itertools::Itertools;
 use log::info;
 use nvim_rs::{error::CallError, Neovim, UiAttachOptions, Value};
 use rmpv::Utf8String;
-use std::{
-    io::Error,
-    ops::Add,
-    time::{Duration, Instant},
-};
+use std::{io::Error, ops::Add, time::Duration};
 use tokio::{
     runtime::{Builder, Runtime},
     select,
-    task::JoinSet,
     time::timeout,
 };
 use winit::event_loop::EventLoopProxy;
@@ -37,16 +32,13 @@ pub use api_info::*;
 pub use command::create_nvim_command;
 pub use events::*;
 pub use session::NeovimWriter;
-pub use ui_commands::{
-    send_ui, shutdown_ui, start_ui_command_handler, ParallelCommand, SerialCommand,
-};
+pub use ui_commands::{send_ui, start_ui_command_handler, ParallelCommand, SerialCommand};
 
 const INTRO_MESSAGE_LUA: &str = include_str!("../../lua/intro.lua");
 const NEOVIM_REQUIRED_VERSION: &str = "0.9.2";
 
 pub struct NeovimRuntime {
     runtime: Runtime,
-    join_set: JoinSet<()>,
 }
 
 fn neovim_instance() -> Result<NeovimInstance> {
@@ -89,11 +81,7 @@ pub async fn show_error_message(
     nvim.echo(prepared_lines, true, vec![]).await
 }
 
-async fn launch(
-    handler: NeovimHandler,
-    grid_size: Option<GridSize<u32>>,
-    join_set: &mut JoinSet<()>,
-) -> Result<NeovimSession> {
+async fn launch(handler: NeovimHandler, grid_size: Option<GridSize<u32>>) -> Result<NeovimSession> {
     let neovim_instance = neovim_instance()?;
 
     let session = NeovimSession::new(neovim_instance, handler)
@@ -126,7 +114,7 @@ async fn launch(
     setup_neovide_specific_state(&session.neovim, should_handle_clipboard, &api_information)
         .await?;
 
-    start_ui_command_handler(session.neovim.clone(), &api_information, join_set);
+    start_ui_command_handler(session.neovim.clone(), &api_information);
     SETTINGS.read_initial_values(&session.neovim).await?;
 
     let mut options = UiAttachOptions::new();
@@ -179,18 +167,11 @@ async fn run(session: NeovimSession, proxy: EventLoopProxy<UserEvent>) {
     proxy.send_event(UserEvent::NeovimExited).ok();
 }
 
-async fn wait(join_set: &mut JoinSet<()>) {
-    while join_set.join_next().await.is_some() {}
-}
-
 impl NeovimRuntime {
     pub fn new() -> Result<Self, Error> {
         let runtime = Builder::new_multi_thread().enable_all().build()?;
 
-        Ok(Self {
-            runtime,
-            join_set: JoinSet::new(),
-        })
+        Ok(Self { runtime })
     }
 
     pub fn launch(
@@ -199,22 +180,8 @@ impl NeovimRuntime {
         grid_size: Option<GridSize<u32>>,
     ) -> Result<()> {
         let handler = start_editor(event_loop_proxy.clone());
-        let session = self
-            .runtime
-            .block_on(launch(handler, grid_size, &mut self.join_set))?;
-        self.join_set
-            .spawn_on(run(session, event_loop_proxy), self.runtime.handle());
+        let session = self.runtime.block_on(launch(handler, grid_size))?;
+        self.runtime.spawn(run(session, event_loop_proxy));
         Ok(())
     }
 }
-
-impl Drop for NeovimRuntime {
-    fn drop(&mut self) {
-        log::info!("Starting neovim runtime shutdown");
-        let start = Instant::now();
-        shutdown_ui();
-        self.runtime.block_on(wait(&mut self.join_set));
-        let elapsed = start.elapsed().as_millis();
-        log::info!("Neovim runtime shutdown took {elapsed} ms");
-    }
-}

src/bridge/ui_commands.rs

@@ -1,14 +1,11 @@
-use std::sync::Mutex;
+use std::sync::OnceLock;
 
 use log::trace;
 
 use anyhow::{Context, Result};
 use nvim_rs::{call_args, error::CallError, rpc::model::IntoVal, Neovim, Value};
 use strum::AsRefStr;
-use tokio::{
-    sync::mpsc::{unbounded_channel, UnboundedReceiver},
-    task::JoinSet,
-};
+use tokio::sync::mpsc::unbounded_channel;
 
 use super::show_error_message;
 use crate::{
@@ -273,34 +270,16 @@ impl AsRef<str> for UiCommand {
     }
 }
 
-struct UIChannels {
-    sender: Mutex<Option<LoggingSender<UiCommand>>>,
-    receiver: Mutex<Option<UnboundedReceiver<UiCommand>>>,
-}
-
-impl UIChannels {
-    fn new() -> Self {
-        let (sender, receiver) = unbounded_channel();
-        Self {
-            sender: Mutex::new(Some(LoggingSender::attach(sender, "UICommand"))),
-            receiver: Mutex::new(Some(receiver)),
-        }
-    }
-}
-
-lazy_static! {
-    static ref UI_CHANNELS: UIChannels = UIChannels::new();
-}
+static UI_COMMAND_CHANNEL: OnceLock<LoggingSender<UiCommand>> = OnceLock::new();
 
-pub fn start_ui_command_handler(
-    nvim: Neovim<NeovimWriter>,
-    api_information: &ApiInformation,
-    join_set: &mut JoinSet<()>,
-) {
+pub fn start_ui_command_handler(nvim: Neovim<NeovimWriter>, api_information: &ApiInformation) {
     let (serial_tx, mut serial_rx) = unbounded_channel::<SerialCommand>();
     let ui_command_nvim = nvim.clone();
-    join_set.spawn(async move {
-        let mut ui_command_receiver = UI_CHANNELS.receiver.lock().unwrap().take().unwrap();
+    let (sender, mut ui_command_receiver) = unbounded_channel();
+    UI_COMMAND_CHANNEL
+        .set(LoggingSender::attach(sender, "UIComand"))
+        .expect("The UI command channel is already created");
+    tokio::spawn(async move {
         loop {
             match ui_command_receiver.recv().await {
                 Some(UiCommand::Serial(serial_command)) => {
@@ -323,7 +302,7 @@ pub fn start_ui_command_handler(
 
     let has_x_buttons = api_information.version.has_version(0, 10, 0);
 
-    join_set.spawn(async move {
+    tokio::spawn(async move {
         tracy_fiber_enter!("Serial command");
         loop {
             tracy_fiber_leave();
@@ -348,16 +327,8 @@ where
     T: Into<UiCommand>,
 {
     let command: UiCommand = command.into();
-    let _ = UI_CHANNELS
-        .sender
-        .lock()
-        .unwrap()
-        .as_ref()
-        .unwrap()
+    let _ = UI_COMMAND_CHANNEL
+        .get()
+        .expect("The UI command channel has not been initialized")
         .send(command);
 }
-
-pub fn shutdown_ui() {
-    *UI_CHANNELS.sender.lock().unwrap() = None;
-    log::info!("The UI subsystem has been shut down")
-}

src/channel_utils.rs

@@ -5,7 +5,7 @@ use tokio::sync::mpsc::{error::SendError as TokioSendError, UnboundedSender};
 
 use crate::profiling::tracy_dynamic_zone;
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct LoggingSender<T>
 where
     T: Debug + AsRef<str>,