From 7f9eab6eeae864bcf946e8c4e6e9fab9d827f36e Mon Sep 17 00:00:00 2001
From: Jan Tache <j.andreitabs@gmail.com>
Date: Fri, 8 Dec 2023 00:14:16 -0800
Subject: [PATCH] feat(linux): add configurable delay for registering devices

---
 src/main.rs        | 31 ++++++++++++++++++++++++-------
 src/oskbd/linux.rs |  9 +++++++--
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index d986d2977..b74d6776d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,19 +1,21 @@
 use anyhow::{bail, Result};
+use clap::Parser;
 use log::info;
 use simplelog::*;
+
 use std::path::PathBuf;
+use std::sync::atomic::Ordering;
 
 mod kanata;
 mod oskbd;
 mod tcp_server;
 
-#[cfg(test)]
-mod tests;
-
-use clap::Parser;
 use kanata::Kanata;
 use tcp_server::TcpServer;
 
+#[cfg(test)]
+mod tests;
+
 type CfgPath = PathBuf;
 
 pub struct ValidatedArgs {
@@ -98,10 +100,19 @@ kanata.kbd in the current working directory and
     #[arg(short, long)]
     trace: bool,
 
-    /// Remove the startup delay on kanata. In some cases, removing the delay may cause keyboard
-    /// issues on startup.
-    #[arg(short, long)]
+    /// Remove the startup delay on kanata.
+    /// In some cases, removing the delay may cause keyboard issues on startup.
+    #[arg(short, long, verbatim_doc_comment)]
     nodelay: bool,
+
+    /// Milliseconds to wait before attempting to register a newly connected
+    /// device. The default is 200.
+    ///
+    /// You may wish to increase this if you have a device that is failing
+    /// to register - the device may be taking too long to become ready.
+    #[cfg(target_os = "linux")]
+    #[arg(short, long, verbatim_doc_comment)]
+    wait_device_ms: Option<u64>,
 }
 
 /// Parse CLI arguments and initialize logging.
@@ -149,6 +160,12 @@ fn cli_init() -> Result<ValidatedArgs> {
         bail!("No config files provided\nFor more info, pass the `-h` or `--help` flags.");
     }
 
+    #[cfg(target_os = "linux")]
+    if let Some(wait) = args.wait_device_ms {
+        log::info!("Setting device registration wait time to {wait} ms.");
+        oskbd::WAIT_DEVICE_MS.store(wait, Ordering::SeqCst);
+    }
+
     Ok(ValidatedArgs {
         paths: cfg_paths,
         port: args.port,
diff --git a/src/oskbd/linux.rs b/src/oskbd/linux.rs
index 7ecc245dd..d6e8319ae 100644
--- a/src/oskbd/linux.rs
+++ b/src/oskbd/linux.rs
@@ -15,6 +15,7 @@ use std::fs;
 use std::io;
 use std::os::unix::io::AsRawFd;
 use std::path::PathBuf;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::thread;
 
 use super::*;
@@ -38,6 +39,8 @@ pub struct KbdIn {
 const INOTIFY_TOKEN_VALUE: usize = 0;
 const INOTIFY_TOKEN: Token = Token(INOTIFY_TOKEN_VALUE);
 
+pub static WAIT_DEVICE_MS: AtomicU64 = AtomicU64::new(200);
+
 impl KbdIn {
     pub fn new(
         dev_paths: &[String],
@@ -184,7 +187,7 @@ impl KbdIn {
             let discovered_devices = missing
                 .iter()
                 .filter_map(|dev_path| {
-                    for _ in 0..10 {
+                    for _ in 0..(WAIT_DEVICE_MS.load(Ordering::SeqCst) / 10) {
                         // try a few times with waits in between; device might not be ready
                         if let Ok(device) = Device::open(dev_path) {
                             return Some((device, dev_path.clone()));
@@ -206,7 +209,9 @@ impl KbdIn {
             missing.retain(|path| !paths_registered.contains(path));
         } else {
             log::info!("sleeping for a moment to let devices become ready");
-            std::thread::sleep(std::time::Duration::from_millis(200));
+            std::thread::sleep(std::time::Duration::from_millis(
+                WAIT_DEVICE_MS.load(Ordering::SeqCst),
+            ));
             discover_devices(self.include_names.as_deref(), self.exclude_names.as_deref())
                 .into_iter()
                 .try_for_each(|(dev, path)| {