diff --git a/src/bin/rwho.rs b/src/bin/rwho.rs index 2beb423..38c6db4 100644 --- a/src/bin/rwho.rs +++ b/src/bin/rwho.rs @@ -1,3 +1,5 @@ +use std::collections::BTreeMap; + use anyhow::Context; use clap::{CommandFactory, Parser}; use clap_complete::{Shell, generate}; @@ -63,7 +65,7 @@ async fn main() -> anyhow::Result<()> { .await .expect("Failed to connect to rwhod server"); - let mut reply = conn + let reply = conn .rwho(max_idle_seconds) .await .context("Failed to send rwho request")? @@ -79,23 +81,40 @@ async fn main() -> anyhow::Result<()> { } })?; - reply.sort_by(|(host, user), (host2, user2)| { - user.user_id - .cmp(&user2.user_id) - .then_with(|| host.cmp(host2)) - .then_with(|| user.tty.cmp(&user2.tty)) - }); - if args.json { - println!("{}", serde_json::to_string_pretty(&reply).unwrap()); - } else if args.old { - old_format_user_entries(&reply) - .iter() - .for_each(|line| println!("{}", line)); + let sorted: BTreeMap> = reply + .into_iter() + .map(|(hostname, mut users)| { + users.sort_by(|a, b| a.user_id.cmp(&b.user_id).then_with(|| a.tty.cmp(&b.tty))); + (hostname, users) + }) + .collect(); + println!("{}", serde_json::to_string_pretty(&sorted).unwrap()); } else { - old_format_user_entries(&reply) - .iter() - .for_each(|line| println!("{}", line)); + let mut entries: Vec<(String, WhodUserEntry)> = reply + .into_iter() + .flat_map(|(hostname, users)| { + users.into_iter().map(move |user| (hostname.clone(), user)) + }) + .collect(); + + entries.sort_by(|(host, user), (host2, user2)| { + user.user_id + .cmp(&user2.user_id) + .then_with(|| host.cmp(host2)) + .then_with(|| user.tty.cmp(&user2.tty)) + }); + + if args.old { + old_format_user_entries(&entries) + .iter() + .for_each(|line| println!("{}", line)); + } else { + // TODO: add a newer and nicer format here + old_format_user_entries(&entries) + .iter() + .for_each(|line| println!("{}", line)); + } } Ok(()) diff --git a/src/server/varlink_api.rs b/src/server/varlink_api.rs index 6ed6396..36239bc 100644 --- a/src/server/varlink_api.rs +++ b/src/server/varlink_api.rs @@ -1,4 +1,4 @@ -use std::{os::fd::OwnedFd, time::Duration}; +use std::{collections::HashMap, os::fd::OwnedFd, time::Duration}; use anyhow::Context; use futures_util::stream; @@ -66,7 +66,7 @@ pub enum VarlinkRwhodClientResponse { Ruptime(VarlinkRuptimeResponse), } -pub type VarlinkRwhoResponse = Vec<(String, WhodUserEntry)>; +pub type VarlinkRwhoResponse = HashMap>; pub type VarlinkRuptimeResponse = Vec; #[derive(Debug, Clone, PartialEq, ReplyError)] @@ -180,19 +180,19 @@ impl VarlinkRoowhoo2ClientServer { tracing::debug!(?max_idle_time, "Handling Rwho request"); let store = self.whod_status_store.read().await; - let mut all_user_entries = Vec::with_capacity(store.len()); - for status_update in store.values() { - all_user_entries.extend( - status_update + store + .values() + .filter_map(|status_update| { + let users: Vec = status_update .users .iter() .filter(|user| max_idle_time.is_none_or(|max| user.idle_time < max)) .cloned() - .map(|user| (status_update.hostname.clone(), user)), - ); - } + .collect(); - all_user_entries + (!users.is_empty()).then(|| (status_update.hostname.clone(), users)) + }) + .collect() } async fn handle_ruptime_request(