From c406b335c576acedf05774fa48658549e6d1eb6d Mon Sep 17 00:00:00 2001
From: Demmie <2e3s19@gmail.com>
Date: Tue, 23 Jan 2024 23:37:11 -0500
Subject: [PATCH] Make a few retries if Gnome Wayland is certain

---
 watchers/src/watchers.rs               |  2 ++
 watchers/src/watchers/gnome_idle.rs    | 14 +++++++---
 watchers/src/watchers/gnome_wayland.rs | 38 ++++++++++++++++++++++++++
 watchers/src/watchers/gnome_window.rs  | 17 ++++++++----
 4 files changed, 62 insertions(+), 9 deletions(-)
 create mode 100644 watchers/src/watchers/gnome_wayland.rs

diff --git a/watchers/src/watchers.rs b/watchers/src/watchers.rs
index 8ad9405..42c7a0a 100644
--- a/watchers/src/watchers.rs
+++ b/watchers/src/watchers.rs
@@ -1,6 +1,8 @@
 #[cfg(feature = "gnome")]
 mod gnome_idle;
 #[cfg(feature = "gnome")]
+mod gnome_wayland;
+#[cfg(feature = "gnome")]
 mod gnome_window;
 mod idle;
 #[cfg(feature = "kwin_window")]
diff --git a/watchers/src/watchers/gnome_idle.rs b/watchers/src/watchers/gnome_idle.rs
index 3c3adfc..cf6df35 100644
--- a/watchers/src/watchers/gnome_idle.rs
+++ b/watchers/src/watchers/gnome_idle.rs
@@ -1,4 +1,4 @@
-use super::{idle, Watcher};
+use super::{gnome_wayland::load_watcher, gnome_wayland::GnomeWatcher, idle, Watcher};
 use crate::report_client::ReportClient;
 use anyhow::Context;
 use async_trait::async_trait;
@@ -28,9 +28,8 @@ impl idle::SinceLastInput for IdleWatcher {
     }
 }
 
-#[async_trait]
-impl Watcher for IdleWatcher {
-    async fn new(_: &Arc<ReportClient>) -> anyhow::Result<Self> {
+impl GnomeWatcher for IdleWatcher {
+    async fn load() -> anyhow::Result<Self> {
         let mut watcher = Self {
             dbus_connection: Connection::session().await?,
             is_idle: false,
@@ -39,6 +38,13 @@ impl Watcher for IdleWatcher {
 
         Ok(watcher)
     }
+}
+
+#[async_trait]
+impl Watcher for IdleWatcher {
+    async fn new(_: &Arc<ReportClient>) -> anyhow::Result<Self> {
+        load_watcher().await
+    }
 
     async fn run_iteration(&mut self, client: &Arc<ReportClient>) -> anyhow::Result<()> {
         self.is_idle = idle::ping_since_last_input(self, self.is_idle, client).await?;
diff --git a/watchers/src/watchers/gnome_wayland.rs b/watchers/src/watchers/gnome_wayland.rs
new file mode 100644
index 0000000..3f2e302
--- /dev/null
+++ b/watchers/src/watchers/gnome_wayland.rs
@@ -0,0 +1,38 @@
+// The extension may not be loaded and available right away in Gnome, this mod will retry a few times.
+pub trait GnomeWatcher {
+    async fn load() -> anyhow::Result<Self>
+    where
+        Self: Sized;
+}
+
+fn is_gnome() -> bool {
+    if let Ok(de) = std::env::var("XDG_CURRENT_DESKTOP") {
+        de.to_lowercase().contains("gnome")
+    } else {
+        false
+    }
+}
+
+fn is_wayland() -> bool {
+    std::env::var("WAYLAND_DISPLAY").is_ok()
+        && !std::env::var("XDG_SESSION_TYPE")
+            .unwrap_or("".into())
+            .to_lowercase()
+            .contains("x11")
+}
+
+pub async fn load_watcher<T: GnomeWatcher>() -> anyhow::Result<T> {
+    if is_gnome() && is_wayland() {
+        let mut watcher = Err(anyhow::anyhow!(""));
+        for _ in 0..3 {
+            watcher = T::load().await;
+            if watcher.is_err() {
+                tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
+            }
+        }
+
+        watcher
+    } else {
+        T::load().await
+    }
+}
diff --git a/watchers/src/watchers/gnome_window.rs b/watchers/src/watchers/gnome_window.rs
index 8769369..4f5bebf 100644
--- a/watchers/src/watchers/gnome_window.rs
+++ b/watchers/src/watchers/gnome_window.rs
@@ -5,7 +5,7 @@ use serde::Deserialize;
 use std::sync::Arc;
 use zbus::Connection;
 
-use super::Watcher;
+use super::{gnome_wayland::load_watcher, gnome_wayland::GnomeWatcher, Watcher};
 
 pub struct WindowWatcher {
     dbus_connection: Connection,
@@ -38,7 +38,7 @@ impl WindowWatcher {
                     .body::<String>()
                     .with_context(|| "DBus interface cannot be parsed as string")?;
                 serde_json::from_str(&json).with_context(|| {
-                    "DBus interface org.gnome.shell.extensions.FocusedWindow returned wrong JSON"
+                    format!("DBus interface org.gnome.shell.extensions.FocusedWindow returned wrong JSON: {json}")
                 })
             }
             Err(e) => {
@@ -79,17 +79,24 @@ impl WindowWatcher {
     }
 }
 
-#[async_trait]
-impl Watcher for WindowWatcher {
-    async fn new(_: &Arc<ReportClient>) -> anyhow::Result<Self> {
+impl GnomeWatcher for WindowWatcher {
+    async fn load() -> anyhow::Result<Self> {
         let watcher = Self {
             dbus_connection: Connection::session().await?,
             last_app_id: String::new(),
             last_title: String::new(),
         };
+        watcher.get_window_data().await?;
 
         Ok(watcher)
     }
+}
+
+#[async_trait]
+impl Watcher for WindowWatcher {
+    async fn new(_: &Arc<ReportClient>) -> anyhow::Result<Self> {
+        load_watcher().await
+    }
 
     async fn run_iteration(&mut self, client: &Arc<ReportClient>) -> anyhow::Result<()> {
         self.send_active_window(client).await