1use std::collections::HashSet;
4use std::hash::Hash;
5use std::sync::{Arc, PoisonError, RwLock};
6
7use crate::backend::NodeId;
8use hidpp::channel::HidppChannel;
9use tokio::sync::watch;
10
11use crate::{DeviceRoute, SharedChannel};
12
13struct Publication<Node, Channel> {
14 node: Node,
15 sequence: u64,
16 routes: Vec<DeviceRoute>,
17 channel: Channel,
18}
19
20struct NodeRegistry<Node, Channel> {
21 publications: Vec<Publication<Node, Channel>>,
22 next_sequence: u64,
23}
24
25impl<Node, Channel> Default for NodeRegistry<Node, Channel> {
26 fn default() -> Self {
27 Self {
28 publications: Vec::new(),
29 next_sequence: 0,
30 }
31 }
32}
33
34impl<Node: Eq, Channel> NodeRegistry<Node, Channel> {
35 fn replace_node(
36 &mut self,
37 node: Node,
38 routes: impl IntoIterator<Item = DeviceRoute>,
39 channel: Channel,
40 same_channel: fn(&Channel, &Channel) -> bool,
41 ) -> bool {
42 let routes: Vec<_> = routes.into_iter().collect();
43 if let Some(publication) = self
44 .publications
45 .iter_mut()
46 .find(|publication| publication.node == node)
47 {
48 let same_routes = publication
49 .routes
50 .iter()
51 .all(|route| routes.contains(route))
52 && routes
53 .iter()
54 .all(|route| publication.routes.contains(route));
55 if same_routes && same_channel(&publication.channel, &channel) {
56 return false;
57 }
58 publication.routes = routes;
59 publication.channel = channel;
60 return true;
61 }
62
63 let sequence = self.next_sequence;
64 self.next_sequence = self.next_sequence.wrapping_add(1);
65 self.publications.push(Publication {
66 node,
67 sequence,
68 routes,
69 channel,
70 });
71 true
72 }
73
74 fn remove_node(&mut self, node: &Node) -> bool {
75 let previous_len = self.publications.len();
76 self.publications
77 .retain(|publication| publication.node != *node);
78 self.publications.len() != previous_len
79 }
80
81 fn lookup(&self, route: &DeviceRoute) -> Option<&Channel> {
82 self.publications
83 .iter()
84 .filter(|publication| publication.routes.contains(route))
85 .min_by_key(|publication| publication.sequence)
86 .map(|publication| &publication.channel)
87 }
88
89 fn any_current(&self, mut predicate: impl FnMut(&DeviceRoute, &Channel) -> bool) -> bool {
90 self.publications.iter().any(|publication| {
91 publication.routes.iter().any(|route| {
92 predicate(route, &publication.channel)
93 && self
94 .lookup(route)
95 .is_some_and(|winner| std::ptr::eq(winner, &raw const publication.channel))
96 })
97 })
98 }
99}
100
101impl<Node: Eq + Hash, Channel> NodeRegistry<Node, Channel> {
102 fn retain_nodes(&mut self, nodes: &HashSet<Node>) -> bool {
103 let previous_len = self.publications.len();
104 self.publications
105 .retain(|publication| nodes.contains(&publication.node));
106 self.publications.len() != previous_len
107 }
108}
109
110struct Registry<Node, Channel> {
111 state: Arc<RwLock<NodeRegistry<Node, Channel>>>,
112 changes: watch::Sender<()>,
113 same_channel: fn(&Channel, &Channel) -> bool,
114}
115
116impl<Node, Channel> Clone for Registry<Node, Channel> {
117 fn clone(&self) -> Self {
118 Self {
119 state: Arc::clone(&self.state),
120 changes: self.changes.clone(),
121 same_channel: self.same_channel,
122 }
123 }
124}
125
126impl<Node, Channel: PartialEq> Default for Registry<Node, Channel> {
127 fn default() -> Self {
128 Self::new(PartialEq::eq)
129 }
130}
131
132impl<Node, Channel> Registry<Node, Channel> {
133 fn new(same_channel: fn(&Channel, &Channel) -> bool) -> Self {
134 let (changes, _) = watch::channel(());
135 Self {
136 state: Arc::new(RwLock::new(NodeRegistry::default())),
137 changes,
138 same_channel,
139 }
140 }
141
142 fn subscribe(&self) -> watch::Receiver<()> {
143 self.changes.subscribe()
144 }
145
146 fn notify_change(&self) {
147 self.changes.send_modify(|()| {});
150 }
151}
152
153impl<Node: Eq, Channel> Registry<Node, Channel> {
154 fn replace_node(
155 &self,
156 node: Node,
157 routes: impl IntoIterator<Item = DeviceRoute>,
158 channel: Channel,
159 ) {
160 let changed = self
161 .state
162 .write()
163 .unwrap_or_else(PoisonError::into_inner)
164 .replace_node(node, routes, channel, self.same_channel);
165 if changed {
166 self.notify_change();
167 }
168 }
169
170 fn remove_node(&self, node: &Node) {
171 let changed = self
172 .state
173 .write()
174 .unwrap_or_else(PoisonError::into_inner)
175 .remove_node(node);
176 if changed {
177 self.notify_change();
178 }
179 }
180}
181
182impl<Node: Eq + Hash, Channel> Registry<Node, Channel> {
183 fn retain_nodes(&self, nodes: &HashSet<Node>) {
184 let changed = self
185 .state
186 .write()
187 .unwrap_or_else(PoisonError::into_inner)
188 .retain_nodes(nodes);
189 if changed {
190 self.notify_change();
191 }
192 }
193}
194
195impl<Node: Eq, Channel: Clone> Registry<Node, Channel> {
196 fn lookup(&self, route: &DeviceRoute) -> Option<Channel> {
197 self.state.read().ok()?.lookup(route).cloned()
198 }
199}
200
201impl<Node: Eq, Channel> Registry<Node, Channel> {
202 fn any_current(&self, predicate: impl FnMut(&DeviceRoute, &Channel) -> bool) -> bool {
203 self.state
204 .read()
205 .is_ok_and(|state| state.any_current(predicate))
206 }
207}
208
209#[derive(Clone)]
215pub struct ChannelRegistry {
216 inner: Registry<NodeId, Arc<HidppChannel>>,
217}
218
219impl Default for ChannelRegistry {
220 fn default() -> Self {
221 Self {
222 inner: Registry::new(Arc::ptr_eq),
223 }
224 }
225}
226
227impl ChannelRegistry {
228 pub(crate) fn replace_node(
231 &self,
232 node: NodeId,
233 routes: impl IntoIterator<Item = DeviceRoute>,
234 channel: Arc<HidppChannel>,
235 ) {
236 self.inner.replace_node(node, routes, channel);
237 }
238
239 pub(crate) fn remove_node(&self, node: &NodeId) {
241 self.inner.remove_node(node);
242 }
243
244 pub(crate) fn retain_nodes(&self, nodes: &HashSet<NodeId>) {
246 self.inner.retain_nodes(nodes);
247 }
248
249 #[must_use]
256 pub fn subscribe(&self) -> watch::Receiver<()> {
257 self.inner.subscribe()
258 }
259
260 #[must_use]
262 pub fn lookup(&self, route: &DeviceRoute) -> Option<SharedChannel> {
263 self.inner
264 .lookup(route)
265 .map(|channel| SharedChannel::new(channel, route.clone()))
266 }
267
268 #[must_use]
271 pub fn is_current(&self, shared: &SharedChannel) -> bool {
272 self.inner.any_current(|route, channel| {
273 shared.matches(route) && Arc::ptr_eq(channel, shared.channel())
274 })
275 }
276}
277
278#[cfg(test)]
279mod tests {
280 use std::collections::HashSet;
281 use std::panic::{AssertUnwindSafe, catch_unwind};
282 use std::sync::Arc;
283 use std::time::Duration;
284
285 use crate::DeviceRoute;
286
287 use super::{PoisonError, Registry};
288
289 impl<Node, Channel> Registry<Node, Channel> {
290 fn poison_for_test(&self) {
291 let _ = catch_unwind(AssertUnwindSafe(|| {
292 let _guard = self.state.write().unwrap_or_else(PoisonError::into_inner);
293 panic!("poison registry for test");
294 }));
295 }
296 }
297
298 impl<Node: Eq, Channel: Clone> Registry<Node, Channel> {
299 fn publisher_lookup_for_test(&self, route: &DeviceRoute) -> Option<Channel> {
300 self.state
301 .read()
302 .unwrap_or_else(PoisonError::into_inner)
303 .lookup(route)
304 .cloned()
305 }
306 }
307
308 fn direct(product_id: u16) -> DeviceRoute {
309 DeviceRoute::Direct {
310 vendor_id: 0x046d,
311 product_id,
312 }
313 }
314
315 fn bolt(uid: &str, slot: u8) -> DeviceRoute {
316 DeviceRoute::Bolt {
317 receiver_uid: uid.into(),
318 slot,
319 }
320 }
321
322 #[test]
323 fn lookup_rejects_every_non_exact_route_field() {
324 let registry = Registry::<u8, &'static str>::default();
325 registry.replace_node(1, [bolt("AABB", 2)], "channel-a");
326
327 assert_eq!(registry.lookup(&bolt("AABB", 2)), Some("channel-a"));
328 assert_eq!(registry.lookup(&bolt("AABB", 3)), None);
329 assert_eq!(registry.lookup(&bolt("CCDD", 2)), None);
330 assert_eq!(registry.lookup(&direct(0xb35b)), None);
331 }
332
333 #[test]
334 fn one_node_can_publish_multiple_receiver_slots() {
335 let registry = Registry::<u8, &'static str>::default();
336 registry.replace_node(1, [bolt("AABB", 1), bolt("AABB", 4)], "receiver-channel");
337
338 assert_eq!(registry.lookup(&bolt("AABB", 1)), Some("receiver-channel"));
339 assert_eq!(registry.lookup(&bolt("AABB", 4)), Some("receiver-channel"));
340 }
341
342 #[test]
343 fn current_check_uses_only_the_exact_winning_publication() {
344 let route = direct(0xb35b);
345 let registry = Registry::<u8, &'static str>::default();
346 registry.replace_node(1, [route.clone()], "a");
347 registry.replace_node(2, [route.clone()], "b");
348
349 assert!(
350 registry.any_current(|candidate, channel| { candidate == &route && *channel == "a" })
351 );
352 assert!(
353 !registry.any_current(|candidate, channel| { candidate == &route && *channel == "b" })
354 );
355 }
356
357 #[test]
358 fn same_route_with_a_different_arc_is_not_current() {
359 let route = direct(0xb35b);
360 let published = Arc::new(());
361 let stale = Arc::new(());
362 let registry = Registry::<u8, Arc<()>>::default();
363 registry.replace_node(1, [route.clone()], Arc::clone(&published));
364
365 assert!(registry.any_current(|candidate, channel| {
366 candidate == &route && Arc::ptr_eq(channel, &published)
367 }));
368 assert!(!registry.any_current(|candidate, channel| {
369 candidate == &route && Arc::ptr_eq(channel, &stale)
370 }));
371 }
372
373 #[test]
374 fn replacing_winner_preserves_priority_then_removal_promotes_next_owner() {
375 let route = direct(0xb35b);
376 let registry = Registry::<u8, &'static str>::default();
377 registry.replace_node(1, [route.clone()], "a-v1");
378 registry.replace_node(2, [route.clone()], "b");
379
380 assert_eq!(registry.lookup(&route), Some("a-v1"));
381
382 registry.replace_node(1, [route.clone()], "a-v2");
383 assert_eq!(registry.lookup(&route), Some("a-v2"));
384
385 registry.remove_node(&1);
386 assert_eq!(registry.lookup(&route), Some("b"));
387 }
388
389 #[test]
390 fn replacement_and_removal_wake_every_subscriber() {
391 let route = direct(0xb35b);
392 let first_channel = Arc::new(());
393 let replacement = Arc::new(());
394 let registry = Registry::<u8, Arc<()>>::new(Arc::ptr_eq);
395 registry.replace_node(1, [route.clone()], Arc::clone(&first_channel));
396 let mut first = registry.subscribe();
397 let mut second = registry.subscribe();
398
399 registry.replace_node(1, [route.clone()], Arc::clone(&replacement));
400
401 assert!(first.has_changed().expect("registry sender remains alive"));
402 assert!(second.has_changed().expect("registry sender remains alive"));
403 first.borrow_and_update();
404 second.borrow_and_update();
405 assert!(!registry.any_current(|candidate, channel| {
406 candidate == &route && Arc::ptr_eq(channel, &first_channel)
407 }));
408 assert!(registry.any_current(|candidate, channel| {
409 candidate == &route && Arc::ptr_eq(channel, &replacement)
410 }));
411
412 registry.remove_node(&1);
413
414 assert!(first.has_changed().expect("registry sender remains alive"));
415 assert!(second.has_changed().expect("registry sender remains alive"));
416 assert_eq!(registry.lookup(&route), None);
417 }
418
419 #[test]
420 fn no_op_publications_do_not_wake_subscribers() {
421 let first_route = bolt("AABB", 1);
422 let second_route = bolt("AABB", 2);
423 let channel = Arc::new(());
424 let registry = Registry::<u8, Arc<()>>::new(Arc::ptr_eq);
425 registry.replace_node(
426 1,
427 [first_route.clone(), second_route.clone()],
428 Arc::clone(&channel),
429 );
430 let changes = registry.subscribe();
431
432 registry.replace_node(1, [second_route, first_route], Arc::clone(&channel));
433 registry.remove_node(&2);
434 registry.retain_nodes(&HashSet::from([1]));
435
436 assert!(
437 !changes
438 .has_changed()
439 .expect("registry sender remains alive")
440 );
441 }
442
443 #[tokio::test]
444 async fn change_between_current_check_and_wait_is_not_missed() {
445 let route = direct(0xb35b);
446 let first_channel = Arc::new(());
447 let replacement = Arc::new(());
448 let registry = Registry::<u8, Arc<()>>::new(Arc::ptr_eq);
449 registry.replace_node(1, [route.clone()], Arc::clone(&first_channel));
450
451 let mut changes = registry.subscribe();
454 assert!(registry.any_current(|candidate, channel| {
455 candidate == &route && Arc::ptr_eq(channel, &first_channel)
456 }));
457 registry.replace_node(1, [route.clone()], Arc::clone(&replacement));
458
459 tokio::time::timeout(Duration::from_millis(100), changes.changed())
460 .await
461 .expect("a pre-wait publication must be observed")
462 .expect("registry sender remains alive");
463 assert!(registry.any_current(|candidate, channel| {
464 candidate == &route && Arc::ptr_eq(channel, &replacement)
465 }));
466 }
467
468 #[test]
469 fn replacing_one_node_is_atomic_and_does_not_touch_another() {
470 let registry = Registry::<u8, &'static str>::default();
471 registry.replace_node(1, [bolt("A", 1), bolt("A", 2)], "a");
472 registry.replace_node(2, [bolt("B", 1)], "b");
473
474 registry.replace_node(1, [bolt("A", 3)], "a-new");
475
476 assert_eq!(registry.lookup(&bolt("A", 1)), None);
477 assert_eq!(registry.lookup(&bolt("A", 2)), None);
478 assert_eq!(registry.lookup(&bolt("A", 3)), Some("a-new"));
479 assert_eq!(registry.lookup(&bolt("B", 1)), Some("b"));
480 }
481
482 #[test]
483 fn retaining_nodes_removes_only_absent_owners() {
484 let registry = Registry::<u8, &'static str>::default();
485 registry.replace_node(1, [direct(0xb35b)], "a");
486 registry.replace_node(2, [direct(0xb36b)], "b");
487
488 registry.retain_nodes(&HashSet::from([2]));
489
490 assert_eq!(registry.lookup(&direct(0xb35b)), None);
491 assert_eq!(registry.lookup(&direct(0xb36b)), Some("b"));
492 }
493
494 #[test]
495 fn poisoned_read_fails_closed_but_publishers_can_clean_up() {
496 let registry = Registry::<u8, &'static str>::default();
497 registry.replace_node(1, [direct(0xb35b)], "a");
498 registry.poison_for_test();
499
500 assert_eq!(registry.lookup(&direct(0xb35b)), None);
501 assert!(!registry.any_current(|_, _| true));
502
503 registry.remove_node(&1);
504 registry.replace_node(2, [direct(0xb36b)], "b");
505 registry.retain_nodes(&HashSet::from([2]));
506
507 assert_eq!(registry.publisher_lookup_for_test(&direct(0xb35b)), None);
508 assert_eq!(
509 registry.publisher_lookup_for_test(&direct(0xb36b)),
510 Some("b")
511 );
512 }
513}