feat: debounce notifications

Author: link2xt

src/debouncer.rs

@@ -0,0 +1,126 @@
+//! # Debouncer for notifications.
+//!
+//! Sometimes the client application may be reinstalled
+//! while keeping the notification token.
+//! In this case the same token is stored twice
+//! for the same mailbox on a chatmail relay
+//! and is notified twice for the same message.
+//! Since it is not possible for the chatmail relay
+//! to deduplicate the tokens in this case
+//! as only the notification gateway
+//! can decrypt them, notification gateway needs
+//! to debounce notifications to the same token.
+
+use std::cmp::Reverse;
+use std::collections::{BinaryHeap, HashSet};
+use std::sync::Mutex;
+use std::time::{Duration, Instant};
+
+#[derive(Default)]
+pub(crate) struct Debouncer {
+    state: Mutex<DebouncerState>,
+}
+
+#[derive(Default)]
+struct DebouncerState {
+    /// Set of recently notified tokens.
+    ///
+    /// The tokens are stored in plaintext,
+    /// not hashed or encrypted.
+    /// No token is stored for a long time anyway.
+    tokens: HashSet<String>,
+
+    /// Binary heap storing tokens
+    /// sorted by the timestamp of the recent notifications.
+    ///
+    /// `Reverse` is used to turn max-heap into min-heap.
+    heap: BinaryHeap<Reverse<(Instant, String)>>,
+}
+
+impl DebouncerState {
+    /// Removes old entries for tokens that can be notified again.
+    fn cleanup(&mut self, now: Instant) {
+        loop {
+            let Some(Reverse((timestamp, token))) = self.heap.pop() else {
+                break;
+            };
+
+            if now.duration_since(timestamp) < Duration::from_secs(1) {
+                self.heap.push(Reverse((timestamp, token)));
+                break;
+            }
+
+            self.tokens.remove(&token);
+        }
+    }
+
+    fn is_debounced(&mut self, now: Instant, token: &String) -> bool {
+        self.cleanup(now);
+        self.tokens.contains(token)
+    }
+
+    fn notify(&mut self, now: Instant, token: String) {
+        if self.tokens.insert(token.clone()) {
+            self.heap.push(Reverse((now, token)));
+        }
+    }
+
+    fn count(&self) -> usize {
+        let res = self.tokens.len();
+        debug_assert_eq!(res, self.heap.len());
+        res
+    }
+}
+
+impl Debouncer {
+    /// Returns true if the token was notified recently
+    /// and should not be notified again.
+    pub(crate) fn is_debounced(&self, now: Instant, token: &String) -> bool {
+        let mut state = self.state.lock().unwrap();
+        state.is_debounced(now, token)
+    }
+
+    pub(crate) fn notify(&self, now: Instant, token: String) {
+        self.state.lock().unwrap().notify(now, token);
+    }
+
+    /// Returns number of currently debounced notification tokens.
+    ///
+    /// This is used for metrics to display the size of the set.
+    ///
+    /// This function does not remove expired tokens.
+    pub(crate) fn count(&self) -> usize {
+        self.state.lock().unwrap().count()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_debouncer() {
+        let mut now = Instant::now();
+
+        let debouncer = Debouncer::default();
+
+        let token1 = "foobar".to_string();
+        let token2 = "barbaz".to_string();
+
+        assert!(!debouncer.is_debounced(now, &token1));
+        assert!(!debouncer.is_debounced(now, &token2));
+        assert_eq!(debouncer.count(), 0);
+
+        debouncer.notify(now, token1.clone());
+
+        assert!(debouncer.is_debounced(now, &token1));
+        assert!(!debouncer.is_debounced(now, &token2));
+        assert_eq!(debouncer.count(), 1);
+
+        now += Duration::from_secs(5);
+
+        assert!(!debouncer.is_debounced(now, &token1));
+        assert!(!debouncer.is_debounced(now, &token2));
+        assert_eq!(debouncer.count(), 0);
+    }
+}

src/lib.rs

@@ -1,3 +1,4 @@
+mod debouncer;
 pub mod metrics;
 pub mod notifier;
 mod openpgp;

src/metrics.rs

@@ -30,6 +30,12 @@ pub struct Metrics {
     /// Number of successfully sent visible UBports notifications.
     pub ubports_notifications_total: Counter,
 
+    /// Number of debounced notifications.
+    pub debounced_notifications_total: Counter,
+
+    /// Number of tokens notified recently.
+    pub debounced_set_size: Gauge<i64, AtomicI64>,
+
     /// Number of successfully sent heartbeat notifications.
     pub heartbeat_notifications_total: Counter,
 
@@ -68,6 +74,20 @@ impl Metrics {
             ubports_notifications_total.clone(),
         );
 
+        let debounced_notifications_total = Counter::default();
+        registry.register(
+            "debounced_notifications",
+            "Number of debounced notifications",
+            debounced_notifications_total.clone(),
+        );
+
+        let debounced_set_size = Gauge::<i64, AtomicI64>::default();
+        registry.register(
+            "debounced_set_size",
+            "Number of tokens notified recently.",
+            debounced_set_size.clone(),
+        );
+
         let heartbeat_notifications_total = Counter::default();
         registry.register(
             "heartbeat_notifications",
@@ -101,6 +121,8 @@ impl Metrics {
             direct_notifications_total,
             fcm_notifications_total,
             ubports_notifications_total,
+            debounced_notifications_total,
+            debounced_set_size,
             heartbeat_notifications_total,
             heartbeat_registrations_total,
             heartbeat_tokens,

src/server.rs

@@ -10,6 +10,7 @@ use chrono::{Local, TimeDelta};
 use log::*;
 use serde::Deserialize;
 use std::str::FromStr;
+use std::time::Instant;
 
 use crate::metrics::Metrics;
 use crate::state::State;
@@ -293,6 +294,19 @@ async fn notify_device(
     }
 
     info!("Got direct notification for {device_token}.");
+    let now = Instant::now();
+    if state.debouncer().is_debounced(now, &device_token) {
+        let metrics = state.metrics();
+        metrics.debounced_notifications_total.inc();
+        metrics
+            .debounced_set_size
+            .set(state.debouncer().count() as i64);
+        return Ok(StatusCode::OK);
+    }
+    state.debouncer().notify(now, device_token.clone());
+    state.metrics()
+        .debounced_set_size
+        .set(state.debouncer().count() as i64);
     let device_token: NotificationToken = device_token.as_str().parse()?;
 
     let status_code = match device_token {

src/state.rs

@@ -6,6 +6,7 @@ use std::time::Duration;
 use a2::{Client, Endpoint};
 use anyhow::{Context as _, Result};
 
+use crate::debouncer::Debouncer;
 use crate::metrics::Metrics;
 use crate::openpgp::PgpDecryptor;
 use crate::schedule::Schedule;
@@ -36,6 +37,8 @@ pub struct InnerState {
     /// Decryptor for incoming tokens
     /// storing the secret keyring inside.
     openpgp_decryptor: PgpDecryptor,
+
+    debouncer: Debouncer,
 }
 
 impl State {
@@ -88,6 +91,7 @@ impl State {
                 interval,
                 fcm_authenticator,
                 openpgp_decryptor,
+                debouncer: Default::default(),
             }),
         })
     }
@@ -134,4 +138,8 @@ impl State {
     pub fn openpgp_decryptor(&self) -> &PgpDecryptor {
         &self.inner.openpgp_decryptor
     }
+
+    pub(crate) fn debouncer(&self) -> &Debouncer {
+        &self.inner.debouncer
+    }
 }