fix: show repository owner avatars for workspaces on SSH devices

crates/herdr-gpui/src/avatars.rs

@@ -1,111 +1,332 @@
 use gpui::{Image, ImageFormat};
 mod cache;
+mod links;
+use herdr_client::{ConnectTarget, protocol::ClientShellWorkspace};
 use std::{
     collections::HashMap,
-    sync::{Arc, mpsc},
+    sync::{Arc, Mutex, mpsc},
     thread,
-    time::{Duration, SystemTime},
+    time::{Duration, Instant, SystemTime},
 };
 
 const MAX_REQUESTS: usize = 128;
 
 /// A bound on each remote path segment, above GitHub's own name limits.
 const MAX_SEGMENT: usize = 100;
 
+/// The whole budget for reading one remote checkout's Git configuration.
+const REMOTE_TIMEOUT: Duration = Duration::from_secs(10);
+
+/// How long a checkout whose host could not be read, while its device stayed
+/// connected, waits before a sidebar update may ask again. A device that
+/// reconnects is asked at once instead (see `Avatars::follow`).
+const RETRY_AFTER: Duration = Duration::from_secs(60);
+
+/// SSH reads that may run at once. A host that stops answering holds a
+/// reader for up to `REMOTE_TIMEOUT`, so a few readers keep one slow host
+/// from stalling every other device's avatars.
+const REMOTE_READERS: usize = 4;
+
+/// A checkout whose repository owner's avatar a sidebar row shows.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub(crate) enum Checkout<S> {
+    /// A workspace directory on this machine, read by local Git.
+    Local(S),
+    /// The Git directory a saved SSH device's daemon reported, read there.
+    Ssh { target: S, git_dir: S },
+}
+
+impl<'a> Checkout<&'a str> {
+    /// The checkout behind `workspace` on the device reached by `target`.
+    /// Only the local daemon's directories are read by this machine's Git; a
+    /// saved SSH device is read over SSH by the Git directory its daemon
+    /// reported, so it needs worktree metadata. Other devices may be forwarded
+    /// from an unknown machine and get none.
+    pub(crate) fn of(
+        local: bool,
+        target: &'a ConnectTarget,
+        workspace: &'a ClientShellWorkspace,
+    ) -> Option<Self> {
+        match target {
+            ConnectTarget::Ssh { target, .. } => {
+                workspace.worktree.as_ref().map(|tree| Self::Ssh {
+                    target,
+                    git_dir: &tree.key,
+                })
+            }
+            _ if local => Some(Self::Local(&workspace.new_workspace_cwd)),
+            _ => None,
+        }
+    }
+
+    fn to_owned(self) -> Checkout<String> {
+        match self {
+            Self::Local(cwd) => Checkout::Local(cwd.to_owned()),
+            Self::Ssh { target, git_dir } => Checkout::Ssh {
+                target: target.to_owned(),
+                git_dir: git_dir.to_owned(),
+            },
+        }
+    }
+}
+
+impl Checkout<String> {
+    fn as_deref(&self) -> Checkout<&str> {
+        match self {
+            Self::Local(cwd) => Checkout::Local(cwd),
+            Self::Ssh { target, git_dir } => Checkout::Ssh { target, git_dir },
+        }
+    }
+}
+
+/// What is known about one checkout's avatar.
+#[derive(Clone, Debug)]
+enum Slot {
+    /// Requested at `asked` and not yet answered.
+    Pending { asked: Instant },
+    /// The owner's avatar, or `None` for a checkout without a GitHub owner.
+    Resolved(Option<Arc<Image>>),
+    /// The host could not be read; asked again from `retry_at`, or as soon as
+    /// its device reconnects.
+    Unreachable { retry_at: Instant },
+}
+
+/// A worker's answer for one checkout: an error only when the host could
+/// not be read, never when it was read and has no GitHub owner.
+type Lookup = herdr_client::Result<Option<Arc<Image>>>;
+
+type Results = mpsc::SyncSender<(Checkout<String>, Lookup)>;
+type Refresh = mpsc::SyncSender<(Checkout<String>, String)>;
+
 /// Memory front end to the shared public-image disk cache.
 pub struct Avatars {
-    requests: mpsc::SyncSender<String>,
-    results: mpsc::Receiver<(String, Option<Arc<Image>>)>,
+    // Local Git answers at once while an SSH read may take `REMOTE_TIMEOUT`,
+    // so each has its own queue and a slow host never delays a local avatar.
+    local_requests: mpsc::SyncSender<Checkout<String>>,
+    remote_requests: mpsc::SyncSender<Checkout<String>>,
+    results: mpsc::Receiver<(Checkout<String>, Lookup)>,
     // Presence deduplicates pending requests as well as resolved hits and misses.
-    images: HashMap<String, Option<Arc<Image>>>,
+    local: HashMap<String, Slot>,
+    // Keyed by SSH target, then Git directory, so a row looks up without allocating.
+    remote: HashMap<String, HashMap<String, Slot>>,
+    // Each connected SSH endpoint's host and when its connection came up.
+    links: HashMap<String, (String, Instant)>,
+    offline: links::Offline,
 }
 
 impl Avatars {
     pub fn new() -> Self {
-        let (requests, incoming) = mpsc::sync_channel::<String>(MAX_REQUESTS);
+        let (local_requests, local_incoming) = mpsc::sync_channel(MAX_REQUESTS);
+        let (remote_requests, remote_incoming) = mpsc::sync_channel(MAX_REQUESTS);
         let (outgoing, results) = mpsc::sync_channel(MAX_REQUESTS * 2);
-        let (refresh, network) = mpsc::sync_channel::<(String, String)>(MAX_REQUESTS);
+        let (refresh, network) = mpsc::sync_channel::<(Checkout<String>, String)>(MAX_REQUESTS);
         let updated = outgoing.clone();
         thread::spawn(move || {
             let agent = avatar_agent();
             let mut owners = HashMap::new();
-            for (cwd, url) in network {
+            for (checkout, url) in network {
                 let image = owners
                     .entry(url.clone())
                     .or_insert_with(|| download(&agent, &url))
                     .clone();
                 if let Some(image) = image
-                    && updated.send((cwd, Some(image))).is_err()
+                    && updated.send((checkout, Ok(Some(image)))).is_err()
                 {
                     break;
                 }
             }
         });
-        // Deliberately detach: dropping the sender ends the worker without a UI join.
+        let offline = links::Offline::default();
+        // Deliberately detach: dropping the senders ends the workers without a UI join.
+        let remote_incoming = Arc::new(Mutex::new(remote_incoming));
+        for _ in 0..REMOTE_READERS {
+            let (incoming, offline, outgoing, refresh) = (
+                remote_incoming.clone(),
+                offline.clone(),
+                outgoing.clone(),
+                refresh.clone(),
+            );
+            thread::spawn(move || {
+                // The lock is held only while waiting, never across a read.
+                while let Some(checkout) = incoming.lock().ok().and_then(|rx| rx.recv().ok()) {
+                    if !answer(checkout, &offline, &outgoing, &refresh) {
+                        break;
+                    }
+                }
+            });
+        }
+        let local_offline = offline.clone();
         thread::spawn(move || {
-            for cwd in incoming {
-                let url = repo_owner(&cwd)
-                    .map(|owner| format!("https://avatars.githubusercontent.com/{owner}?size=48"));
-                let cached = url.as_deref().and_then(cached_image);
-                let fresh = cached.as_ref().is_some_and(|(_, fresh)| *fresh);
-                if outgoing
-                    .send((cwd.clone(), cached.map(|(image, _)| image)))
-                    .is_err()
-                {
+            for checkout in local_incoming {
+                if !answer(checkout, &local_offline, &outgoing, &refresh) {
                     break;
                 }
-                if !fresh && let Some(url) = url {
-                    let _ = refresh.try_send((cwd, url));
-                }
             }
         });
         Self {
-            requests,
+            local_requests,
+            remote_requests,
             results,
-            images: HashMap::new(),
+            local: HashMap::new(),
+            remote: HashMap::new(),
+            links: HashMap::new(),
+            offline,
         }
     }
 
-    /// Enqueue each exact cwd at most once, including failed resolutions.
-    pub fn request(&mut self, cwd: &str) {
-        if self.images.len() >= MAX_REQUESTS {
-            return;
+    /// Enqueue each exact checkout at most once, including ones without a
+    /// GitHub owner. A host that could not be read is asked again once its
+    /// retry time has passed at `now`.
+    pub(crate) fn request(&mut self, checkout: Checkout<&str>, now: Instant) {
+        match self.slot(checkout) {
+            Some(Slot::Unreachable { retry_at }) if now >= *retry_at => {}
+            Some(_) => return,
+            None => {
+                let len = self.local.len() + self.remote.values().map(HashMap::len).sum::<usize>();
+                if len >= MAX_REQUESTS {
+                    return;
+                }
+            }
         }
-        if let std::collections::hash_map::Entry::Vacant(entry) = self.images.entry(cwd.to_owned())
-            && self.requests.try_send(entry.key().clone()).is_ok()
-        {
-            entry.insert(None);
+        let requests = match checkout {
+            Checkout::Local(_) => &self.local_requests,
+            Checkout::Ssh { .. } => &self.remote_requests,
+        };
+        if requests.try_send(checkout.to_owned()).is_ok() {
+            self.insert(checkout.to_owned(), Slot::Pending { asked: now });
         }
     }
 
-    /// Drain completed requests without blocking. Misses also count as updates.
-    pub fn poll(&mut self) -> bool {
+    /// Drain completed requests without blocking. Misses also count as
+    /// updates; a host that could not be read is retried after `now`.
+    pub fn poll(&mut self, now: Instant) -> bool {
         let mut changed = false;
-        for (cwd, image) in self.results.try_iter() {
-            self.images.insert(cwd, image);
+        while let Ok((checkout, lookup)) = self.results.try_recv() {
+            // A read that failed after its device reconnected missed that
+            // connection's snapshot, which found it still pending, so it is
+            // asked again here rather than waiting for another snapshot.
+            let reconnected = lookup.is_err() && self.reconnected_since_asked(checkout.as_deref());
+            let slot = match lookup {
+                Ok(image) => Slot::Resolved(image),
+                // Stopped because its device disconnected: due again as soon
+                // as one is connected, since only then is it asked for.
+                Err(herdr_client::Error::SshCancelled) => Slot::Unreachable { retry_at: now },
+                Err(_) if reconnected => Slot::Unreachable { retry_at: now },
+                Err(_) => Slot::Unreachable {
+                    retry_at: now + RETRY_AFTER,
+                },
+            };
+            if reconnected {
+                self.insert(checkout.clone(), slot);
+                self.request(checkout.as_deref(), now);
+            } else {
+                self.insert(checkout, slot);
+            }
             changed = true;
         }
         changed
     }
 
-    pub fn image(&self, cwd: &str) -> Option<Arc<Image>> {
-        self.images.get(cwd).cloned().flatten()
+    pub(crate) fn image(&self, checkout: Checkout<&str>) -> Option<Arc<Image>> {
+        match self.slot(checkout)? {
+            Slot::Resolved(image) => image.clone(),
+            Slot::Pending { .. } | Slot::Unreachable { .. } => None,
+        }
+    }
+
+    fn slot(&self, checkout: Checkout<&str>) -> Option<&Slot> {
+        match checkout {
+            Checkout::Local(cwd) => self.local.get(cwd),
+            Checkout::Ssh { target, git_dir } => self.remote.get(target)?.get(git_dir),
+        }
+    }
+
+    fn insert(&mut self, checkout: Checkout<String>, image: Slot) {
+        match checkout {
+            Checkout::Local(cwd) => {
+                self.local.insert(cwd, image);
+            }
+            Checkout::Ssh { target, git_dir } => {
+                self.remote
+                    .entry(target)
+                    .or_default()
+                    .insert(git_dir, image);
+            }
+        }
+    }
+}
+
+/// Look up one checkout's owner and send what the disk cache holds for it,
+/// queueing a download when that is missing or stale. False once the UI end
+/// is gone.
+fn answer(
+    checkout: Checkout<String>,
+    offline: &links::Offline,
+    outgoing: &Results,
+    refresh: &Refresh,
+) -> bool {
+    let url = match repo_owner(checkout.as_deref(), offline) {
+        Ok(owner) => {
+            owner.map(|owner| format!("https://avatars.githubusercontent.com/{owner}?size=48"))
+        }
+        Err(error) => return outgoing.send((checkout, Err(error))).is_ok(),
+    };
+    let cached = url.as_deref().and_then(cached_image);
+    let fresh = cached.as_ref().is_some_and(|(_, fresh)| *fresh);
+    if outgoing
+        .send((checkout.clone(), Ok(cached.map(|(image, _)| image))))
+        .is_err()
+    {
+        return false;
+    }
+    if !fresh && let Some(url) = url {
+        let _ = refresh.try_send((checkout, url));
     }
+    true
 }
 
-fn repo_owner(cwd: &str) -> Option<String> {
+/// The GitHub owner behind a checkout's `upstream` remote, else its `origin`.
+/// Fails only when a remote host could not be read, or went offline before
+/// or during the read; local Git failures and absent remotes are `None`.
+fn repo_owner(
+    checkout: Checkout<&str>,
+    offline: &links::Offline,
+) -> herdr_client::Result<Option<String>> {
+    let deadline = Instant::now() + REMOTE_TIMEOUT;
     for key in ["remote.upstream.url", "remote.origin.url"] {
-        if let Ok(output) = herdr_client::new_command("git")
-            .args(["-C", cwd, "config", "--default", "", "--get", key])
-            .output()
-            && output.status.success()
-            && let Ok(text) = std::str::from_utf8(&output.stdout)
-            && let Some(owner) = github_owner(text.trim_end_matches(['\r', '\n']))
-        {
-            return Some(owner);
+        let url = match checkout {
+            Checkout::Local(cwd) => local_config(cwd, key),
+            // The upstream read may use most of the budget; origin gets only
+            // what is left.
+            // A read queued before its device disconnected never dials.
+            Checkout::Ssh { target, .. } if links::is_offline(offline, target) => {
+                return Err(herdr_client::Error::SshCancelled);
+            }
+            Checkout::Ssh { target, git_dir } => herdr_client::remote_config_value(
+                target,
+                git_dir,
+                key,
+                deadline
+                    .checked_duration_since(Instant::now())
+                    .ok_or(herdr_client::Error::SshTimeout)?,
+                || links::is_offline(offline, target),
+            )?,
+        };
+        if let Some(owner) = url.as_deref().and_then(github_owner) {
+            return Ok(Some(owner));
         }
     }
-    None
+    Ok(None)
+}
+
+fn local_config(cwd: &str, key: &str) -> Option<String> {
+    let output = herdr_client::new_command("git")
+        .args(["-C", cwd, "config", "--default", "", "--get", key])
+        .output()
+        .ok()
+        .filter(|output| output.status.success())?;
+    let text = std::str::from_utf8(&output.stdout).ok()?;
+    Some(text.trim_end_matches(['\r', '\n']).to_owned())
 }
 
 fn github_owner(remote: &str) -> Option<String> {