autopush_common/db/redis/redis_client/
mod.rs

1use std::collections::HashSet;
2use std::fmt;
3use std::fmt::Display;
4use std::str::FromStr;
5use std::sync::Arc;
6use std::time::SystemTime;
7
8use async_trait::async_trait;
9use cadence::{CountedExt, StatsdClient};
10use deadpool_redis::Config;
11use deadpool_redis::redis::{AsyncCommands, SetExpiry, SetOptions, pipe};
12use uuid::Uuid;
13
14use crate::db::redis::StorableNotification;
15use crate::db::{
16    DbSettings, Notification, User,
17    client::{DbClient, FetchMessageResponse},
18    error::{DbError, DbResult},
19};
20use crate::util::{ms_since_epoch, sec_since_epoch};
21
22use super::RedisDbSettings;
23
24fn now_secs() -> u64 {
25    // Return the current time in seconds since EPOCH
26    SystemTime::now()
27        .duration_since(SystemTime::UNIX_EPOCH)
28        .unwrap()
29        .as_secs()
30}
31
32/// Semi convenience wrapper to ensure that the UAID is formatted and displayed consistently.
33struct Uaid<'a>(&'a Uuid);
34
35impl<'a> Display for Uaid<'a> {
36    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
37        write!(f, "{}", self.0.as_hyphenated())
38    }
39}
40
41impl<'a> From<Uaid<'a>> for String {
42    fn from(uaid: Uaid) -> String {
43        uaid.0.as_hyphenated().to_string()
44    }
45}
46
47struct ChannelId<'a>(&'a Uuid);
48
49impl<'a> Display for ChannelId<'a> {
50    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
51        write!(f, "{}", self.0.as_hyphenated())
52    }
53}
54
55impl<'a> From<ChannelId<'a>> for String {
56    fn from(chid: ChannelId) -> String {
57        chid.0.as_hyphenated().to_string()
58    }
59}
60
61#[derive(Clone)]
62/// Wrapper for the Redis connection
63pub struct RedisClientImpl {
64    /// Database connector string
65    pub pool: deadpool_redis::Pool,
66    /// Metrics client
67    metrics: Arc<StatsdClient>,
68    router_opts: SetOptions,
69    // Default notification options (mostly TTL)
70    notification_opts: SetOptions,
71}
72
73impl RedisClientImpl {
74    pub fn new(metrics: Arc<StatsdClient>, settings: &DbSettings) -> DbResult<Self> {
75        debug!("🐰 New redis client");
76        let dsn = settings.dsn.clone().ok_or(DbError::General(
77            "Redis DSN not configured. Set `db_dsn` to `redis://HOST:PORT` in settings.".to_owned(),
78        ))?;
79        let db_settings = RedisDbSettings::try_from(settings.db_settings.as_ref())?;
80        info!("🐰 {:#?}", db_settings);
81        let router_ttl_secs = db_settings.router_ttl.unwrap_or_default().as_secs();
82        let notification_ttl_secs = db_settings.notification_ttl.unwrap_or_default().as_secs();
83
84        let config = Config::from_url(dsn);
85        let pool = config
86            .builder()
87            .map_err(|e| DbError::General(format!("Could not create Redis pool: {:?}", e)))?
88            .create_timeout(db_settings.create_timeout)
89            .runtime(deadpool_redis::Runtime::Tokio1)
90            .build()
91            .map_err(|e| DbError::General(format!("Could not create Redis pool: {:?}", e)))?;
92        /* We have the option of using either a OneCell wrapped get_multiplexed_async_connection or
93         * a pool. Reliability already uses a pool, so for consistency we use a pool here as well.
94         */
95
96        // We specify different TTLs for router vs message.
97        Ok(Self {
98            pool,
99            metrics,
100            router_opts: SetOptions::default().with_expiration(SetExpiry::EX(router_ttl_secs)),
101            notification_opts: SetOptions::default()
102                .with_expiration(SetExpiry::EX(notification_ttl_secs)),
103        })
104    }
105
106    /// Return a [ConnectionLike], which implement redis [Commands] and can be
107    /// used in pipes.
108    ///
109    /// Pools also return a ConnectionLike, so we can add support for pools later.
110    async fn connection(&self) -> DbResult<deadpool_redis::Connection> {
111        self.pool.get().await.map_err(|e| {
112            DbError::RedisError(redis::RedisError::from((
113                redis::ErrorKind::Io,
114                "Could not get Redis connection from pool",
115                format!("{:?}", e),
116            )))
117        })
118    }
119
120    fn user_key(&self, uaid: &Uaid) -> String {
121        format!("autopush/user/{}", uaid)
122    }
123
124    /// This store the last connection record, but doesn't update User
125    fn last_co_key(&self, uaid: &Uaid) -> String {
126        format!("autopush/co/{}", uaid)
127    }
128
129    /// This store the last timestamp incremented by the server once messages are ACK'ed
130    fn storage_timestamp_key(&self, uaid: &Uaid) -> String {
131        format!("autopush/timestamp/{}", uaid)
132    }
133
134    fn channel_list_key(&self, uaid: &Uaid) -> String {
135        format!("autopush/channels/{}", uaid)
136    }
137
138    fn message_list_key(&self, uaid: &Uaid) -> String {
139        format!("autopush/msgs/{}", uaid)
140    }
141
142    fn message_exp_list_key(&self, uaid: &Uaid) -> String {
143        format!("autopush/msgs_exp/{}", uaid)
144    }
145
146    fn message_key(&self, uaid: &Uaid, chidmessageid: &str) -> String {
147        format!("autopush/msg/{}/{}", uaid, chidmessageid)
148    }
149
150    #[cfg(feature = "reliable_report")]
151    fn reliability_key(
152        &self,
153        reliability_id: &str,
154        state: &crate::reliability::ReliabilityState,
155    ) -> String {
156        format!("autopush/reliability/{}/{}", reliability_id, state)
157    }
158
159    #[cfg(test)]
160    /// Return a single "raw" message (used by testing and validation)
161    async fn fetch_message(&self, uaid: &Uuid, chidmessageid: &str) -> DbResult<Option<String>> {
162        let message_key = self.message_key(&Uaid(uaid), chidmessageid);
163        let mut con = self.connection().await?;
164        debug!("🐰 Fetching message from {}", &message_key);
165        let message = con.get::<String, Option<String>>(message_key).await?;
166        Ok(message)
167    }
168}
169
170#[async_trait]
171impl DbClient for RedisClientImpl {
172    /// add user to the database
173    async fn add_user(&self, user: &User) -> DbResult<()> {
174        let uaid = Uaid(&user.uaid);
175        let user_key = self.user_key(&uaid);
176        let mut con = self.connection().await?;
177        let co_key = self.last_co_key(&uaid);
178        trace!("🐰 Adding user {} at {}:{}", &user.uaid, &user_key, &co_key);
179        trace!("🐰 Logged at {}", &user.connected_at);
180        pipe()
181            .set_options(co_key, ms_since_epoch(), self.router_opts.clone())
182            .set_options(
183                user_key,
184                serde_json::to_string(user)?,
185                self.router_opts.clone(),
186            )
187            .exec_async(&mut con)
188            .await?;
189        Ok(())
190    }
191
192    /// To update the TTL of the Redis entry we just have to SET again, with the new expiry
193    ///
194    /// NOTE: This function is called by mobile during the daily
195    /// [autoendpoint::routes::update_token_route] handling, and by desktop
196    /// [autoconnect-ws-sm::get_or_create_user]` which is called
197    /// during the `HELLO` handler. This should be enough to ensure that the ROUTER records
198    /// are properly refreshed for "lively" clients.
199    ///
200    /// NOTE: There is some, very small, potential risk that a desktop client that can
201    /// somehow remain connected the duration of MAX_ROUTER_TTL, may be dropped as not being
202    /// "lively".
203    async fn update_user(&self, user: &mut User) -> DbResult<bool> {
204        trace!("🐰 Updating user");
205        let mut con = self.connection().await?;
206        let co_key = self.last_co_key(&Uaid(&user.uaid));
207        let last_co: Option<u64> = con.get(&co_key).await?;
208        if last_co.is_some_and(|c| c < user.connected_at) {
209            trace!(
210                "🐰 Was connected at {}, now at {}",
211                last_co.unwrap(),
212                &user.connected_at
213            );
214            self.add_user(user).await?;
215            Ok(true)
216        } else {
217            Ok(false)
218        }
219    }
220
221    async fn get_user(&self, uaid: &Uuid) -> DbResult<Option<User>> {
222        let mut con = self.connection().await?;
223        let user_key = self.user_key(&Uaid(uaid));
224        let user: Option<User> = con
225            .get::<&str, Option<String>>(&user_key)
226            .await?
227            .and_then(|s| serde_json::from_str(s.as_ref()).ok());
228        if user.is_some() {
229            trace!("🐰 Found a record for {}", &uaid);
230        }
231        Ok(user)
232    }
233
234    async fn remove_user(&self, uaid: &Uuid) -> DbResult<()> {
235        let uaid = Uaid(uaid);
236        let mut con = self.connection().await?;
237        let user_key = self.user_key(&uaid);
238        let co_key = self.last_co_key(&uaid);
239        let chan_list_key = self.channel_list_key(&uaid);
240        let msg_list_key = self.message_list_key(&uaid);
241        let exp_list_key = self.message_exp_list_key(&uaid);
242        let timestamp_key = self.storage_timestamp_key(&uaid);
243        pipe()
244            .del(&user_key)
245            .del(&co_key)
246            .del(&chan_list_key)
247            .del(&msg_list_key)
248            .del(&exp_list_key)
249            .del(&timestamp_key)
250            .exec_async(&mut con)
251            .await?;
252        Ok(())
253    }
254
255    async fn add_channel(&self, uaid: &Uuid, channel_id: &Uuid) -> DbResult<()> {
256        let uaid = Uaid(uaid);
257        let mut con = self.connection().await?;
258        let co_key = self.last_co_key(&uaid);
259        let chan_list_key = self.channel_list_key(&uaid);
260
261        let _: () = pipe()
262            .rpush(chan_list_key, channel_id.as_hyphenated().to_string())
263            .set_options(co_key, ms_since_epoch(), self.router_opts.clone())
264            .exec_async(&mut con)
265            .await?;
266        Ok(())
267    }
268
269    /// Add channels in bulk (used mostly during migration)
270    async fn add_channels(&self, uaid: &Uuid, channels: HashSet<Uuid>) -> DbResult<()> {
271        let uaid = Uaid(uaid);
272        // channel_ids are stored a list within a single redis key
273        let mut con = self.connection().await?;
274        let co_key = self.last_co_key(&uaid);
275        let chan_list_key = self.channel_list_key(&uaid);
276        pipe()
277            .set_options(co_key, ms_since_epoch(), self.router_opts.clone())
278            .rpush(
279                chan_list_key,
280                channels
281                    .into_iter()
282                    .map(|c| c.as_hyphenated().to_string())
283                    .collect::<Vec<String>>(),
284            )
285            .exec_async(&mut con)
286            .await?;
287        Ok(())
288    }
289
290    async fn get_channels(&self, uaid: &Uuid) -> DbResult<HashSet<Uuid>> {
291        let uaid = Uaid(uaid);
292        let mut con = self.connection().await?;
293        let chan_list_key = self.channel_list_key(&uaid);
294        let channels: HashSet<Uuid> = con
295            .lrange::<&str, HashSet<String>>(&chan_list_key, 0, -1)
296            .await?
297            .into_iter()
298            .filter_map(|s| Uuid::from_str(&s).ok())
299            .collect();
300        trace!("🐰 Found {} channels for {}", channels.len(), &uaid);
301        Ok(channels)
302    }
303
304    /// Delete the channel. Does not delete its associated pending messages.
305    async fn remove_channel(&self, uaid: &Uuid, channel_id: &Uuid) -> DbResult<bool> {
306        let uaid = Uaid(uaid);
307        let channel_id = ChannelId(channel_id);
308        let mut con = self.connection().await?;
309        let co_key = self.last_co_key(&uaid);
310        let chan_list_key = self.channel_list_key(&uaid);
311        // Remove {channel_id} from autopush/channel/{auid}
312        trace!("🐰 Removing channel {}", channel_id);
313        let (status,): (bool,) = pipe()
314            .set_options(co_key, ms_since_epoch(), self.router_opts.clone())
315            .ignore()
316            .lrem(&chan_list_key, 1, channel_id.to_string())
317            .query_async(&mut con)
318            .await?;
319        Ok(status)
320    }
321
322    /// Remove the node_id
323    async fn remove_node_id(
324        &self,
325        uaid: &Uuid,
326        _node_id: &str,
327        _connected_at: u64,
328        _version: &Option<Uuid>,
329    ) -> DbResult<bool> {
330        if let Some(mut user) = self.get_user(uaid).await? {
331            user.node_id = None;
332            self.update_user(&mut user).await?;
333        }
334        Ok(true)
335    }
336
337    /// Write the notification to storage.
338    ///
339    /// If the message contains a topic, we remove the old message
340    async fn save_message(&self, uaid: &Uuid, message: Notification) -> DbResult<()> {
341        let uaid = Uaid(uaid);
342        let mut con = self.connection().await?;
343        let msg_list_key = self.message_list_key(&uaid);
344        let exp_list_key = self.message_exp_list_key(&uaid);
345        let msg_id = &message.chidmessageid();
346        let msg_key = self.message_key(&uaid, msg_id);
347        let storable: StorableNotification = message.into();
348
349        debug!("🐰 Saving message {} :: {:?}", &msg_key, &storable);
350        trace!(
351            "🐰 timestamp: {:?}",
352            &storable.timestamp.to_be_bytes().to_vec()
353        );
354
355        // Remember, `timestamp` is effectively the time to kill the message, not the
356        // current time.
357        let expiry = now_secs() + storable.ttl;
358        trace!("🐰 Message Expiry {}, currently:{} ", expiry, now_secs());
359
360        let mut pipe = pipe();
361
362        // If this is a topic message:
363        // zadd(msg_list_key) and zadd(exp_list_key) will replace their old entry
364        // in the sorted set if one already exists
365        // and set(msg_key, message) will override it too: nothing to do.
366        let is_topic = storable.topic.is_some();
367
368        // notification_ttl is already min(headers.ttl, MAX_NOTIFICATION_TTL)
369        // see autoendpoint/src/extractors/notification_headers.rs
370        let notif_opts = self
371            .notification_opts
372            .clone()
373            .with_expiration(SetExpiry::EXAT(expiry));
374
375        // Store notification record in autopush/msg/{uaid}/{chidmessageid}
376        // And store {chidmessageid} in autopush/msgs/{uaid}
377        debug!("🐰 Saving to {}", &msg_key);
378        pipe.set_options(msg_key, serde_json::to_string(&storable)?, notif_opts)
379            // The function [fetch_timestamp_messages] takes a timestamp in input,
380            // here we use the timestamp of the record
381            .zadd(&exp_list_key, msg_id, expiry)
382            .zadd(&msg_list_key, msg_id, sec_since_epoch());
383
384        let _: () = pipe.exec_async(&mut con).await?;
385        self.metrics
386            .incr_with_tags("notification.message.stored")
387            .with_tag("topic", &is_topic.to_string())
388            .with_tag("database", &self.name())
389            .send();
390        Ok(())
391    }
392
393    /// Save a batch of messages to the database.
394    ///
395    /// Currently just iterating through the list and saving one at a time. There's a bulk way
396    /// to save messages, but there are other considerations (e.g. mutation limits)
397    async fn save_messages(&self, uaid: &Uuid, messages: Vec<Notification>) -> DbResult<()> {
398        // A plate simple way of solving this:
399        for message in messages {
400            self.save_message(uaid, message).await?;
401        }
402        Ok(())
403    }
404
405    /// Delete expired messages
406    async fn increment_storage(&self, uaid: &Uuid, timestamp: u64) -> DbResult<()> {
407        let uaid = Uaid(uaid);
408        debug!("🐰🔥 Incrementing storage to {}", timestamp);
409        let msg_list_key = self.message_list_key(&uaid);
410        let exp_list_key = self.message_exp_list_key(&uaid);
411        let storage_timestamp_key = self.storage_timestamp_key(&uaid);
412        let mut con = self.connection().await?;
413        trace!("🐇 SEARCH: increment: {:?} - {}", &exp_list_key, timestamp);
414        let exp_id_list: Vec<String> = con.zrangebyscore(&exp_list_key, 0, timestamp).await?;
415        if !exp_id_list.is_empty() {
416            // Remember, we store just the message_ids in the exp and msg lists, but need to convert back to
417            // the full message keys for deletion.
418            let delete_msg_keys: Vec<String> = exp_id_list
419                .clone()
420                .into_iter()
421                .map(|msg_id| self.message_key(&uaid, &msg_id))
422                .collect();
423
424            trace!(
425                "🐰🔥:rem: Deleting {} : [{:?}]",
426                msg_list_key, &delete_msg_keys
427            );
428            trace!("🐰🔥:rem: Deleting {} : [{:?}]", exp_list_key, &exp_id_list);
429            pipe()
430                .set_options::<_, _>(&storage_timestamp_key, timestamp, self.router_opts.clone())
431                .del(&delete_msg_keys)
432                .zrem(&msg_list_key, &exp_id_list)
433                .zrem(&exp_list_key, &exp_id_list)
434                .exec_async(&mut con)
435                .await?;
436        } else {
437            con.set_options::<_, _, ()>(
438                &storage_timestamp_key,
439                timestamp,
440                self.router_opts.clone(),
441            )
442            .await?;
443        }
444        Ok(())
445    }
446
447    /// Delete the notification from storage.
448    async fn remove_message(&self, uaid: &Uuid, chidmessageid: &str) -> DbResult<()> {
449        let uaid = Uaid(uaid);
450        trace!(
451            "🐰 attemping to delete {:?} :: {:?}",
452            uaid.to_string(),
453            chidmessageid
454        );
455        let msg_key = self.message_key(&uaid, chidmessageid);
456        let msg_list_key = self.message_list_key(&uaid);
457        let exp_list_key = self.message_exp_list_key(&uaid);
458        debug!("🐰🔥 Deleting message {}", &msg_key);
459        let mut con = self.connection().await?;
460        // We remove the id from the exp list at the end, to be sure
461        // it can't be removed from the list before the message is removed
462        trace!(
463            "🐰🔥:remsg: Deleting {} : {:?}",
464            msg_list_key, &chidmessageid
465        );
466        trace!(
467            "🐰🔥:remsg: Deleting {} : {:?}",
468            exp_list_key, &chidmessageid
469        );
470        pipe()
471            .del(&msg_key)
472            .zrem(&msg_list_key, chidmessageid)
473            .zrem(&exp_list_key, chidmessageid)
474            .exec_async(&mut con)
475            .await?;
476        self.metrics
477            .incr_with_tags("notification.message.deleted")
478            .with_tag("database", &self.name())
479            .send();
480        Ok(())
481    }
482
483    /// Topic messages are handled as other messages with redis, we return nothing.
484    async fn fetch_topic_messages(
485        &self,
486        _uaid: &Uuid,
487        _limit: usize,
488    ) -> DbResult<FetchMessageResponse> {
489        Ok(FetchMessageResponse {
490            messages: vec![],
491            timestamp: None,
492        })
493    }
494
495    /// Return [`limit`] messages pending for a [`uaid`] that have a record timestamp
496    /// after [`timestamp`] (secs).
497    ///
498    /// If [`limit`] = 0, we fetch all messages after [`timestamp`].
499    ///
500    /// This can return expired messages, following bigtables behavior
501    async fn fetch_timestamp_messages(
502        &self,
503        uaid: &Uuid,
504        timestamp: Option<u64>,
505        limit: usize,
506    ) -> DbResult<FetchMessageResponse> {
507        let uaid = Uaid(uaid);
508        trace!("🐰 Fetching {} messages since {:?}", limit, timestamp);
509        let mut con = self.connection().await?;
510        let msg_list_key = self.message_list_key(&uaid);
511        let timestamp = if let Some(timestamp) = timestamp {
512            timestamp
513        } else {
514            let storage_timestamp_key = self.storage_timestamp_key(&uaid);
515            con.get(&storage_timestamp_key).await.unwrap_or(0)
516        };
517        // ZRANGE Key (x) +inf LIMIT 0 limit
518        trace!(
519            "🐇 SEARCH: zrangebyscore {:?} {} +inf withscores limit 0 {:?}",
520            &msg_list_key, timestamp, limit,
521        );
522        let results = con
523            .zrangebyscore_limit_withscores::<&str, &str, &str, Vec<(String, u64)>>(
524                &msg_list_key,
525                &timestamp.to_string(),
526                "+inf",
527                0,
528                limit as isize,
529            )
530            .await?;
531        let (messages_id, mut scores): (Vec<String>, Vec<u64>) = results
532            .into_iter()
533            .map(|(id, s): (String, u64)| (self.message_key(&uaid, &id), s))
534            .unzip();
535        if messages_id.is_empty() {
536            trace!("🐰 No message found");
537            return Ok(FetchMessageResponse {
538                messages: vec![],
539                timestamp: None,
540            });
541        }
542        let messages: Vec<Notification> = con
543            .mget::<&Vec<String>, Vec<Option<String>>>(&messages_id)
544            .await?
545            .into_iter()
546            .filter_map(|opt: Option<String>| {
547                if let Some(m) = opt {
548                    serde_json::from_str(&m)
549                        .inspect_err(|e| {
550                            // Since we can't raise the error here, at least record it
551                            // so that it's not lost.
552                            // Mind you, if there is an error here, it's probably due to
553                            // some developmental issue since the unit and integration tests
554                            // should fail.
555                            error!("🐰 ERROR parsing entry: {:?}", e);
556                        })
557                        .ok()
558                } else {
559                    None
560                }
561            })
562            .collect();
563        if messages.is_empty() {
564            trace!("🐰 No Valid messages found");
565            return Ok(FetchMessageResponse {
566                timestamp: None,
567                messages: vec![],
568            });
569        }
570        let timestamp = scores.pop();
571        trace!("🐰 Found {} messages until {:?}", messages.len(), timestamp);
572        Ok(FetchMessageResponse {
573            messages,
574            timestamp,
575        })
576    }
577
578    #[cfg(feature = "reliable_report")]
579    async fn log_report(
580        &self,
581        reliability_id: &str,
582        state: crate::reliability::ReliabilityState,
583    ) -> DbResult<()> {
584        use crate::MAX_NOTIFICATION_TTL_SECS;
585
586        trace!("🐰 Logging reliability report");
587        let mut con = self.connection().await?;
588        // TODO: Should this be a hash key per reliability_id?
589        let reliability_key = self.reliability_key(reliability_id, &state);
590        // Reports should last about as long as the notifications they're tied to.
591        let expiry = MAX_NOTIFICATION_TTL_SECS;
592        let opts = SetOptions::default().with_expiration(SetExpiry::EX(expiry));
593        let mut pipe = pipe();
594        pipe.set_options(reliability_key, sec_since_epoch(), opts)
595            .exec_async(&mut con)
596            .await?;
597        Ok(())
598    }
599
600    async fn health_check(&self) -> DbResult<bool> {
601        let _: () = self.connection().await?.ping().await?;
602        Ok(true)
603    }
604
605    /// Returns true, because there's no table in Redis
606    async fn router_table_exists(&self) -> DbResult<bool> {
607        Ok(true)
608    }
609
610    /// Returns true, because there's no table in Redis
611    async fn message_table_exists(&self) -> DbResult<bool> {
612        Ok(true)
613    }
614
615    fn box_clone(&self) -> Box<dyn DbClient> {
616        Box::new(self.clone())
617    }
618
619    fn name(&self) -> String {
620        "Redis".to_owned()
621    }
622
623    fn pool_status(&self) -> Option<deadpool::Status> {
624        None
625    }
626}
627
628#[cfg(test)]
629mod tests {
630    use crate::{logging::init_test_logging, util::ms_since_epoch};
631    use rand::prelude::*;
632    use std::env;
633
634    use super::*;
635    const TEST_CHID: &str = "DECAFBAD-0000-0000-0000-0123456789AB";
636    const TOPIC_CHID: &str = "DECAFBAD-1111-0000-0000-0123456789AB";
637
638    fn new_client() -> DbResult<RedisClientImpl> {
639        // Use an environment variable to potentially override the default redis test host.
640        let host = env::var("REDIS_HOST").unwrap_or("localhost".into());
641        let env_dsn = format!("redis://{host}");
642        debug!("🐰 Connecting to {env_dsn}");
643        let settings = DbSettings {
644            dsn: Some(env_dsn),
645            db_settings: "".into(),
646        };
647        let metrics = Arc::new(StatsdClient::builder("", cadence::NopMetricSink).build());
648        RedisClientImpl::new(metrics, &settings)
649    }
650
651    fn gen_test_user() -> String {
652        // Create a semi-unique test user to avoid conflicting test values.
653        let mut rng = rand::rng();
654        let test_num = rng.random::<u8>();
655        format!("DEADBEEF-0000-0000-{:04}-{:012}", test_num, now_secs())
656    }
657
658    #[actix_rt::test]
659    async fn health_check() {
660        let client = new_client().unwrap();
661
662        let result = client.health_check().await;
663        assert!(result.is_ok());
664        assert!(result.unwrap());
665    }
666
667    /// Test if [increment_storage] correctly wipe expired messages
668    #[actix_rt::test]
669    async fn wipe_expired() -> DbResult<()> {
670        init_test_logging();
671        let client = new_client()?;
672
673        let connected_at = ms_since_epoch();
674
675        let uaid = Uuid::parse_str(&gen_test_user()).unwrap();
676        let chid = Uuid::parse_str(TEST_CHID).unwrap();
677
678        let node_id = "test_node".to_owned();
679
680        // purge the user record if it exists.
681        let _ = client.remove_user(&uaid).await;
682
683        let test_user = User {
684            uaid,
685            router_type: "webpush".to_owned(),
686            connected_at,
687            router_data: None,
688            node_id: Some(node_id.clone()),
689            ..Default::default()
690        };
691
692        // purge the old user (if present)
693        // in case a prior test failed for whatever reason.
694        let _ = client.remove_user(&uaid).await;
695
696        // can we add the user?
697        let timestamp = now_secs();
698        client.add_user(&test_user).await?;
699        let test_notification = crate::db::Notification {
700            channel_id: chid,
701            version: "test".to_owned(),
702            ttl: 1,
703            timestamp,
704            data: Some("Encrypted".into()),
705            sortkey_timestamp: Some(timestamp),
706            ..Default::default()
707        };
708        client.save_message(&uaid, test_notification).await?;
709        client.increment_storage(&uaid, timestamp + 1).await?;
710        let msgs = client.fetch_timestamp_messages(&uaid, None, 999).await?;
711        assert_eq!(msgs.messages.len(), 0);
712        assert!(client.remove_user(&uaid).await.is_ok());
713        Ok(())
714    }
715
716    /// run a gauntlet of testing. These are a bit linear because they need
717    /// to run in sequence.
718    #[actix_rt::test]
719    async fn run_gauntlet() -> DbResult<()> {
720        init_test_logging();
721        let client = new_client()?;
722
723        let connected_at = ms_since_epoch();
724
725        let user_id = &gen_test_user();
726        let uaid = Uuid::parse_str(user_id).unwrap();
727        let chid = Uuid::parse_str(TEST_CHID).unwrap();
728        let topic_chid = Uuid::parse_str(TOPIC_CHID).unwrap();
729
730        let node_id = "test_node".to_owned();
731
732        // purge the user record if it exists.
733        let _ = client.remove_user(&uaid).await;
734
735        let test_user = User {
736            uaid,
737            router_type: "webpush".to_owned(),
738            connected_at,
739            router_data: None,
740            node_id: Some(node_id.clone()),
741            ..Default::default()
742        };
743
744        // purge the old user (if present)
745        // in case a prior test failed for whatever reason.
746        let _ = client.remove_user(&uaid).await;
747
748        // can we add the user?
749        client.add_user(&test_user).await?;
750        let fetched = client.get_user(&uaid).await?;
751        assert!(fetched.is_some());
752        let fetched = fetched.unwrap();
753        assert_eq!(fetched.router_type, "webpush".to_owned());
754
755        // Simulate a connected_at occuring before the following writes
756        let connected_at = ms_since_epoch();
757
758        // can we add channels?
759        client.add_channel(&uaid, &chid).await?;
760        let channels = client.get_channels(&uaid).await?;
761        assert!(channels.contains(&chid));
762
763        // can we add lots of channels?
764        let mut new_channels: HashSet<Uuid> = HashSet::new();
765        new_channels.insert(chid);
766        for _ in 1..10 {
767            new_channels.insert(uuid::Uuid::new_v4());
768        }
769        let chid_to_remove = uuid::Uuid::new_v4();
770        new_channels.insert(chid_to_remove);
771        client.add_channels(&uaid, new_channels.clone()).await?;
772        let channels = client.get_channels(&uaid).await?;
773        assert_eq!(channels, new_channels);
774
775        // can we remove a channel?
776        assert!(client.remove_channel(&uaid, &chid_to_remove).await?);
777        assert!(!client.remove_channel(&uaid, &chid_to_remove).await?);
778        new_channels.remove(&chid_to_remove);
779        let channels = client.get_channels(&uaid).await?;
780        assert_eq!(channels, new_channels);
781
782        // now ensure that we can update a user that's after the time we set
783        // prior. first ensure that we can't update a user that's before the
784        // time we set prior to the last write
785        let mut updated = User {
786            connected_at,
787            ..test_user.clone()
788        };
789        let result = client.update_user(&mut updated).await;
790        assert!(result.is_ok());
791        assert!(!result.unwrap());
792
793        // Make sure that the `connected_at` wasn't modified
794        let fetched2 = client.get_user(&fetched.uaid).await?.unwrap();
795        assert_eq!(fetched.connected_at, fetched2.connected_at);
796
797        // and make sure we can update a record with a later connected_at time.
798        let mut updated = User {
799            connected_at: fetched.connected_at + 300,
800            ..fetched2
801        };
802        let result = client.update_user(&mut updated).await;
803        assert!(result.is_ok());
804        assert!(result.unwrap());
805        assert_ne!(
806            fetched2.connected_at,
807            client.get_user(&uaid).await?.unwrap().connected_at
808        );
809
810        // can we increment the storage for the user?
811        client
812            .increment_storage(
813                &fetched.uaid,
814                SystemTime::now()
815                    .duration_since(SystemTime::UNIX_EPOCH)
816                    .unwrap()
817                    .as_secs(),
818            )
819            .await?;
820
821        let test_data = "An_encrypted_pile_of_crap".to_owned();
822        let timestamp = now_secs();
823        let sort_key = now_secs();
824        let fetch_timestamp = timestamp;
825        // Can we store a message?
826        let test_notification = crate::db::Notification {
827            channel_id: chid,
828            version: "test".to_owned(),
829            ttl: 300,
830            timestamp,
831            data: Some(test_data.clone()),
832            sortkey_timestamp: Some(sort_key),
833            ..Default::default()
834        };
835        let res = client.save_message(&uaid, test_notification.clone()).await;
836        assert!(res.is_ok());
837
838        let mut fetched = client.fetch_timestamp_messages(&uaid, None, 999).await?;
839        assert_ne!(fetched.messages.len(), 0);
840        let fm = fetched.messages.pop().unwrap();
841        assert_eq!(fm.channel_id, test_notification.channel_id);
842        assert_eq!(fm.data, Some(test_data));
843
844        // Grab all 1 of the messages that were submitted within the past 10 seconds.
845        let fetched = client
846            .fetch_timestamp_messages(&uaid, Some(fetch_timestamp - 10), 999)
847            .await?;
848        assert_ne!(fetched.messages.len(), 0);
849
850        // Try grabbing a message for 10 seconds from now.
851        let fetched = client
852            .fetch_timestamp_messages(&uaid, Some(fetch_timestamp + 10), 999)
853            .await?;
854        assert_eq!(fetched.messages.len(), 0);
855
856        // can we clean up our toys?
857        assert!(
858            client
859                .remove_message(&uaid, &test_notification.chidmessageid())
860                .await
861                .is_ok()
862        );
863
864        assert!(client.remove_channel(&uaid, &chid).await.is_ok());
865
866        let msgs = client
867            .fetch_timestamp_messages(&uaid, None, 999)
868            .await?
869            .messages;
870        assert!(msgs.is_empty());
871
872        // Now, can we do all that with topic messages
873        // Unlike bigtable, we don't use [fetch_topic_messages]: it always return None:
874        // they are handled as usuals messages.
875        client.add_channel(&uaid, &topic_chid).await?;
876        let test_data = "An_encrypted_pile_of_crap_with_a_topic".to_owned();
877        let timestamp = now_secs();
878        let sort_key = now_secs();
879
880        // We store 2 messages, with a single topic
881        let test_notification_0 = crate::db::Notification {
882            channel_id: topic_chid,
883            version: "version0".to_owned(),
884            ttl: 300,
885            topic: Some("topic".to_owned()),
886            timestamp,
887            data: Some(test_data.clone()),
888            sortkey_timestamp: Some(sort_key),
889            ..Default::default()
890        };
891        assert!(
892            client
893                .save_message(&uaid, test_notification_0.clone())
894                .await
895                .is_ok()
896        );
897
898        let test_notification = crate::db::Notification {
899            timestamp: now_secs(),
900            version: "version1".to_owned(),
901            sortkey_timestamp: Some(sort_key + 10),
902            ..test_notification_0
903        };
904
905        assert!(
906            client
907                .save_message(&uaid, test_notification.clone())
908                .await
909                .is_ok()
910        );
911
912        let mut fetched = client.fetch_timestamp_messages(&uaid, None, 999).await?;
913        assert_eq!(fetched.messages.len(), 1);
914        let fm = fetched.messages.pop().unwrap();
915        assert_eq!(fm.channel_id, test_notification.channel_id);
916        assert_eq!(fm.data, Some(test_data));
917
918        // Grab the message that was submitted.
919        let fetched = client.fetch_timestamp_messages(&uaid, None, 999).await?;
920        assert_ne!(fetched.messages.len(), 0);
921
922        // can we clean up our toys?
923        assert!(
924            client
925                .remove_message(&uaid, &test_notification.chidmessageid())
926                .await
927                .is_ok()
928        );
929
930        assert!(client.remove_channel(&uaid, &topic_chid).await.is_ok());
931
932        let msgs = client
933            .fetch_timestamp_messages(&uaid, None, 999)
934            .await?
935            .messages;
936        assert!(msgs.is_empty());
937
938        let fetched = client.get_user(&uaid).await?.unwrap();
939        assert!(
940            client
941                .remove_node_id(&uaid, &node_id, connected_at, &fetched.version)
942                .await
943                .is_ok()
944        );
945        // did we remove it?
946        let fetched = client.get_user(&uaid).await?.unwrap();
947        assert_eq!(fetched.node_id, None);
948
949        assert!(client.remove_user(&uaid).await.is_ok());
950
951        assert!(client.get_user(&uaid).await?.is_none());
952        Ok(())
953    }
954
955    #[actix_rt::test]
956    async fn test_expiry() -> DbResult<()> {
957        // Make sure that we really are purging messages correctly
958        init_test_logging();
959        let client = new_client()?;
960
961        let uaid = Uuid::parse_str(&gen_test_user()).unwrap();
962        let chid = Uuid::parse_str(TEST_CHID).unwrap();
963        let now = now_secs();
964
965        let test_notification = crate::db::Notification {
966            channel_id: chid,
967            version: "test".to_owned(),
968            ttl: 2,
969            timestamp: now,
970            data: Some("SomeData".into()),
971            sortkey_timestamp: Some(now),
972            ..Default::default()
973        };
974        debug!("Writing test notif");
975        let res = client.save_message(&uaid, test_notification.clone()).await;
976        assert!(res.is_ok());
977        let key = client.message_key(&Uaid(&uaid), &test_notification.chidmessageid());
978        debug!("Checking {}...", &key);
979        let msg = client
980            .fetch_message(&uaid, &test_notification.chidmessageid())
981            .await?;
982        assert!(!msg.unwrap().is_empty());
983        debug!("Purging...");
984        client.increment_storage(&uaid, now + 2).await?;
985        debug!("Checking {}...", &key);
986        let cc = client
987            .fetch_message(&uaid, &test_notification.chidmessageid())
988            .await;
989        assert_eq!(cc.unwrap(), None);
990        // clean up after the test.
991        assert!(client.remove_user(&uaid).await.is_ok());
992        Ok(())
993    }
994}
autopush_common/db/redis/redis_client/mod.rs

autopush_common/db/redis/redis_client/
mod.rs
