p2p_chat/storage/
seen.rs

1//! This module provides an interface and implementation for tracking messages that have been seen.
2use anyhow::Result;
3use async_trait::async_trait;
4use chrono::Utc;
5use sled::Db;
6use std::time::Duration;
7use uuid::Uuid;
8
9/// A trait for tracking seen messages.
10#[async_trait]
11pub trait SeenTracker {
12    /// Marks a message as seen.
13    ///
14    /// # Arguments
15    ///
16    /// * `msg_id` - The `Uuid` of the message to mark as seen.
17    ///
18    /// # Errors
19    ///
20    /// This function will return an error if the message cannot be marked as seen.
21    async fn mark_seen(&self, msg_id: Uuid) -> Result<()>;
22
23    /// Checks if a message has been seen.
24    ///
25    /// # Arguments
26    ///
27    /// * `msg_id` - The `Uuid` of the message to check.
28    ///
29    /// # Returns
30    ///
31    /// `true` if the message has been seen, `false` otherwise.
32    ///
33    /// # Errors
34    ///
35    /// This function will return an error if the seen status cannot be retrieved.
36    async fn is_seen(&self, msg_id: &Uuid) -> Result<bool>;
37
38    /// Cleans up old seen message records.
39    ///
40    /// # Arguments
41    ///
42    /// * `max_age` - The maximum age for seen records to be retained.
43    ///
44    /// # Errors
45    ///
46    /// This function will return an error if cleanup fails.
47    async fn cleanup_old(&self, max_age: Duration) -> Result<()>;
48}
49
50/// A `SeenTracker` implementation using `sled` for storage.
51pub struct SledSeenTracker {
52    tree: sled::Tree,
53}
54
55impl SledSeenTracker {
56    /// Creates a new `SledSeenTracker`.
57    ///
58    /// # Arguments
59    ///
60    /// * `db` - The `sled::Db` instance to use for storage.
61    ///
62    /// # Errors
63    ///
64    /// This function will return an error if the `seen` tree cannot be opened.
65    pub fn new(db: Db) -> Result<Self> {
66        let tree = db.open_tree("seen")?;
67        Ok(Self { tree })
68    }
69}
70
71#[async_trait]
72impl SeenTracker for SledSeenTracker {
73    async fn mark_seen(&self, msg_id: Uuid) -> Result<()> {
74        let key = msg_id.to_string();
75        let timestamp = Utc::now().timestamp_millis();
76
77        self.tree.insert(key.as_bytes(), &timestamp.to_be_bytes())?;
78        self.tree.flush_async().await?;
79        Ok(())
80    }
81
82    async fn is_seen(&self, msg_id: &Uuid) -> Result<bool> {
83        let key = msg_id.to_string();
84        Ok(self.tree.contains_key(key.as_bytes())?)
85    }
86
87    async fn cleanup_old(&self, max_age: Duration) -> Result<()> {
88        let cutoff = Utc::now().timestamp_millis() - max_age.as_millis() as i64;
89        let mut keys_to_remove = Vec::new();
90
91        for result in self.tree.iter() {
92            let (key, value) = result?;
93            if value.len() >= 8 {
94                let timestamp = i64::from_be_bytes(value[..8].try_into().unwrap());
95                if timestamp < cutoff {
96                    keys_to_remove.push(key.to_vec());
97                }
98            }
99        }
100
101        for key in keys_to_remove {
102            self.tree.remove(key)?;
103        }
104
105        self.tree.flush_async().await?;
106        Ok(())
107    }
108}