Skip to main content

autopush_common/db/bigtable/bigtable_client/
mod.rs

1use std::borrow::Cow;
2use std::collections::{BTreeMap, HashMap, HashSet};
3use std::error::Error as StdError;
4use std::fmt;
5use std::fmt::Display;
6use std::future::Future;
7use std::str::FromStr;
8use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
9use std::sync::{Arc, Mutex};
10use std::time::{Duration, SystemTime};
11
12use again::RetryPolicy;
13use async_trait::async_trait;
14use cadence::{Counted, StatsdClient};
15#[cfg(feature = "reliable_report")]
16use chrono::TimeDelta;
17use gcp_auth::TokenProvider;
18use googleapis_tonic_google_bigtable_v2::google::bigtable::v2 as bigtable;
19use googleapis_tonic_google_bigtable_v2::google::bigtable::v2::bigtable_client::BigtableClient;
20use serde_json::{from_str, json};
21use tonic::metadata::{AsciiMetadataValue, MetadataMap};
22use tonic::transport::Channel;
23use tonic::{Code, Request, Status};
24use uuid::Uuid;
25
26use crate::MAX_ROUTER_TTL_SECS;
27use crate::db::{
28    DbSettings, Notification, USER_RECORD_VERSION, User,
29    client::{DbClient, FetchMessageResponse},
30    error::{DbError, DbResult},
31    models::RangeKey,
32};
33use crate::metric_name::MetricName;
34use crate::metrics::StatsdClientExt;
35
36pub use self::metadata::MetadataBuilder;
37use self::row::{Row, RowCells};
38use super::BigTableDbSettings;
39use super::pool::BigTablePool;
40
41pub mod cell;
42pub mod error;
43pub(crate) mod merge;
44pub mod metadata;
45pub mod row;
46
47// these are normally Vec<u8>
48pub type RowKey = String;
49
50// These are more for code clarity than functional types.
51// Rust will happily swap between the two in any case.
52// See [super::row::Row] for discussion about how these
53// are overloaded in order to simplify fetching data.
54pub type Qualifier = String;
55pub type FamilyId = String;
56
57const ROUTER_FAMILY: &str = "router";
58const MESSAGE_FAMILY: &str = "message"; // The default family for messages
59const MESSAGE_TOPIC_FAMILY: &str = "message_topic";
60#[cfg(feature = "reliable_report")]
61const RELIABLE_LOG_FAMILY: &str = "reliability";
62#[cfg(feature = "reliable_report")]
63/// The maximum TTL for reliability logging (60 days).
64/// /// In most use cases, converted to seconds through .num_seconds().
65pub const RELIABLE_LOG_TTL: TimeDelta = TimeDelta::days(60);
66
67/// Default number of retries after the initial Bigtable RPC attempt.
68///
69/// So a connectivity failure cannot fan out into a
70/// retry storm, we try to keep this number small.
71/// Can be overridden through `db_settings.retry_count`; health checks use the
72/// same configured value as other point reads.
73pub(crate) const RETRY_COUNT: usize = 2;
74
75/// Maximum gRPC message size (256MB), matching the prior grpcio configuration.
76const MAX_MESSAGE_LEN: usize = 1 << 28;
77
78/// OAuth2 scopes requested for the Bigtable data API.
79const BIGTABLE_DATA_SCOPES: &[&str] = &["https://www.googleapis.com/auth/bigtable.data"];
80
81/// Simple circuit breaker to prevent retry storms during BigTable outages.
82///
83/// After `failure_threshold` consecutive failures, the circuit opens and
84/// requests fail fast for `cooldown_secs` seconds before allowing a retry.
85#[derive(Debug)]
86pub struct CircuitBreaker {
87    consecutive_failures: AtomicU32,
88    opened_at_epoch_secs: AtomicU64,
89    failure_threshold: u32,
90    cooldown_secs: u64,
91}
92
93impl CircuitBreaker {
94    pub fn new(failure_threshold: u32, cooldown_secs: u64) -> Self {
95        Self {
96            consecutive_failures: AtomicU32::new(0),
97            opened_at_epoch_secs: AtomicU64::new(0),
98            failure_threshold,
99            cooldown_secs,
100        }
101    }
102
103    /// Check if the circuit is allowing requests through.
104    /// Returns true if the request should proceed, false if it should fail fast.
105    pub fn allow_request(&self) -> bool {
106        let failures = self.consecutive_failures.load(Ordering::Relaxed);
107        if failures < self.failure_threshold {
108            return true;
109        }
110        // Circuit is open — check if cooldown has elapsed
111        let opened_at = self.opened_at_epoch_secs.load(Ordering::Relaxed);
112        let now = SystemTime::now()
113            .duration_since(SystemTime::UNIX_EPOCH)
114            .unwrap_or_default()
115            .as_secs();
116        if now.saturating_sub(opened_at) >= self.cooldown_secs {
117            // Allow a single probe request (half-open state)
118            true
119        } else {
120            false
121        }
122    }
123
124    /// Record a successful operation, resetting the circuit breaker.
125    pub fn record_success(&self) {
126        self.consecutive_failures.store(0, Ordering::Relaxed);
127    }
128
129    /// Record a failed operation.
130    pub fn record_failure(&self) {
131        let prev = self.consecutive_failures.fetch_add(1, Ordering::Relaxed);
132        if prev + 1 >= self.failure_threshold {
133            let now = SystemTime::now()
134                .duration_since(SystemTime::UNIX_EPOCH)
135                .unwrap_or_default()
136                .as_secs();
137            self.opened_at_epoch_secs.store(now, Ordering::Relaxed);
138        }
139    }
140}
141
142impl Default for CircuitBreaker {
143    fn default() -> Self {
144        // Open after 5 consecutive failures, cooldown for 30 seconds
145        Self::new(5, 30)
146    }
147}
148
149/// Semi convenience wrapper to ensure that the UAID is formatted and displayed consistently.
150// TODO:Should we create something similar for ChannelID?
151struct Uaid(Uuid);
152
153impl Display for Uaid {
154    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
155        write!(f, "{}", self.0.as_simple())
156    }
157}
158
159impl From<Uaid> for String {
160    fn from(uaid: Uaid) -> String {
161        uaid.0.as_simple().to_string()
162    }
163}
164
165#[derive(Clone)]
166/// Bigtable-backed implementation of the application's database interface.
167pub struct BigTableClientImpl {
168    pub(crate) settings: BigTableDbSettings,
169    /// Metrics client
170    metrics: Arc<StatsdClient>,
171    /// Logical-operation pool and shared tonic channel set.
172    pool: BigTablePool,
173    metadata: MetadataMap,
174    /// Circuit breaker to prevent retry storms during BigTable outages
175    circuit_breaker: Arc<CircuitBreaker>,
176}
177
178/// Return a RowFilter matching the `maxversions=1` portion of a family's GC
179/// policy (keep only the most recent cell per column). Shared by the router and
180/// message family filters below.
181fn max_versions_filter() -> bigtable::RowFilter {
182    bigtable::RowFilter {
183        filter: Some(bigtable::row_filter::Filter::CellsPerColumnLimitFilter(1)),
184    }
185}
186
187/// Return a RowFilter that excludes already-expired cells.
188///
189/// Bigtable's server-side garbage collection can lag the expiry time by days,
190/// so this filters out any cells that have expired in the past but haven't yet
191/// been garbage collected.
192fn expiry_filter() -> Result<bigtable::RowFilter, error::BigTableError> {
193    let bt_now: i64 = SystemTime::now()
194        .duration_since(SystemTime::UNIX_EPOCH)
195        .map_err(error::BigTableError::WriteTime)?
196        .as_millis() as i64;
197    Ok(bigtable::RowFilter {
198        filter: Some(bigtable::row_filter::Filter::TimestampRangeFilter(
199            bigtable::TimestampRange {
200                start_timestamp_micros: bt_now * 1000,
201                end_timestamp_micros: 0,
202            },
203        )),
204    })
205}
206
207/// Return a RowFilter matching the GC policy of the router Column Family.
208///
209/// Deliberately *not* chained with [expiry_filter], unlike
210/// [message_gc_policy_filter]. The router family currently has no server-side
211/// `max_age` policy. Ahead of applying server-side GC, [DbClient::get_user]
212/// counts what would have been expired/hidden, so the population can be sized
213/// first.
214fn router_gc_policy_filter() -> bigtable::RowFilter {
215    max_versions_filter()
216}
217
218/// Return a chain of RowFilters matching the GC policy of the message Column
219/// Families
220fn message_gc_policy_filter() -> Result<Vec<bigtable::RowFilter>, error::BigTableError> {
221    Ok(vec![max_versions_filter(), expiry_filter()?])
222}
223
224/// Return a Column family regex RowFilter
225fn family_filter(regex: String) -> bigtable::RowFilter {
226    bigtable::RowFilter {
227        filter: Some(bigtable::row_filter::Filter::FamilyNameRegexFilter(regex)),
228    }
229}
230
231/// Escape bytes for RE values
232///
233/// Based off google-re2/perl's quotemeta function
234fn escape_bytes(bytes: &[u8]) -> Vec<u8> {
235    let mut vec = Vec::with_capacity(bytes.len() * 2);
236    for &b in bytes {
237        if !b.is_ascii_alphanumeric() && b != b'_' && (b & 128) == 0 {
238            if b == b'\0' {
239                // Special handling for null: Note that this special handling
240                // is not strictly required for RE2, but this quoting is
241                // required for other regexp libraries such as PCRE.
242                // Can't use "\\0" since the next character might be a digit.
243                vec.extend("\\x00".as_bytes());
244                continue;
245            }
246            vec.push(b'\\');
247        }
248        vec.push(b);
249    }
250    vec
251}
252
253/// Return a chain of RowFilters limiting to a match of the specified
254/// `version`'s column value
255fn version_filter(version: &Uuid) -> Vec<bigtable::RowFilter> {
256    let cq_filter = bigtable::RowFilter {
257        filter: Some(bigtable::row_filter::Filter::ColumnQualifierRegexFilter(
258            "^version$".as_bytes().to_vec(),
259        )),
260    };
261    let value_filter = bigtable::RowFilter {
262        filter: Some(bigtable::row_filter::Filter::ValueRegexFilter(
263            escape_bytes(version.as_bytes()),
264        )),
265    };
266
267    vec![
268        family_filter(format!("^{ROUTER_FAMILY}$")),
269        cq_filter,
270        value_filter,
271    ]
272}
273
274/// Return a newly generated `version` column `Cell`
275fn new_version_cell(timestamp: SystemTime) -> cell::Cell {
276    cell::Cell {
277        qualifier: "version".to_owned(),
278        value: Uuid::new_v4().into(),
279        timestamp,
280        ..Default::default()
281    }
282}
283
284/// Return a RowFilter chain from multiple RowFilters
285fn filter_chain(filters: Vec<bigtable::RowFilter>) -> bigtable::RowFilter {
286    bigtable::RowFilter {
287        filter: Some(bigtable::row_filter::Filter::Chain(
288            bigtable::row_filter::Chain { filters },
289        )),
290    }
291}
292
293/// Return a ReadRowsRequest against table for a given row key
294fn read_row_request(
295    table_name: &str,
296    app_profile_id: &str,
297    row_key: &str,
298) -> bigtable::ReadRowsRequest {
299    bigtable::ReadRowsRequest {
300        table_name: table_name.to_owned(),
301        app_profile_id: app_profile_id.to_owned(),
302        rows: Some(bigtable::RowSet {
303            row_keys: vec![row_key.as_bytes().to_vec()],
304            row_ranges: Vec::new(),
305        }),
306        ..Default::default()
307    }
308}
309
310fn to_u64(value: Vec<u8>, name: &str) -> Result<u64, DbError> {
311    let v: [u8; 8] = value
312        .try_into()
313        .map_err(|_| DbError::DeserializeU64(name.to_owned()))?;
314    Ok(u64::from_be_bytes(v))
315}
316
317fn to_string(value: Vec<u8>, name: &str) -> Result<String, DbError> {
318    String::from_utf8(value).map_err(|e| {
319        debug!("🉑 cannot read string {}: {:?}", name, e);
320        DbError::DeserializeString(name.to_owned())
321    })
322}
323
324/// Parse the "set" (see [DbClient::add_channels]) of channel ids in a bigtable Row.
325///
326/// Cells should solely contain the set of channels otherwise an Error is returned.
327fn channels_from_cells(cells: &RowCells) -> DbResult<HashSet<Uuid>> {
328    let mut result = HashSet::new();
329    for cells in cells.values() {
330        let Some(cell) = cells.last() else {
331            continue;
332        };
333        let Some((_, chid)) = cell.qualifier.split_once("chid:") else {
334            return Err(DbError::Integrity(
335                "get_channels expected: chid:<chid>".to_owned(),
336                None,
337            ));
338        };
339        result.insert(Uuid::from_str(chid).map_err(|e| DbError::General(e.to_string()))?);
340    }
341    Ok(result)
342}
343
344/// Convert the [HashSet] of channel ids to cell entries for a bigtable Row
345fn channels_to_cells(channels: Cow<HashSet<Uuid>>, expiry: SystemTime) -> Vec<cell::Cell> {
346    let channels = channels.into_owned();
347    let mut cells = Vec::with_capacity(channels.len().min(100_000));
348    for (i, channel_id) in channels.into_iter().enumerate() {
349        // There is a limit of 100,000 mutations per batch for bigtable.
350        // https://cloud.google.com/bigtable/quotas
351        // If you have 100,000 channels, you have too many.
352        if i >= 100_000 {
353            break;
354        }
355        cells.push(cell::Cell {
356            qualifier: format!("chid:{}", channel_id.as_hyphenated()),
357            timestamp: expiry,
358            ..Default::default()
359        });
360    }
361    cells
362}
363
364pub fn retry_policy(max: usize) -> RetryPolicy {
365    RetryPolicy::default()
366        .with_max_retries(max)
367        .with_jitter(true)
368}
369
370#[derive(Clone, Copy, Debug, Eq, PartialEq)]
371enum RetryKind {
372    Read,
373    IdempotentWrite,
374    ConditionalWrite,
375}
376
377#[derive(Clone, Copy, Debug)]
378enum RpcClass {
379    Point,
380    Scan,
381}
382
383#[derive(Clone, Copy, Debug, Eq, PartialEq)]
384enum BreakerPolicy {
385    Track,
386    Ignore,
387}
388
389#[derive(Clone, Copy, Debug)]
390struct RpcPolicy {
391    attempt_timeout: Duration,
392    total_timeout: Duration,
393    retry_kind: RetryKind,
394    breaker: BreakerPolicy,
395}
396
397async fn rpc_attempt_until<T>(
398    deadline: tokio::time::Instant,
399    future: impl Future<Output = Result<T, error::BigTableError>>,
400) -> Result<T, error::BigTableError> {
401    tokio::time::timeout_at(deadline, future)
402        .await
403        .map_err(|_| error::BigTableError::AttemptTimeout)?
404}
405
406async fn operation_budget_until<T>(
407    deadline: tokio::time::Instant,
408    future: impl Future<Output = Result<T, error::BigTableError>>,
409) -> Result<T, error::BigTableError> {
410    tokio::time::timeout_at(deadline, future)
411        .await
412        .map_err(|_| error::BigTableError::OperationTimeout)?
413}
414
415fn retryable_pre_send_err(status: &Status) -> bool {
416    (status.code() == Code::Unknown
417        && status
418            .message()
419            .to_ascii_lowercase()
420            .starts_with("service was not ready:"))
421        || has_connect_source(status)
422}
423
424fn retryable_transient_err(status: &Status) -> bool {
425    if retryable_pre_send_err(status) {
426        return true;
427    }
428    match status.code() {
429        Code::Internal => {
430            let message = status.message().to_ascii_lowercase();
431            [
432                "rst_stream",
433                "rst stream",
434                "received unexpected eos on data frame from server",
435            ]
436            .iter()
437            .any(|fragment| message.contains(fragment))
438        }
439        Code::Unavailable | Code::DeadlineExceeded => true,
440        _ => false,
441    }
442}
443
444/// Return whether a failed read is safe to retry.
445///
446/// In addition to the statuses that are retryable for every operation, tonic
447/// can surface HTTP/2 connection retirement as either:
448///
449/// - `Internal("h2 protocol error: http2 error")` for a remote GOAWAY, or
450/// - `Cancelled("operation was canceled")` when hyper closes the connection, or
451/// - `Unknown("transport error")` when a connection-level failure (e.g. a
452///   GFE-reaped idle connection surfacing as a broken pipe) is written to.
453///
454/// Reads are idempotent, so replaying these bounded attempts is safe. Autopush
455/// `MutateRow` requests are also replay-safe: set-cell retries reuse the same
456/// explicit timestamp, delete mutations are idempotent, and every attempt
457/// clones the same request. Conditional `CheckAndMutateRow` operations use the
458/// narrower pre-send-only policy because their returned predicate result can
459/// change after an applied-but-lost attempt.
460fn retryable_read_err(status: &Status) -> bool {
461    if retryable_transient_err(status) {
462        return true;
463    }
464    // Tonic 0.14.6 uses both Internal (for some header/GOAWAY paths) and
465    // Unknown (for a connection lost while streaming a body) for this wrapper
466    // message. The exact Hyper/H2 suffix is not a stable API.
467    if matches!(status.code(), Code::Internal | Code::Unknown)
468        && status
469            .message()
470            .to_ascii_lowercase()
471            .starts_with("h2 protocol error:")
472    {
473        return true;
474    }
475    match status.code() {
476        Code::Cancelled => status
477            .message()
478            .eq_ignore_ascii_case("operation was canceled"),
479        // Tonic retains the transport and OS errors in the source chain. Use
480        // their types rather than matching an unstable display string.
481        Code::Unknown => has_transport_source(status),
482        _ => false,
483    }
484}
485
486fn has_transport_source(status: &Status) -> bool {
487    let mut source = StdError::source(status);
488    while let Some(error) = source {
489        if error.downcast_ref::<tonic::transport::Error>().is_some()
490            || error.downcast_ref::<std::io::Error>().is_some()
491        {
492            return true;
493        }
494        source = error.source();
495    }
496    false
497}
498
499fn has_connect_source(status: &Status) -> bool {
500    let mut source = StdError::source(status);
501    while let Some(error) = source {
502        if error.downcast_ref::<tonic::ConnectError>().is_some() {
503            return true;
504        }
505        source = error.source();
506    }
507    false
508}
509
510fn counts_toward_breaker(error: &error::BigTableError, operation_timed_out_in_rpc: bool) -> bool {
511    matches!(
512        error,
513        error::BigTableError::InvalidRowResponse(_)
514            | error::BigTableError::InvalidChunk(_)
515            | error::BigTableError::Read(_)
516            | error::BigTableError::Write(_)
517            | error::BigTableError::AttemptTimeout
518            | error::BigTableError::Status(_, _)
519    ) || matches!(error, error::BigTableError::OperationTimeout) && operation_timed_out_in_rpc
520}
521
522#[cfg(test)]
523mod retry_tests {
524    use std::convert::Infallible;
525    use std::sync::Mutex;
526    use std::sync::atomic::AtomicUsize;
527
528    use futures::StreamExt;
529    use http_body_util::StreamBody;
530    use hyper::body::{Bytes, Frame, Incoming};
531    use hyper::server::conn::http2;
532    use hyper::service::service_fn;
533    use hyper::{Request as HyperRequest, Response as HyperResponse};
534    use hyper_util::rt::{TokioExecutor, TokioIo};
535    use tokio::net::TcpListener;
536    use tokio::sync::oneshot;
537
538    use super::*;
539    #[derive(Debug, thiserror::Error)]
540    #[error("transport error")]
541    struct TestTransportError {
542        #[source]
543        source: std::io::Error,
544    }
545
546    fn broken_pipe_status() -> Status {
547        Status::from_error(Box::new(TestTransportError {
548            source: std::io::Error::from(std::io::ErrorKind::BrokenPipe),
549        }))
550    }
551
552    #[test]
553    fn retries_tonic_pre_send_readiness_errors() {
554        let status = Status::unknown("Service was not ready: transport error");
555        assert!(retryable_pre_send_err(&status));
556    }
557
558    #[test]
559    fn retries_source_backed_connect_errors_before_conditional_writes() {
560        let connect = tonic::ConnectError(Box::new(std::io::Error::from(
561            std::io::ErrorKind::ConnectionRefused,
562        )));
563        let status = Status::from_error(Box::new(connect));
564
565        assert_eq!(status.code(), Code::Unavailable);
566        assert!(retryable_pre_send_err(&status));
567    }
568
569    #[test]
570    fn does_not_retry_arbitrary_unknown_errors() {
571        let status = Status::unknown("operation outcome is unknown");
572        assert!(!retryable_pre_send_err(&status));
573    }
574
575    #[test]
576    fn retries_transient_bigtable_statuses() {
577        assert!(retryable_transient_err(&Status::unavailable(
578            "No zones were available"
579        )));
580        assert!(retryable_transient_err(&Status::deadline_exceeded(
581            "deadline exceeded"
582        )));
583        assert!(retryable_transient_err(&Status::internal(
584            "stream terminated by RST_STREAM before headers"
585        )));
586    }
587
588    #[test]
589    fn retries_observed_tonic_transport_failures_for_reads() {
590        let goaway = Status::internal("h2 protocol error: http2 error");
591        assert!(retryable_read_err(&goaway));
592        assert!(classify_retry(RetryKind::Read, &error::BigTableError::Read(goaway)).is_some());
593
594        let connection_closed = Status::cancelled("operation was canceled");
595        assert!(retryable_read_err(&connection_closed));
596        assert!(
597            classify_retry(
598                RetryKind::Read,
599                &error::BigTableError::Read(connection_closed)
600            )
601            .is_some()
602        );
603
604        // A GFE-reaped idle connection written to surfaces as
605        // `Unknown("transport error")` (the "broken pipe" io error is in the
606        // source, not the status message).
607        let transport_error = broken_pipe_status();
608        assert_eq!(transport_error.code(), Code::Unknown);
609        assert_eq!(transport_error.message(), "transport error");
610        assert!(retryable_read_err(&transport_error));
611        assert!(
612            classify_retry(
613                RetryKind::Read,
614                &error::BigTableError::Read(transport_error)
615            )
616            .is_some()
617        );
618    }
619
620    #[test]
621    fn does_not_retry_non_transport_unknown_reads() {
622        // `Unknown` is a catch-all that also covers server-returned application
623        // errors; require a transport or IO error in the source chain.
624        assert!(!retryable_read_err(&Status::unknown(
625            "some application error"
626        )));
627    }
628
629    #[test]
630    fn retries_idempotent_but_not_conditional_writes() {
631        let goaway =
632            error::BigTableError::Write(Status::internal("h2 protocol error: http2 error"));
633        assert!(classify_retry(RetryKind::IdempotentWrite, &goaway).is_some());
634        let conditional_goaway =
635            error::BigTableError::Write(Status::internal("h2 protocol error: http2 error"));
636        assert!(classify_retry(RetryKind::ConditionalWrite, &conditional_goaway).is_none());
637
638        let connection_closed =
639            error::BigTableError::Write(Status::cancelled("operation was canceled"));
640        assert!(classify_retry(RetryKind::IdempotentWrite, &connection_closed).is_some());
641
642        let transport_error = error::BigTableError::Write(broken_pipe_status());
643        assert!(classify_retry(RetryKind::IdempotentWrite, &transport_error).is_some());
644
645        let unavailable = error::BigTableError::Write(Status::unavailable("try again"));
646        assert!(classify_retry(RetryKind::IdempotentWrite, &unavailable).is_some());
647        assert!(classify_retry(RetryKind::ConditionalWrite, &unavailable).is_none());
648
649        let attempt_timeout = error::BigTableError::AttemptTimeout;
650        assert!(classify_retry(RetryKind::Read, &attempt_timeout).is_some());
651        assert!(classify_retry(RetryKind::ConditionalWrite, &attempt_timeout).is_none());
652    }
653
654    #[actix_rt::test]
655    async fn attempt_and_operation_timeouts_are_enforced() {
656        let attempt = rpc_attempt_until::<()>(
657            tokio::time::Instant::now() + Duration::from_millis(1),
658            async { futures::future::pending().await },
659        )
660        .await;
661        assert!(matches!(attempt, Err(error::BigTableError::AttemptTimeout)));
662
663        let operation = operation_budget_until::<()>(
664            tokio::time::Instant::now() + Duration::from_millis(1),
665            async { futures::future::pending().await },
666        )
667        .await;
668        assert!(matches!(
669            operation,
670            Err(error::BigTableError::OperationTimeout)
671        ));
672    }
673
674    #[actix_rt::test]
675    async fn retries_a_midstream_drop_on_the_next_channel() {
676        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
677        let address = listener.local_addr().unwrap();
678        let (drop_connection, wait_for_drop) = oneshot::channel::<()>();
679        let (finish_server, wait_for_finish) = oneshot::channel::<()>();
680
681        let server = tokio::spawn(async move {
682            // The first channel begins a valid streaming response and then
683            // loses its transport. This is the failure mode observed when a
684            // load balancer reaps an idle HTTP/2 connection.
685            let (socket, _) = listener.accept().await.unwrap();
686            let service = service_fn(|_request: HyperRequest<Incoming>| async move {
687                // One valid, empty ReadRowsResponse followed by a body that
688                // never completes. The client signals only after tonic has
689                // decoded this message, guaranteeing a mid-stream drop.
690                let first_message = Frame::data(Bytes::from_static(&[0, 0, 0, 0, 0]));
691                let body = StreamBody::new(
692                    futures::stream::once(
693                        async move { Ok::<Frame<Bytes>, Infallible>(first_message) },
694                    )
695                    .chain(futures::stream::pending::<Result<Frame<Bytes>, Infallible>>()),
696                );
697                Ok::<_, Infallible>(
698                    HyperResponse::builder()
699                        .status(200)
700                        .header("content-type", "application/grpc")
701                        .body(body)
702                        .unwrap(),
703                )
704            });
705            {
706                let connection = http2::Builder::new(TokioExecutor::new())
707                    .serve_connection(TokioIo::new(socket), service);
708                tokio::pin!(connection);
709                tokio::select! {
710                    result = &mut connection => {
711                        panic!("test HTTP/2 connection ended before it was dropped: {result:?}");
712                    }
713                    _ = wait_for_drop => {}
714                }
715            }
716
717            // A retry must select the second channel. Return a complete,
718            // successful gRPC stream there so success proves both retry
719            // classification and fresh-channel selection end to end.
720            let (socket, _) = listener.accept().await.unwrap();
721            let service = service_fn(|_request: HyperRequest<Incoming>| async move {
722                // A trailers-only success is the canonical empty gRPC
723                // streaming response and ends the stream immediately.
724                let body =
725                    StreamBody::new(futures::stream::empty::<Result<Frame<Bytes>, Infallible>>());
726                Ok::<_, Infallible>(
727                    HyperResponse::builder()
728                        .status(200)
729                        .header("content-type", "application/grpc")
730                        .header("grpc-status", "0")
731                        .body(body)
732                        .unwrap(),
733                )
734            });
735            let connection = http2::Builder::new(TokioExecutor::new())
736                .serve_connection(TokioIo::new(socket), service);
737            tokio::pin!(connection);
738            tokio::select! {
739                result = &mut connection => {
740                    panic!("successful HTTP/2 connection ended unexpectedly: {result:?}");
741                }
742                _ = wait_for_finish => {}
743            }
744        });
745
746        let metrics = Arc::new(StatsdClient::builder("", cadence::NopMetricSink).build());
747        let settings = DbSettings {
748            dsn: Some(format!("grpc://localhost:{}", address.port())),
749            db_settings: json!({
750                "table_name": "projects/test/instances/test/tables/test",
751                "grpc_channel_count": 2,
752                "retry_count": 1,
753                "grpc_connect_timeout": 2,
754                "grpc_point_attempt_timeout": 5,
755                "grpc_point_total_timeout": 10
756            })
757            .to_string(),
758        };
759        let client = BigTableClientImpl::new(metrics, &settings).unwrap();
760        let drop_connection = Arc::new(Mutex::new(Some(drop_connection)));
761        let policy = client.rpc_policy(RpcClass::Point, RetryKind::Read, BreakerPolicy::Track);
762        let started = tokio::time::Instant::now();
763        tokio::time::timeout(
764            Duration::from_secs(8),
765            client.execute_rpc(
766                bigtable::ReadRowsRequest::default(),
767                policy,
768                move |channel, request| {
769                    let drop_connection = drop_connection.clone();
770                    async move {
771                        let mut bigtable = BigtableDb::client(channel);
772                        let mut stream = bigtable
773                            .read_rows(request)
774                            .await
775                            .map_err(error::BigTableError::Read)?
776                            .into_inner();
777
778                        let signal = drop_connection
779                            .lock()
780                            .unwrap_or_else(|poisoned| poisoned.into_inner())
781                            .take();
782                        if let Some(signal) = signal {
783                            assert!(stream.message().await.unwrap().is_some());
784                            signal.send(()).unwrap();
785                        }
786                        while stream
787                            .message()
788                            .await
789                            .map_err(error::BigTableError::Read)?
790                            .is_some()
791                        {}
792                        Ok(())
793                    }
794                },
795            ),
796        )
797        .await
798        .expect("Bigtable retry exceeded the test budget")
799        .expect("the second channel should complete the read");
800        assert!(
801            started.elapsed() < Duration::from_secs(5),
802            "the retry waited for the attempt deadline instead of classifying the transport drop"
803        );
804
805        finish_server.send(()).unwrap();
806        server.await.unwrap();
807    }
808
809    #[test]
810    fn default_retry_budget_is_bounded() {
811        assert_eq!(RETRY_COUNT, 2);
812    }
813
814    #[actix_rt::test]
815    async fn request_preserves_metadata_and_grpc_timeout() {
816        let db = BigtableDb::new(None);
817        let mut metadata = MetadataMap::new();
818        metadata.insert("x-goog-request-params", "table_name=test".parse().unwrap());
819
820        let request = db
821            .request(
822                bigtable::ReadRowsRequest::default(),
823                &metadata,
824                tokio::time::Instant::now() + Duration::from_secs(5),
825            )
826            .await
827            .unwrap();
828
829        assert_eq!(
830            request.metadata().get("x-goog-request-params").unwrap(),
831            "table_name=test"
832        );
833        assert!(request.metadata().get("grpc-timeout").is_some());
834    }
835
836    #[actix_rt::test]
837    async fn retry_metric_is_consumed_only_by_a_followup_attempt() {
838        let pending = Arc::new(Mutex::new(None));
839        let attempt_pending = pending.clone();
840        let condition_pending = pending.clone();
841        let attempts = Arc::new(AtomicUsize::new(0));
842        let attempt_count = attempts.clone();
843        let emitted = Arc::new(AtomicUsize::new(0));
844        let emitted_count = emitted.clone();
845
846        let result = RetryPolicy::fixed(Duration::ZERO)
847            .with_max_retries(1)
848            .retry_if(
849                || {
850                    attempt_count.fetch_add(1, Ordering::Relaxed);
851                    if take_pending_retry_metric(&attempt_pending).is_some() {
852                        emitted_count.fetch_add(1, Ordering::Relaxed);
853                    }
854                    async { Err::<(), _>(error::BigTableError::AttemptTimeout) }
855                },
856                move |error: &error::BigTableError| {
857                    if let Some(retry_metric) = classify_retry(RetryKind::Read, error) {
858                        set_pending_retry_metric(&condition_pending, retry_metric);
859                        true
860                    } else {
861                        false
862                    }
863                },
864            )
865            .await;
866
867        assert!(matches!(result, Err(error::BigTableError::AttemptTimeout)));
868        assert_eq!(attempts.load(Ordering::Relaxed), 2);
869        assert_eq!(emitted.load(Ordering::Relaxed), 1);
870    }
871
872    #[test]
873    fn local_failures_do_not_count_toward_the_backend_breaker() {
874        assert!(!counts_toward_breaker(
875            &error::BigTableError::PreSendTimeout,
876            false,
877        ));
878        assert!(!counts_toward_breaker(
879            &error::BigTableError::CircuitBreakerOpen,
880            false,
881        ));
882        assert!(!counts_toward_breaker(
883            &error::BigTableError::OperationTimeout,
884            false,
885        ));
886        assert!(counts_toward_breaker(
887            &error::BigTableError::OperationTimeout,
888            true,
889        ));
890        assert!(counts_toward_breaker(
891            &error::BigTableError::Read(Status::unavailable("backend unavailable")),
892            false,
893        ));
894    }
895}
896
897pub fn metric(metrics: &Arc<StatsdClient>, err_type: &str, code: Option<&str>) {
898    let mut metric = metrics
899        .incr_with_tags(MetricName::DatabaseRetry)
900        .with_tag("error", err_type)
901        .with_tag("type", "bigtable");
902    if let Some(code) = code {
903        metric = metric.with_tag("code", code);
904    }
905    metric.send();
906}
907
908#[derive(Clone, Copy, Debug)]
909struct RetryMetric {
910    error: &'static str,
911    code: Option<Code>,
912}
913
914impl RetryMetric {
915    fn send(self, metrics: &Arc<StatsdClient>) {
916        let code = self.code.map(|code| format!("{code:?}"));
917        metric(metrics, self.error, code.as_deref());
918    }
919}
920
921fn take_pending_retry_metric(pending: &Mutex<Option<RetryMetric>>) -> Option<RetryMetric> {
922    pending
923        .lock()
924        .unwrap_or_else(|poisoned| poisoned.into_inner())
925        .take()
926}
927
928fn set_pending_retry_metric(pending: &Mutex<Option<RetryMetric>>, metric: RetryMetric) {
929    *pending
930        .lock()
931        .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(metric);
932}
933
934/// Classify a safe retry and describe the metric to emit if its next attempt
935/// actually starts.
936fn classify_retry(kind: RetryKind, err: &error::BigTableError) -> Option<RetryMetric> {
937    debug!("🉑 Checking BigTableError...{err}");
938    match err {
939        error::BigTableError::InvalidRowResponse(status)
940        | error::BigTableError::Read(status)
941        | error::BigTableError::Write(status) => {
942            let retry = match kind {
943                RetryKind::Read | RetryKind::IdempotentWrite => retryable_read_err(status),
944                RetryKind::ConditionalWrite => retryable_pre_send_err(status),
945            };
946            info!("GRPC Failure: {:?}", status);
947            retry.then_some(RetryMetric {
948                error: "RpcFailure",
949                code: Some(status.code()),
950            })
951        }
952        error::BigTableError::AttemptTimeout if kind != RetryKind::ConditionalWrite => {
953            Some(RetryMetric {
954                error: "AttemptTimeout",
955                code: None,
956            })
957        }
958        // Authentication and request construction happen before the RPC is
959        // dispatched, so their deadline is safe to retry for every method,
960        // including conditional mutations.
961        error::BigTableError::PreSendTimeout => Some(RetryMetric {
962            error: "PreSendTimeout",
963            code: None,
964        }),
965        // Failures to fetch an OAuth token (e.g. a transient metadata server
966        // hiccup) are retryable, matching grpcio's prior behavior.
967        error::BigTableError::Auth(_) => Some(RetryMetric {
968            error: "Auth",
969            code: None,
970        }),
971        _ => None,
972    }
973}
974
975/// Determine if a router record is "incomplete" (doesn't include [User]
976/// columns):
977///
978/// They can be incomplete for a couple reasons:
979///
980/// 1) A migration code bug caused a few incomplete migrations where
981///    `add_channels` and `increment_storage` calls occurred when the migration's
982///    initial `add_user` was never completed:
983///    https://github.com/mozilla-services/autopush-rs/pull/640
984///
985/// 2) When router TTLs are eventually enabled: `add_channel` and
986///    `increment_storage` can write cells with later expiry times than the other
987///    router cells
988fn is_incomplete_router_record(cells: &RowCells) -> bool {
989    cells
990        .keys()
991        .all(|k| ["current_timestamp", "version"].contains(&k.as_str()) || k.starts_with("chid:"))
992}
993
994/// Normalize a record's `router_type` into a metric tag: the notification's
995/// eventual destination.
996///
997/// Anything unrecognized (a corrupt cell, or a router type retired since the
998/// record was written) collapses to `unknown` so a bad record can't blow up
999/// the tag's cardinality.
1000fn router_type_tag(router_type: &str) -> &'static str {
1001    match router_type {
1002        "webpush" => "webpush",
1003        "fcm" => "fcm",
1004        "gcm" => "gcm",
1005        "apns" => "apns",
1006        "stub" => "stub",
1007        _ => "unknown",
1008    }
1009}
1010
1011/// The bridge application id from a user's `router_data`, for use as a metric
1012/// tag. This is the closest the router record comes to identifying a device
1013/// type: it's the `{app_id}` the client registered under, which distinguishes
1014/// the Firefox flavor (and so, in practice, the platform build) behind a
1015/// bridged record.
1016///
1017/// APNS stores it as `rel_channel` and FCM as `app_id`; WebPush records carry
1018/// neither. Both are validated against the configured bridge clients at
1019/// registration, so the value stays low cardinality.
1020fn app_id_tag(router_data: Option<&HashMap<String, serde_json::Value>>) -> &str {
1021    router_data
1022        .and_then(|data| {
1023            data.get("rel_channel")
1024                .or_else(|| data.get("app_id"))
1025                .and_then(serde_json::Value::as_str)
1026        })
1027        .unwrap_or("none")
1028}
1029
1030#[cfg(test)]
1031mod metric_tag_tests {
1032    use super::{app_id_tag, router_type_tag};
1033    use serde_json::json;
1034    use std::collections::HashMap;
1035
1036    /// Known router types pass through; anything else collapses to a single
1037    /// bucket so the tag stays bounded.
1038    #[test]
1039    fn router_type_clamps_to_known_destinations() {
1040        assert_eq!(router_type_tag("apns"), "apns");
1041        assert_eq!(router_type_tag("webpush"), "webpush");
1042        assert_eq!(router_type_tag("gcm"), "gcm");
1043        assert_eq!(router_type_tag("APNS"), "unknown");
1044        assert_eq!(router_type_tag(""), "unknown");
1045    }
1046
1047    /// APNS keys the app id as `rel_channel` and FCM as `app_id`; WebPush
1048    /// records carry no `router_data` at all.
1049    #[test]
1050    fn app_id_reads_either_bridge_key() {
1051        let apns: HashMap<_, _> = [
1052            ("token".to_owned(), json!("deadbeef")),
1053            ("rel_channel".to_owned(), json!("firefox")),
1054        ]
1055        .into();
1056        assert_eq!(app_id_tag(Some(&apns)), "firefox");
1057
1058        let fcm: HashMap<_, _> = [("app_id".to_owned(), json!("fennec"))].into();
1059        assert_eq!(app_id_tag(Some(&fcm)), "fennec");
1060
1061        assert_eq!(app_id_tag(None), "none");
1062        assert_eq!(app_id_tag(Some(&HashMap::new())), "none");
1063        // a non-string value must not panic or leak a serialized blob
1064        let odd: HashMap<_, _> = [("app_id".to_owned(), json!(7))].into();
1065        assert_eq!(app_id_tag(Some(&odd)), "none");
1066    }
1067}
1068
1069/// Connect to a BigTable storage model.
1070///
1071/// BigTable is available via the Google Console, and is a schema less storage system.
1072///
1073/// The `db_dsn` string should be in the form of
1074/// `grpc://{BigTableEndpoint}`
1075///
1076/// The settings contains the `table_name` which is the GRPC path to the data.
1077/// (e.g. `projects/{project_id}/instances/{instance_id}/tables/{table_id}`)
1078///
1079/// where:
1080/// _BigTableEndpoint_ is the endpoint domain to use (the default is `bigtable.googleapis.com`) See
1081/// [BigTable Endpoints](https://cloud.google.com/bigtable/docs/regional-endpoints) for more details.
1082/// _project-id_ is the Google project identifier (see the Google developer console (e.g. 'autopush-dev'))
1083/// _instance-id_ is the Google project instance, (see the Google developer console (e.g. 'development-1'))
1084/// _table_id_ is the Table Name (e.g. 'autopush')
1085///
1086/// This will automatically bring in the default credentials specified by the `GOOGLE_APPLICATION_CREDENTIALS`
1087/// environment variable.
1088///
1089/// NOTE: Some configurations may look for the default credential file (pointed to by
1090/// `GOOGLE_APPLICATION_CREDENTIALS`) to be stored in
1091/// `$HOME/.config/gcloud/application_default_credentials.json`)
1092///
1093impl BigTableClientImpl {
1094    pub fn new(metrics: Arc<StatsdClient>, settings: &DbSettings) -> DbResult<Self> {
1095        debug!("🏊 BT Pool new");
1096        let db_settings = BigTableDbSettings::try_from(settings.db_settings.as_ref())?;
1097        info!("🉑 {:#?}", db_settings);
1098        let pool = BigTablePool::new(settings, &metrics)?;
1099
1100        // create the metadata header blocks required by Google for accessing GRPC resources.
1101        let metadata = db_settings.metadata()?;
1102        Ok(Self {
1103            settings: db_settings,
1104            metrics,
1105            metadata,
1106            pool,
1107            circuit_breaker: Arc::new(CircuitBreaker::default()),
1108        })
1109    }
1110
1111    fn router_ttl(&self) -> Duration {
1112        self.settings
1113            .max_router_ttl
1114            .unwrap_or(Duration::from_secs(MAX_ROUTER_TTL_SECS))
1115    }
1116
1117    /// Return a ReadRowsRequest for a given row key
1118    fn read_row_request(&self, row_key: &str) -> bigtable::ReadRowsRequest {
1119        read_row_request(
1120            &self.settings.table_name,
1121            &self.settings.app_profile_id,
1122            row_key,
1123        )
1124    }
1125
1126    /// Return a MutateRowRequest for a given row key
1127    fn mutate_row_request(&self, row_key: &str) -> bigtable::MutateRowRequest {
1128        bigtable::MutateRowRequest {
1129            table_name: self.settings.table_name.clone(),
1130            app_profile_id: self.settings.app_profile_id.clone(),
1131            row_key: row_key.as_bytes().to_vec(),
1132            ..Default::default()
1133        }
1134    }
1135
1136    /// Return a CheckAndMutateRowRequest for a given row key
1137    fn check_and_mutate_row_request(&self, row_key: &str) -> bigtable::CheckAndMutateRowRequest {
1138        bigtable::CheckAndMutateRowRequest {
1139            table_name: self.settings.table_name.clone(),
1140            app_profile_id: self.settings.app_profile_id.clone(),
1141            row_key: row_key.as_bytes().to_vec(),
1142            ..Default::default()
1143        }
1144    }
1145
1146    fn rpc_policy(
1147        &self,
1148        class: RpcClass,
1149        retry_kind: RetryKind,
1150        breaker: BreakerPolicy,
1151    ) -> RpcPolicy {
1152        let (attempt_timeout, total_timeout) = match class {
1153            RpcClass::Point => (
1154                self.settings.grpc_point_attempt_timeout,
1155                self.settings.grpc_point_total_timeout,
1156            ),
1157            RpcClass::Scan => (
1158                self.settings.grpc_scan_attempt_timeout,
1159                self.settings.grpc_scan_total_timeout,
1160            ),
1161        };
1162        RpcPolicy {
1163            attempt_timeout,
1164            total_timeout,
1165            retry_kind,
1166            breaker,
1167        }
1168    }
1169
1170    /// Execute one logical Bigtable operation within a total budget.
1171    ///
1172    /// A deadpool checkout is a local concurrency permit. It is included in the
1173    /// caller's total budget, but checkout failures do not affect the Bigtable
1174    /// circuit breaker. Once checked out, every retry selects a fresh channel
1175    /// slot and rebuilds request metadata and credentials.
1176    async fn execute_rpc<M, T, Call, CallFuture>(
1177        &self,
1178        message: M,
1179        policy: RpcPolicy,
1180        call: Call,
1181    ) -> Result<T, error::BigTableError>
1182    where
1183        M: Clone,
1184        Call: Fn(Channel, Request<M>) -> CallFuture + Clone,
1185        CallFuture: Future<Output = Result<T, error::BigTableError>>,
1186    {
1187        if policy.breaker == BreakerPolicy::Track && !self.circuit_breaker.allow_request() {
1188            return Err(error::BigTableError::CircuitBreakerOpen);
1189        }
1190
1191        let operation_deadline = tokio::time::Instant::now() + policy.total_timeout;
1192        let pooled = match tokio::time::timeout_at(operation_deadline, self.pool.get()).await {
1193            Ok(result) => result?,
1194            Err(_) => return Err(error::BigTableError::OperationTimeout),
1195        };
1196        let bigtable = (*pooled).clone();
1197        let metadata = self.metadata.clone();
1198        let retry_policy = retry_policy(self.settings.retry_count);
1199        let pending_retry_metric = Arc::new(Mutex::new(None));
1200        let attempt_retry_metric = pending_retry_metric.clone();
1201        let condition_retry_metric = pending_retry_metric;
1202        let metrics = self.metrics.clone();
1203        // If the total budget cancels the retry future, this remains true only
1204        // when cancellation happened inside a dispatched RPC. Request
1205        // preparation and retry backoff leave it false.
1206        let rpc_in_flight = Arc::new(AtomicBool::new(false));
1207        let attempt_rpc_in_flight = rpc_in_flight.clone();
1208        let result = operation_budget_until(
1209            operation_deadline,
1210            retry_policy.retry_if(
1211                || {
1212                    // The retry predicate runs even for a terminal failure.
1213                    // Emit only when the next attempt actually starts.
1214                    if let Some(retry_metric) = take_pending_retry_metric(&attempt_retry_metric) {
1215                        retry_metric.send(&metrics);
1216                    }
1217                    // The final attempt may start near the end of the total
1218                    // retry budget. Clamp its grpc-timeout to that remaining
1219                    // budget so the server does not keep abandoned work alive
1220                    // after this logical operation has already returned.
1221                    let attempt_deadline = std::cmp::min(
1222                        tokio::time::Instant::now() + policy.attempt_timeout,
1223                        operation_deadline,
1224                    );
1225                    let bigtable = bigtable.clone();
1226                    let metadata = metadata.clone();
1227                    let message = message.clone();
1228                    let channel = self.pool.next_channel();
1229                    let call = call.clone();
1230                    let rpc_in_flight = attempt_rpc_in_flight.clone();
1231                    async move {
1232                        let request = bigtable
1233                            .request(message, &metadata, attempt_deadline)
1234                            .await?;
1235                        rpc_in_flight.store(true, Ordering::Relaxed);
1236                        let result =
1237                            rpc_attempt_until(attempt_deadline, call(channel, request)).await;
1238                        rpc_in_flight.store(false, Ordering::Relaxed);
1239                        result
1240                    }
1241                },
1242                move |error: &error::BigTableError| {
1243                    if let Some(retry_metric) = classify_retry(policy.retry_kind, error) {
1244                        set_pending_retry_metric(&condition_retry_metric, retry_metric);
1245                        true
1246                    } else {
1247                        false
1248                    }
1249                },
1250            ),
1251        )
1252        .await;
1253
1254        if policy.breaker == BreakerPolicy::Track {
1255            let operation_timed_out_in_rpc = rpc_in_flight.load(Ordering::Relaxed);
1256            match &result {
1257                Ok(_) => self.circuit_breaker.record_success(),
1258                Err(error) if counts_toward_breaker(error, operation_timed_out_in_rpc) => {
1259                    self.circuit_breaker.record_failure();
1260                }
1261                Err(_) => {}
1262            }
1263        }
1264        result
1265    }
1266
1267    /// Apply an idempotent mutation to one row.
1268    async fn mutate_row(
1269        &self,
1270        req: bigtable::MutateRowRequest,
1271    ) -> Result<(), error::BigTableError> {
1272        let policy = self.rpc_policy(
1273            RpcClass::Point,
1274            RetryKind::IdempotentWrite,
1275            BreakerPolicy::Track,
1276        );
1277        self.execute_rpc(req, policy, |channel, request| async move {
1278            let mut client = BigtableDb::client(channel);
1279            client
1280                .mutate_row(request)
1281                .await
1282                .map_err(error::BigTableError::Write)?;
1283            Ok(())
1284        })
1285        .await
1286    }
1287
1288    /// Read one row for the [ReadRowsRequest] (assuming only a single row was requested).
1289    async fn read_row(
1290        &self,
1291        req: bigtable::ReadRowsRequest,
1292    ) -> Result<Option<row::Row>, error::BigTableError> {
1293        let mut rows = self.read_rows_with_class(req, RpcClass::Point).await?;
1294        Ok(rows.pop_first().map(|(_, v)| v))
1295    }
1296
1297    /// Count the router cells that a `max_age` GC policy on the router family
1298    /// would currently hide, tagged by destination and bridge app id.
1299    ///
1300    /// Reads deliberately don't apply that policy (see
1301    /// [router_gc_policy_filter]), so these records are still served. This
1302    /// sizes the population that enabling it would turn into 410s, and says
1303    /// which platform would absorb them, without actually dropping anyone.
1304    ///
1305    /// `expires_at` is the `connected_at` cell's Bigtable timestamp, which is
1306    /// the record's expiry rather than its write time. `channel_cells` is what
1307    /// remains of the row once the core user cells have been taken: the
1308    /// `chid:` set, whose cells expire independently of the record itself.
1309    fn report_expired_cells(&self, user: &User, expires_at: SystemTime, channel_cells: &RowCells) {
1310        let now = SystemTime::now();
1311        let expired_channels = channel_cells
1312            .iter()
1313            .filter(|(qualifier, _)| qualifier.starts_with("chid:"))
1314            .filter(|(_, cells)| cells.last().is_some_and(|cell| cell.timestamp <= now))
1315            .count();
1316        let record_expired = expires_at <= now;
1317        if !record_expired && expired_channels == 0 {
1318            return;
1319        }
1320
1321        let router_type = router_type_tag(&user.router_type);
1322        let app_id = app_id_tag(user.router_data.as_ref());
1323        if record_expired {
1324            self.metrics
1325                .incr_with_tags(MetricName::DatabaseExpiredUser)
1326                .with_tag("router_type", router_type)
1327                .with_tag("app_id", app_id)
1328                .send();
1329        }
1330        if expired_channels > 0 {
1331            self.metrics
1332                .count_with_tags(
1333                    MetricName::DatabaseExpiredChannels.as_ref(),
1334                    expired_channels as i64,
1335                )
1336                .with_tag("router_type", router_type)
1337                .with_tag("app_id", app_id)
1338                .send();
1339        }
1340    }
1341
1342    /// Take a big table ReadRowsRequest (containing the keys and filters) and return a set of row data indexed by row key.
1343    ///
1344    ///
1345    async fn read_rows(
1346        &self,
1347        req: bigtable::ReadRowsRequest,
1348    ) -> Result<BTreeMap<RowKey, row::Row>, error::BigTableError> {
1349        self.read_rows_with_class(req, RpcClass::Scan).await
1350    }
1351
1352    async fn read_rows_with_class(
1353        &self,
1354        req: bigtable::ReadRowsRequest,
1355        class: RpcClass,
1356    ) -> Result<BTreeMap<RowKey, row::Row>, error::BigTableError> {
1357        self.read_rows_with_policy(req, class, BreakerPolicy::Track)
1358            .await
1359    }
1360
1361    async fn read_rows_with_policy(
1362        &self,
1363        req: bigtable::ReadRowsRequest,
1364        class: RpcClass,
1365        breaker: BreakerPolicy,
1366    ) -> Result<BTreeMap<RowKey, row::Row>, error::BigTableError> {
1367        let policy = self.rpc_policy(class, RetryKind::Read, breaker);
1368        self.execute_rpc(req, policy, |channel, request| async move {
1369            let mut client = BigtableDb::client(channel);
1370            let response = client
1371                .read_rows(request)
1372                .await
1373                .map_err(error::BigTableError::Read)?
1374                .into_inner();
1375            merge::RowMerger::process_chunks(response).await
1376        })
1377        .await
1378    }
1379
1380    /// write a given row.
1381    ///
1382    /// there's also `.mutate_rows` which I presume allows multiple.
1383    async fn write_row(&self, row: row::Row) -> Result<(), error::BigTableError> {
1384        let mut req = self.mutate_row_request(&row.row_key);
1385        // compile the mutations.
1386        // It's possible to do a lot here, including altering in process
1387        // mutations, clearing them, etc. It's all up for grabs until we commit
1388        // below. For now, let's just presume a write and be done.
1389        req.mutations = self.get_mutations(row.cells)?;
1390        self.mutate_row(req).await?;
1391        Ok(())
1392    }
1393
1394    /// Compile the list of mutations for this row.
1395    fn get_mutations(
1396        &self,
1397        cells: HashMap<FamilyId, Vec<crate::db::bigtable::bigtable_client::cell::Cell>>,
1398    ) -> Result<Vec<bigtable::Mutation>, error::BigTableError> {
1399        let mut mutations = Vec::new();
1400        for (family_id, cells) in cells {
1401            for cell in cells {
1402                let timestamp = cell
1403                    .timestamp
1404                    .duration_since(SystemTime::UNIX_EPOCH)
1405                    .map_err(error::BigTableError::WriteTime)?;
1406                debug!("🉑 expiring in {:?}", timestamp.as_millis());
1407                mutations.push(bigtable::Mutation {
1408                    mutation: Some(bigtable::mutation::Mutation::SetCell(
1409                        bigtable::mutation::SetCell {
1410                            family_name: family_id.clone(),
1411                            column_qualifier: cell.qualifier.clone().into_bytes(),
1412                            // Bigtable tables use millisecond cell-timestamp granularity;
1413                            // timestamp_micros must be a multiple of 1,000 or the server
1414                            // rejects the mutation with a granularity mismatch.
1415                            timestamp_micros: (timestamp.as_millis() * 1000) as i64,
1416                            value: cell.value,
1417                        },
1418                    )),
1419                });
1420            }
1421        }
1422        Ok(mutations)
1423    }
1424
1425    /// Write mutations if the row meets a condition specified by the filter.
1426    ///
1427    /// Mutations can be applied either when the filter matches (state `true`)
1428    /// or doesn't match (state `false`).
1429    ///
1430    /// Returns whether the filter matched records (which indicates whether the
1431    /// mutations were applied, depending on the state)
1432    async fn check_and_mutate_row(
1433        &self,
1434        row: row::Row,
1435        filter: bigtable::RowFilter,
1436        state: bool,
1437    ) -> Result<bool, error::BigTableError> {
1438        let mut req = self.check_and_mutate_row_request(&row.row_key);
1439        let mutations = self.get_mutations(row.cells)?;
1440        req.predicate_filter = Some(filter);
1441        if state {
1442            req.true_mutations = mutations;
1443        } else {
1444            req.false_mutations = mutations;
1445        }
1446        self.check_and_mutate(req).await
1447    }
1448
1449    async fn check_and_mutate(
1450        &self,
1451        req: bigtable::CheckAndMutateRowRequest,
1452    ) -> Result<bool, error::BigTableError> {
1453        let policy = self.rpc_policy(
1454            RpcClass::Point,
1455            RetryKind::ConditionalWrite,
1456            BreakerPolicy::Track,
1457        );
1458        let predicate_matched = self
1459            .execute_rpc(req, policy, |channel, request| async move {
1460                let mut client = BigtableDb::client(channel);
1461                let response = client
1462                    .check_and_mutate_row(request)
1463                    .await
1464                    .map_err(error::BigTableError::Write)?;
1465                Ok(response.into_inner().predicate_matched)
1466            })
1467            .await?;
1468        debug!("🉑 Predicate Matched: {predicate_matched}");
1469        Ok(predicate_matched)
1470    }
1471
1472    fn get_delete_mutations(
1473        &self,
1474        family: &str,
1475        column_names: &[&str],
1476        time_range: Option<&bigtable::TimestampRange>,
1477    ) -> Result<Vec<bigtable::Mutation>, error::BigTableError> {
1478        let mut mutations = Vec::new();
1479        for column in column_names {
1480            // DeleteFromRow -- Delete all cells for a given row.
1481            // DeleteFromFamily -- Delete all cells from a family for a given row.
1482            // DeleteFromColumn -- Delete all cells from a column name for a given row, restricted by timestamp range.
1483            mutations.push(bigtable::Mutation {
1484                mutation: Some(bigtable::mutation::Mutation::DeleteFromColumn(
1485                    bigtable::mutation::DeleteFromColumn {
1486                        family_name: family.to_owned(),
1487                        column_qualifier: column.as_bytes().to_vec(),
1488                        time_range: time_range.cloned(),
1489                    },
1490                )),
1491            });
1492        }
1493        Ok(mutations)
1494    }
1495
1496    /// Delete all the cells for the given row. NOTE: This will drop the row.
1497    async fn delete_row(&self, row_key: &str) -> Result<(), error::BigTableError> {
1498        let mut req = self.mutate_row_request(row_key);
1499        req.mutations = vec![bigtable::Mutation {
1500            mutation: Some(bigtable::mutation::Mutation::DeleteFromRow(
1501                bigtable::mutation::DeleteFromRow {},
1502            )),
1503        }];
1504        self.mutate_row(req).await
1505    }
1506
1507    fn rows_to_notifications(
1508        &self,
1509        rows: BTreeMap<String, Row>,
1510    ) -> Result<Vec<Notification>, DbError> {
1511        rows.into_iter()
1512            .map(|(row_key, row)| self.row_to_notification(&row_key, row))
1513            .collect()
1514    }
1515
1516    fn row_to_notification(&self, row_key: &str, mut row: Row) -> Result<Notification, DbError> {
1517        let Some((_, chidmessageid)) = row_key.split_once('#') else {
1518            return Err(DbError::Integrity(
1519                "rows_to_notification expected row_key: uaid:chidmessageid ".to_owned(),
1520                None,
1521            ));
1522        };
1523        let range_key = RangeKey::parse_chidmessageid(chidmessageid).map_err(|e| {
1524            DbError::Integrity(
1525                format!("rows_to_notification expected chidmessageid: {e}"),
1526                None,
1527            )
1528        })?;
1529
1530        // Create from the known, required fields.
1531        let mut notif = Notification {
1532            channel_id: range_key.channel_id,
1533            topic: range_key.topic,
1534            sortkey_timestamp: range_key.sortkey_timestamp,
1535            version: to_string(row.take_required_cell("version")?.value, "version")?,
1536            ttl: to_u64(row.take_required_cell("ttl")?.value, "ttl")?,
1537            timestamp: to_u64(row.take_required_cell("timestamp")?.value, "timestamp")?,
1538            ..Default::default()
1539        };
1540
1541        // Backfill the Optional fields
1542        if let Some(cell) = row.take_cell("data") {
1543            notif.data = Some(to_string(cell.value, "data")?);
1544        }
1545        #[cfg(feature = "reliable_report")]
1546        {
1547            if let Some(cell) = row.take_cell("reliability_id") {
1548                notif.reliability_id = Some(to_string(cell.value, "reliability_id")?);
1549            }
1550            if let Some(cell) = row.take_cell("reliable_state") {
1551                notif.reliable_state = Some(
1552                    crate::reliability::ReliabilityState::from_str(&to_string(
1553                        cell.value,
1554                        "reliable_state",
1555                    )?)
1556                    .map_err(|e| {
1557                        DbError::DeserializeString(format!("Could not parse reliable_state {e:?}"))
1558                    })?,
1559                );
1560            }
1561        }
1562        if let Some(cell) = row.take_cell("headers") {
1563            notif.headers = Some(
1564                serde_json::from_str::<HashMap<String, String>>(&to_string(cell.value, "headers")?)
1565                    .map_err(|e| DbError::Serialization(e.to_string()))?,
1566            );
1567        }
1568        #[cfg(feature = "reliable_report")]
1569        if let Some(cell) = row.take_cell("reliability_id") {
1570            trace!("🚣  Is reliable");
1571            notif.reliability_id = Some(to_string(cell.value, "reliability_id")?);
1572        }
1573
1574        trace!("🚣  Deserialized message row: {:?}", &notif);
1575        Ok(notif)
1576    }
1577
1578    /// Return a Row for writing from a [User] and a `version`
1579    ///
1580    /// `version` is specified as an argument (ignoring [User::version]) so
1581    /// that [update_user] may specify a new version to write before modifying
1582    /// the [User] struct
1583    fn user_to_row(&self, user: &User, version: &Uuid) -> Row {
1584        let row_key = user.uaid.simple().to_string();
1585        let mut row = Row::new(row_key);
1586        let expiry = std::time::SystemTime::now() + self.router_ttl();
1587
1588        let mut cells: Vec<cell::Cell> = vec![
1589            cell::Cell {
1590                qualifier: "connected_at".to_owned(),
1591                value: user.connected_at.to_be_bytes().to_vec(),
1592                timestamp: expiry,
1593                ..Default::default()
1594            },
1595            cell::Cell {
1596                qualifier: "router_type".to_owned(),
1597                value: user.router_type.clone().into_bytes(),
1598                timestamp: expiry,
1599                ..Default::default()
1600            },
1601            cell::Cell {
1602                qualifier: "record_version".to_owned(),
1603                value: user
1604                    .record_version
1605                    .unwrap_or(USER_RECORD_VERSION)
1606                    .to_be_bytes()
1607                    .to_vec(),
1608                timestamp: expiry,
1609                ..Default::default()
1610            },
1611            cell::Cell {
1612                qualifier: "version".to_owned(),
1613                value: (*version).into(),
1614                timestamp: expiry,
1615                ..Default::default()
1616            },
1617        ];
1618
1619        if let Some(router_data) = &user.router_data {
1620            cells.push(cell::Cell {
1621                qualifier: "router_data".to_owned(),
1622                value: json!(router_data).to_string().as_bytes().to_vec(),
1623                timestamp: expiry,
1624                ..Default::default()
1625            });
1626        };
1627        if let Some(current_timestamp) = user.current_timestamp {
1628            cells.push(cell::Cell {
1629                qualifier: "current_timestamp".to_owned(),
1630                value: current_timestamp.to_be_bytes().to_vec(),
1631                timestamp: expiry,
1632                ..Default::default()
1633            });
1634        };
1635        if let Some(node_id) = &user.node_id {
1636            cells.push(cell::Cell {
1637                qualifier: "node_id".to_owned(),
1638                value: node_id.as_bytes().to_vec(),
1639                timestamp: expiry,
1640                ..Default::default()
1641            });
1642        };
1643
1644        cells.extend(channels_to_cells(
1645            Cow::Borrowed(&user.priv_channels),
1646            expiry,
1647        ));
1648
1649        row.add_cells(ROUTER_FAMILY, cells);
1650        row
1651    }
1652}
1653
1654#[derive(Clone)]
1655pub struct BigtableDb {
1656    /// Application Default Credentials token provider, shared across the
1657    /// pool. `None` when running against the emulator (which needs no
1658    /// credentials).
1659    auth_provider: Option<Arc<dyn TokenProvider>>,
1660}
1661
1662impl BigtableDb {
1663    pub fn new(auth_provider: Option<Arc<dyn TokenProvider>>) -> Self {
1664        Self { auth_provider }
1665    }
1666
1667    pub(super) fn client(channel: Channel) -> BigtableClient<Channel> {
1668        BigtableClient::new(channel)
1669            .max_decoding_message_size(MAX_MESSAGE_LEN)
1670            .max_encoding_message_size(MAX_MESSAGE_LEN)
1671    }
1672
1673    /// Build a [tonic::Request] for the given message, attaching the standard
1674    /// Google routing metadata and (outside of the emulator) an OAuth2
1675    /// `authorization` token from the Application Default Credentials.
1676    ///
1677    /// The token provider caches tokens internally, so this is cheap to call
1678    /// per-request (and per retry attempt, where it transparently picks up a
1679    /// fresh token if the prior one expired). Token preparation is bounded by
1680    /// `deadline`; the remaining time is also sent to Bigtable as the
1681    /// `grpc-timeout` metadata value. A local deadline expiry is reported as
1682    /// [error::BigTableError::PreSendTimeout].
1683    pub(super) async fn request<T>(
1684        &self,
1685        msg: T,
1686        metadata: &MetadataMap,
1687        deadline: tokio::time::Instant,
1688    ) -> Result<Request<T>, error::BigTableError> {
1689        let authorization = if let Some(provider) = &self.auth_provider {
1690            let token = tokio::time::timeout_at(deadline, provider.token(BIGTABLE_DATA_SCOPES))
1691                .await
1692                .map_err(|_| error::BigTableError::PreSendTimeout)?
1693                .map_err(error::BigTableError::Auth)?;
1694            Some(
1695                format!("Bearer {}", token.as_str())
1696                    .parse::<AsciiMetadataValue>()
1697                    .map_err(|e| {
1698                        error::BigTableError::Config(format!("Invalid auth token: {e}"))
1699                    })?,
1700            )
1701        } else {
1702            None
1703        };
1704
1705        let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
1706        if remaining.is_zero() {
1707            return Err(error::BigTableError::PreSendTimeout);
1708        }
1709
1710        let mut request = Request::new(msg);
1711        *request.metadata_mut() = metadata.clone();
1712        // Set this after copying metadata: set_timeout is implemented by
1713        // inserting the grpc-timeout entry into the request metadata map.
1714        request.set_timeout(remaining);
1715        if let Some(value) = authorization {
1716            request.metadata_mut().insert("authorization", value);
1717        }
1718        Ok(request)
1719    }
1720}
1721
1722#[async_trait]
1723impl DbClient for BigTableClientImpl {
1724    /// add user to the database
1725    async fn add_user(&self, user: &User) -> DbResult<()> {
1726        trace!("🉑 Adding user");
1727        let Some(ref version) = user.version else {
1728            return Err(DbError::General(
1729                "add_user expected a user version field".to_owned(),
1730            ));
1731        };
1732        let row = self.user_to_row(user, version);
1733
1734        // Only add when the user doesn't already exist
1735        let row_key_filter = bigtable::RowFilter {
1736            filter: Some(bigtable::row_filter::Filter::RowKeyRegexFilter(
1737                format!("^{}$", row.row_key).into_bytes(),
1738            )),
1739        };
1740        let filter = filter_chain(vec![router_gc_policy_filter(), row_key_filter]);
1741
1742        if self.check_and_mutate_row(row, filter, false).await? {
1743            return Err(DbError::Conditional);
1744        }
1745        Ok(())
1746    }
1747
1748    /// BigTable doesn't really have the concept of an "update". You simply write the data and
1749    /// the individual cells create a new version. Depending on the garbage collection rules for
1750    /// the family, these can either persist or be automatically deleted.
1751    ///
1752    /// NOTE: This function updates the key ROUTER records for a given UAID. It does this by
1753    /// calling [BigTableClientImpl::user_to_row] which creates a new row with new `cell.timestamp` values set
1754    /// to now + `MAX_ROUTER_TTL`. This function is called by mobile during the daily
1755    /// [autoendpoint::routes::update_token_route] handling, and by desktop
1756    /// [autoconnect-ws-sm::get_or_create_user]` which is called
1757    /// during the `HELLO` handler. This should be enough to ensure that the ROUTER records
1758    /// are properly refreshed for "lively" clients.
1759    ///
1760    /// NOTE: There is some, very small, potential risk that a desktop client that can
1761    /// somehow remain connected the duration of MAX_ROUTER_TTL, may be dropped as not being
1762    /// "lively".
1763    async fn update_user(&self, user: &mut User) -> DbResult<bool> {
1764        let Some(ref version) = user.version else {
1765            return Err(DbError::General(
1766                "update_user expected a user version field".to_owned(),
1767            ));
1768        };
1769
1770        let mut filters = vec![router_gc_policy_filter()];
1771        filters.extend(version_filter(version));
1772        let filter = filter_chain(filters);
1773
1774        let new_version = Uuid::new_v4();
1775        // Always write a newly generated version
1776        let row = self.user_to_row(user, &new_version);
1777
1778        let predicate_matched = self.check_and_mutate_row(row, filter, true).await?;
1779        user.version = Some(new_version);
1780        Ok(predicate_matched)
1781    }
1782
1783    async fn get_user(&self, uaid: &Uuid) -> DbResult<Option<User>> {
1784        let row_key = uaid.as_simple().to_string();
1785        let mut req = self.read_row_request(&row_key);
1786        let mut filters = vec![router_gc_policy_filter()];
1787        filters.push(family_filter(format!("^{ROUTER_FAMILY}$")));
1788        req.filter = Some(filter_chain(filters));
1789        let Some(mut row) = self.read_row(req).await? else {
1790            return Ok(None);
1791        };
1792
1793        trace!("🉑 Found a record for {}", row_key);
1794
1795        let connected_at_cell = match row.take_required_cell("connected_at") {
1796            Ok(cell) => cell,
1797            Err(_) => {
1798                if !is_incomplete_router_record(&row.cells) {
1799                    return Err(DbError::Integrity(
1800                        "Expected column: connected_at".to_owned(),
1801                        Some(format!("{row:#?}")),
1802                    ));
1803                }
1804                // Special case incomplete records: they're equivalent to no
1805                // user exists. Incompletes caused by the migration bug in #640
1806                // will have their migration re-triggered by returning None:
1807                // https://github.com/mozilla-services/autopush-rs/pull/640
1808                trace!("🉑 Dropping an incomplete user record for {}", row_key);
1809                self.metrics
1810                    .incr_with_tags(MetricName::DatabaseDropUser)
1811                    .with_tag("reason", "incomplete_record")
1812                    .send();
1813                self.remove_user(uaid).await?;
1814                return Ok(None);
1815            }
1816        };
1817
1818        // The cell's Bigtable timestamp is the record's expiry, not its write
1819        // time. Read it before the cell is consumed below.
1820        let expires_at = connected_at_cell.timestamp;
1821
1822        let mut result = User {
1823            uaid: *uaid,
1824            connected_at: to_u64(connected_at_cell.value, "connected_at")?,
1825            router_type: to_string(row.take_required_cell("router_type")?.value, "router_type")?,
1826            record_version: Some(to_u64(
1827                row.take_required_cell("record_version")?.value,
1828                "record_version",
1829            )?),
1830            version: Some(
1831                row.take_required_cell("version")?
1832                    .value
1833                    .try_into()
1834                    .map_err(|e| {
1835                        DbError::Serialization(format!("Could not deserialize version: {e:?}"))
1836                    })?,
1837            ),
1838            ..Default::default()
1839        };
1840
1841        if let Some(cell) = row.take_cell("router_data") {
1842            result.router_data = from_str(&to_string(cell.value, "router_type")?).map_err(|e| {
1843                DbError::Serialization(format!("Could not deserialize router_type: {e:?}"))
1844            })?;
1845        }
1846
1847        if let Some(cell) = row.take_cell("node_id") {
1848            result.node_id = Some(to_string(cell.value, "node_id")?);
1849        }
1850
1851        if let Some(cell) = row.take_cell("current_timestamp") {
1852            result.current_timestamp = Some(to_u64(cell.value, "current_timestamp")?)
1853        }
1854
1855        // Read the channels last, after removal of all non channel cells
1856        result.priv_channels = channels_from_cells(&row.cells)?;
1857
1858        self.report_expired_cells(&result, expires_at, &row.cells);
1859
1860        Ok(Some(result))
1861    }
1862
1863    async fn remove_user(&self, uaid: &Uuid) -> DbResult<()> {
1864        let row_key = uaid.simple().to_string();
1865        self.delete_row(&row_key).await?;
1866        Ok(())
1867    }
1868
1869    async fn add_channel(&self, uaid: &Uuid, channel_id: &Uuid) -> DbResult<()> {
1870        let channels = HashSet::from_iter([channel_id.to_owned()]);
1871        self.add_channels(uaid, channels).await
1872    }
1873
1874    /// Add channels in bulk (used mostly during migration)
1875    ///
1876    async fn add_channels(&self, uaid: &Uuid, channels: HashSet<Uuid>) -> DbResult<()> {
1877        // channel_ids are stored as a set within one Bigtable row
1878        //
1879        // Bigtable allows "millions of columns in a table, as long as no row
1880        // exceeds the maximum limit of 256 MB per row" enabling the use of
1881        // column qualifiers as data.
1882        //
1883        // The "set" of channel_ids consists of column qualifiers named
1884        // "chid:<chid value>" as set member entries (with their cell values
1885        // being a single 0 byte).
1886        //
1887        // Storing the full set in a single row makes batch updates
1888        // (particularly to reset the GC expiry timestamps) potentially more
1889        // easy/efficient
1890        let row_key = uaid.simple().to_string();
1891        let mut row = Row::new(row_key);
1892        let expiry = std::time::SystemTime::now() + self.router_ttl();
1893
1894        // Note: updating the version column isn't necessary here because this
1895        // write only adds a new (or updates an existing) column with a 0 byte
1896        // value
1897        row.add_cells(
1898            ROUTER_FAMILY,
1899            channels_to_cells(Cow::Owned(channels), expiry),
1900        );
1901
1902        self.write_row(row).await?;
1903        Ok(())
1904    }
1905
1906    async fn get_channels(&self, uaid: &Uuid) -> DbResult<HashSet<Uuid>> {
1907        let row_key = uaid.simple().to_string();
1908        let mut req = self.read_row_request(&row_key);
1909
1910        let cq_filter = bigtable::RowFilter {
1911            filter: Some(bigtable::row_filter::Filter::ColumnQualifierRegexFilter(
1912                "^chid:.*$".as_bytes().to_vec(),
1913            )),
1914        };
1915        req.filter = Some(filter_chain(vec![
1916            router_gc_policy_filter(),
1917            family_filter(format!("^{ROUTER_FAMILY}$")),
1918            cq_filter,
1919        ]));
1920
1921        let Some(row) = self.read_row(req).await? else {
1922            return Ok(Default::default());
1923        };
1924        channels_from_cells(&row.cells)
1925    }
1926
1927    /// Delete the channel. Does not delete its associated pending messages.
1928    async fn remove_channel(&self, uaid: &Uuid, channel_id: &Uuid) -> DbResult<bool> {
1929        let row_key = uaid.simple().to_string();
1930        let mut req = self.check_and_mutate_row_request(&row_key);
1931
1932        // Delete the column representing the channel_id
1933        let column = format!("chid:{}", channel_id.as_hyphenated());
1934        let mut mutations = self.get_delete_mutations(ROUTER_FAMILY, &[column.as_ref()], None)?;
1935
1936        // and write a new version cell
1937        let mut row = Row::new(row_key);
1938        let expiry = std::time::SystemTime::now() + self.router_ttl();
1939        row.cells
1940            .insert(ROUTER_FAMILY.to_owned(), vec![new_version_cell(expiry)]);
1941        mutations.extend(self.get_mutations(row.cells)?);
1942
1943        // check if the channel existed/was actually removed
1944        let cq_filter = bigtable::RowFilter {
1945            filter: Some(bigtable::row_filter::Filter::ColumnQualifierRegexFilter(
1946                format!("^{column}$").into_bytes(),
1947            )),
1948        };
1949        req.predicate_filter = Some(filter_chain(vec![router_gc_policy_filter(), cq_filter]));
1950        req.true_mutations = mutations;
1951
1952        Ok(self.check_and_mutate(req).await?)
1953    }
1954
1955    /// Remove the node_id
1956    async fn remove_node_id(
1957        &self,
1958        uaid: &Uuid,
1959        _node_id: &str,
1960        _connected_at: u64,
1961        version: &Option<Uuid>,
1962    ) -> DbResult<bool> {
1963        let row_key = uaid.simple().to_string();
1964        trace!("🉑 Removing node_id for: {row_key} (version: {version:?}) ",);
1965        let Some(version) = version else {
1966            return Err(DbError::General("Expected a user version field".to_owned()));
1967        };
1968
1969        let mut req = self.check_and_mutate_row_request(&row_key);
1970
1971        let mut filters = vec![router_gc_policy_filter()];
1972        filters.extend(version_filter(version));
1973        req.predicate_filter = Some(filter_chain(filters));
1974        req.true_mutations = self.get_delete_mutations(ROUTER_FAMILY, &["node_id"], None)?;
1975
1976        Ok(self.check_and_mutate(req).await?)
1977    }
1978
1979    /// Write the notification to storage.
1980    async fn save_message(&self, uaid: &Uuid, message: Notification) -> DbResult<()> {
1981        let row_key = format!("{}#{}", uaid.simple(), message.chidmessageid());
1982        debug!("🗄️ Saving message {} :: {:?}", &row_key, &message);
1983        trace!(
1984            "🉑 timestamp: {:?}",
1985            &message.timestamp.to_be_bytes().to_vec()
1986        );
1987        let mut row = Row::new(row_key);
1988
1989        // Remember, `timestamp` is effectively the time to kill the message, not the
1990        // current time.
1991        // TODO: use message.expiry()
1992        let expiry = SystemTime::now() + Duration::from_secs(message.ttl);
1993        trace!(
1994            "🉑 Message Expiry {}",
1995            expiry
1996                .duration_since(SystemTime::UNIX_EPOCH)
1997                .unwrap_or_default()
1998                .as_millis()
1999        );
2000
2001        let mut cells: Vec<cell::Cell> = Vec::new();
2002
2003        let is_topic = message.topic.is_some();
2004        let family = if is_topic {
2005            MESSAGE_TOPIC_FAMILY
2006        } else {
2007            MESSAGE_FAMILY
2008        };
2009        cells.extend(vec![
2010            cell::Cell {
2011                qualifier: "ttl".to_owned(),
2012                value: message.ttl.to_be_bytes().to_vec(),
2013                timestamp: expiry,
2014                ..Default::default()
2015            },
2016            cell::Cell {
2017                qualifier: "timestamp".to_owned(),
2018                value: message.timestamp.to_be_bytes().to_vec(),
2019                timestamp: expiry,
2020                ..Default::default()
2021            },
2022            cell::Cell {
2023                qualifier: "version".to_owned(),
2024                value: message.version.into_bytes(),
2025                timestamp: expiry,
2026                ..Default::default()
2027            },
2028        ]);
2029        if let Some(headers) = message.headers
2030            && !headers.is_empty()
2031        {
2032            cells.push(cell::Cell {
2033                qualifier: "headers".to_owned(),
2034                value: json!(headers).to_string().into_bytes(),
2035                timestamp: expiry,
2036                ..Default::default()
2037            });
2038        }
2039        #[cfg(feature = "reliable_report")]
2040        {
2041            if let Some(reliability_id) = message.reliability_id {
2042                trace!("🔍 FOUND RELIABILITY ID: {}", reliability_id);
2043                cells.push(cell::Cell {
2044                    qualifier: "reliability_id".to_owned(),
2045                    value: reliability_id.into_bytes(),
2046                    timestamp: expiry,
2047                    ..Default::default()
2048                });
2049            }
2050            if let Some(reliable_state) = message.reliable_state {
2051                cells.push(cell::Cell {
2052                    qualifier: "reliable_state".to_owned(),
2053                    value: reliable_state.to_string().into_bytes(),
2054                    timestamp: expiry,
2055                    ..Default::default()
2056                });
2057            }
2058        }
2059        if let Some(data) = message.data {
2060            cells.push(cell::Cell {
2061                qualifier: "data".to_owned(),
2062                value: data.into_bytes(),
2063                timestamp: expiry,
2064                ..Default::default()
2065            });
2066        }
2067
2068        row.add_cells(family, cells);
2069        trace!("🉑 Adding row");
2070        self.write_row(row).await?;
2071
2072        self.metrics
2073            .incr_with_tags(MetricName::NotificationMessageStored)
2074            .with_tag("topic", &is_topic.to_string())
2075            .with_tag("database", &self.name())
2076            .send();
2077        Ok(())
2078    }
2079
2080    /// Save a batch of messages to the database.
2081    ///
2082    /// Currently just iterating through the list and saving one at a time. There's a bulk way
2083    /// to save messages, but there are other considerations (e.g. mutation limits)
2084    async fn save_messages(&self, uaid: &Uuid, messages: Vec<Notification>) -> DbResult<()> {
2085        // plate simple way of solving this:
2086        for message in messages {
2087            self.save_message(uaid, message).await?;
2088        }
2089        Ok(())
2090    }
2091
2092    /// Set the `current_timestamp` in the meta record for this user agent.
2093    ///
2094    /// This is a bit different for BigTable. Field expiration (technically cell
2095    /// expiration) is determined by the lifetime assigned to the cell once it hits
2096    /// a given date. That means you can't really extend a lifetime by adjusting a
2097    /// single field. You'd have to adjust all the cells that are in the family.
2098    /// So, we're not going to do expiration that way.
2099    ///
2100    /// That leaves the meta "current_timestamp" field. We do not purge ACK'd records,
2101    /// instead we presume that the TTL will kill them off eventually. On reads, we use
2102    /// the `current_timestamp` to determine what records to return, since we return
2103    /// records with timestamps later than `current_timestamp`.
2104    ///
2105    async fn increment_storage(&self, uaid: &Uuid, timestamp: u64) -> DbResult<()> {
2106        let row_key = uaid.simple().to_string();
2107        debug!(
2108            "🉑 Updating {} current_timestamp:  {:?}",
2109            &row_key,
2110            timestamp.to_be_bytes().to_vec()
2111        );
2112        let expiry = std::time::SystemTime::now() + self.router_ttl();
2113        let mut row = Row::new(row_key.clone());
2114
2115        row.cells.insert(
2116            ROUTER_FAMILY.to_owned(),
2117            vec![
2118                cell::Cell {
2119                    qualifier: "current_timestamp".to_owned(),
2120                    value: timestamp.to_be_bytes().to_vec(),
2121                    timestamp: expiry,
2122                    ..Default::default()
2123                },
2124                new_version_cell(expiry),
2125            ],
2126        );
2127
2128        self.write_row(row).await?;
2129
2130        Ok(())
2131    }
2132
2133    /// Delete the notification from storage.
2134    async fn remove_message(&self, uaid: &Uuid, chidmessageid: &str) -> DbResult<()> {
2135        trace!(
2136            "🉑 attemping to delete {:?} :: {:?}",
2137            uaid.to_string(),
2138            chidmessageid
2139        );
2140        let row_key = format!("{}#{}", uaid.simple(), chidmessageid);
2141        debug!("🉑🔥 Deleting message {}", &row_key);
2142        self.delete_row(&row_key).await?;
2143        self.metrics
2144            .incr_with_tags(MetricName::NotificationMessageDeleted)
2145            .with_tag("database", &self.name())
2146            .send();
2147        Ok(())
2148    }
2149
2150    /// Return `limit` pending messages from storage. `limit=0` for all messages.
2151    async fn fetch_topic_messages(
2152        &self,
2153        uaid: &Uuid,
2154        limit: usize,
2155    ) -> DbResult<FetchMessageResponse> {
2156        let start_key = format!("{}#01:", uaid.simple());
2157        let end_key = format!("{}#02:", uaid.simple());
2158        let mut req = bigtable::ReadRowsRequest {
2159            table_name: self.settings.table_name.clone(),
2160            app_profile_id: self.settings.app_profile_id.clone(),
2161            rows: Some(bigtable::RowSet {
2162                row_keys: Vec::new(),
2163                row_ranges: vec![bigtable::RowRange {
2164                    start_key: Some(bigtable::row_range::StartKey::StartKeyOpen(
2165                        start_key.into_bytes(),
2166                    )),
2167                    end_key: Some(bigtable::row_range::EndKey::EndKeyOpen(
2168                        end_key.into_bytes(),
2169                    )),
2170                }],
2171            }),
2172            ..Default::default()
2173        };
2174
2175        let mut filters = message_gc_policy_filter()?;
2176        filters.push(family_filter(format!("^{MESSAGE_TOPIC_FAMILY}$")));
2177
2178        req.filter = Some(filter_chain(filters));
2179        if limit > 0 {
2180            trace!("🉑 Setting limit to {limit}");
2181            req.rows_limit = limit as i64;
2182        }
2183        let rows = self.read_rows(req).await?;
2184        debug!(
2185            "🉑 Fetch Topic Messages. Found {} row(s) of {}",
2186            rows.len(),
2187            limit
2188        );
2189
2190        let messages = self.rows_to_notifications(rows)?;
2191
2192        // Note: Bigtable always returns a timestamp of None.
2193        // Under Bigtable `current_timestamp` is instead initially read
2194        // from [get_user].
2195        Ok(FetchMessageResponse {
2196            messages,
2197            timestamp: None,
2198        })
2199    }
2200
2201    /// Return `limit` messages pending for a UAID that have a sortkey_timestamp after
2202    /// what's specified. `limit=0` for all messages.
2203    async fn fetch_timestamp_messages(
2204        &self,
2205        uaid: &Uuid,
2206        timestamp: Option<u64>,
2207        limit: usize,
2208    ) -> DbResult<FetchMessageResponse> {
2209        let start_key = if let Some(ts) = timestamp {
2210            // Fetch everything after the last message with timestamp: the "z"
2211            // moves past the last message's channel_id's 1st hex digit
2212            format!("{}#02:{}z", uaid.simple(), ts)
2213        } else {
2214            format!("{}#02:", uaid.simple())
2215        };
2216        let end_key = format!("{}#03:", uaid.simple());
2217        let mut req = bigtable::ReadRowsRequest {
2218            table_name: self.settings.table_name.clone(),
2219            app_profile_id: self.settings.app_profile_id.clone(),
2220            rows: Some(bigtable::RowSet {
2221                row_keys: Vec::new(),
2222                row_ranges: vec![bigtable::RowRange {
2223                    start_key: Some(bigtable::row_range::StartKey::StartKeyOpen(
2224                        start_key.into_bytes(),
2225                    )),
2226                    end_key: Some(bigtable::row_range::EndKey::EndKeyOpen(
2227                        end_key.into_bytes(),
2228                    )),
2229                }],
2230            }),
2231            ..Default::default()
2232        };
2233
2234        // We can fetch data and do [some remote filtering](https://cloud.google.com/bigtable/docs/filters),
2235        // unfortunately I don't think the filtering we need will be super helpful.
2236        //
2237        //
2238        /*
2239        //NOTE: if you filter on a given field, BigTable will only
2240        // return that specific field. Adding filters for the rest of
2241        // the known elements may NOT return those elements or may
2242        // cause the message to not be returned because any of
2243        // those elements are not present. It may be preferable to
2244        // therefore run two filters, one to fetch the candidate IDs
2245        // and another to fetch the content of the messages.
2246         */
2247        let mut filters = message_gc_policy_filter()?;
2248        filters.push(family_filter(format!("^{MESSAGE_FAMILY}$")));
2249
2250        req.filter = Some(filter_chain(filters));
2251        if limit > 0 {
2252            req.rows_limit = limit as i64;
2253        }
2254        let rows = self.read_rows(req).await?;
2255        debug!(
2256            "🉑 Fetch Timestamp Messages ({:?}) Found {} row(s) of {}",
2257            timestamp,
2258            rows.len(),
2259            limit,
2260        );
2261
2262        let messages = self.rows_to_notifications(rows)?;
2263        // The timestamp of the last message read
2264        let timestamp = messages.last().and_then(|m| m.sortkey_timestamp);
2265        Ok(FetchMessageResponse {
2266            messages,
2267            timestamp,
2268        })
2269    }
2270
2271    async fn health_check(&self) -> DbResult<bool> {
2272        // Use a random key so health checks do not create a hot tablet. The
2273        // BlockAll filter verifies the data path without returning row data.
2274        let random_uaid = Uuid::new_v4().simple().to_string();
2275        let mut req = read_row_request(
2276            &self.settings.table_name,
2277            &self.settings.app_profile_id,
2278            &random_uaid,
2279        );
2280        req.filter = Some(bigtable::RowFilter {
2281            filter: Some(bigtable::row_filter::Filter::BlockAllFilter(true)),
2282        });
2283        // Health is an independent, bounded backend probe. It must neither be
2284        // blocked by nor mutate the application traffic circuit breaker.
2285        self.read_rows_with_policy(req, RpcClass::Point, BreakerPolicy::Ignore)
2286            .await?;
2287        Ok(true)
2288    }
2289
2290    /// Returns true, because there's only one table in BigTable. We divide things up
2291    /// by `family`.
2292    async fn router_table_exists(&self) -> DbResult<bool> {
2293        Ok(true)
2294    }
2295
2296    /// Returns true, because there's only one table in BigTable. We divide things up
2297    /// by `family`.
2298    async fn message_table_exists(&self) -> DbResult<bool> {
2299        Ok(true)
2300    }
2301
2302    #[cfg(feature = "reliable_report")]
2303    async fn log_report(
2304        &self,
2305        reliability_id: &str,
2306        new_state: crate::reliability::ReliabilityState,
2307    ) -> DbResult<()> {
2308        let row_key = reliability_id.to_owned();
2309
2310        let mut row = Row::new(row_key);
2311        let expiry = SystemTime::now() + Duration::from_secs(RELIABLE_LOG_TTL.num_seconds() as u64);
2312
2313        // Log the latest transition time for this id.
2314        let cells: Vec<cell::Cell> = vec![cell::Cell {
2315            qualifier: new_state.to_string(),
2316            value: crate::util::ms_since_epoch().to_be_bytes().to_vec(),
2317            timestamp: expiry,
2318            ..Default::default()
2319        }];
2320
2321        row.add_cells(RELIABLE_LOG_FAMILY, cells);
2322
2323        self.write_row(row).await?;
2324
2325        Ok(())
2326    }
2327
2328    fn box_clone(&self) -> Box<dyn DbClient> {
2329        Box::new(self.clone())
2330    }
2331
2332    fn name(&self) -> String {
2333        "Bigtable".to_owned()
2334    }
2335
2336    fn pool_status(&self) -> Option<deadpool::Status> {
2337        Some(self.pool.pool.status())
2338    }
2339
2340    fn configured_channel_count(&self) -> Option<usize> {
2341        Some(self.pool.configured_channel_count())
2342    }
2343}
2344
2345#[cfg(all(test, feature = "emulator"))]
2346mod tests {
2347
2348    //! Currently, these test rely on having a BigTable emulator running on the current machine.
2349    //! The tests presume to be able to connect to localhost:8086. See docs/bigtable.md for
2350    //! details and how to set up and initialize an emulator.
2351    //!
2352    use std::sync::Arc;
2353    use std::time::SystemTime;
2354
2355    use cadence::StatsdClient;
2356    use uuid;
2357
2358    use super::*;
2359    use crate::{db::DbSettings, test_support::gen_test_uaid, util::ms_since_epoch};
2360
2361    const TEST_USER: &str = "DEADBEEF-0000-0000-0000-0123456789AB";
2362    const TEST_CHID: &str = "DECAFBAD-0000-0000-0000-0123456789AB";
2363    const TOPIC_CHID: &str = "DECAFBAD-1111-0000-0000-0123456789AB";
2364
2365    fn now() -> u64 {
2366        SystemTime::now()
2367            .duration_since(SystemTime::UNIX_EPOCH)
2368            .unwrap()
2369            .as_secs()
2370    }
2371
2372    fn new_client() -> DbResult<BigTableClientImpl> {
2373        let env_dsn = format!(
2374            "grpc://{}",
2375            std::env::var("BIGTABLE_EMULATOR_HOST").unwrap_or("localhost:8080".to_owned())
2376        );
2377        let settings = DbSettings {
2378            // this presumes the table was created with
2379            // ```
2380            // scripts/setup_bt.sh
2381            // ```
2382            // with `message`, `router`, and `message_topic` families
2383            dsn: Some(env_dsn),
2384            db_settings: json!({"table_name": "projects/test/instances/test/tables/autopush"})
2385                .to_string(),
2386        };
2387
2388        let metrics = Arc::new(StatsdClient::builder("", cadence::NopMetricSink).build());
2389
2390        BigTableClientImpl::new(metrics, &settings)
2391    }
2392
2393    #[test]
2394    fn escape_bytes_for_regex() {
2395        let b = b"hi";
2396        assert_eq!(escape_bytes(b), b.to_vec());
2397        assert_eq!(escape_bytes(b"h.*i!"), b"h\\.\\*i\\!".to_vec());
2398        let b = b"\xe2\x80\xb3";
2399        assert_eq!(escape_bytes(b), b.to_vec());
2400        // clippy::octal-escapes rightly discourages this ("\022") in a byte literal
2401        let b = [b'f', b'o', b'\0', b'2', b'2', b'o'];
2402        assert_eq!(escape_bytes(&b), b"fo\\x0022o".to_vec());
2403        let b = b"\xc0";
2404        assert_eq!(escape_bytes(b), b.to_vec());
2405        assert_eq!(escape_bytes(b"\x03"), b"\\\x03".to_vec());
2406    }
2407
2408    #[actix_rt::test]
2409    async fn health_check() {
2410        let client = new_client().unwrap();
2411
2412        let result = client.health_check().await;
2413        assert!(result.is_ok());
2414        assert!(result.unwrap());
2415    }
2416
2417    /// Bigtable rejects SetCell timestamps that are not aligned to the table's
2418    /// millisecond granularity (timestamp_micros must be a multiple of 1,000).
2419    #[actix_rt::test]
2420    async fn timestamp_granularity_rejects_sub_millisecond_micros() {
2421        let client = new_client().unwrap();
2422        let uaid = gen_test_uaid();
2423        let row_key = uaid.simple().to_string();
2424        let _ = client.remove_user(&uaid).await;
2425
2426        let mut req = client.mutate_row_request(&row_key);
2427        req.mutations = vec![bigtable::Mutation {
2428            mutation: Some(bigtable::mutation::Mutation::SetCell(
2429                bigtable::mutation::SetCell {
2430                    family_name: ROUTER_FAMILY.to_owned(),
2431                    column_qualifier: b"granularity_probe".to_vec(),
2432                    timestamp_micros: 1,
2433                    value: b"x".to_vec(),
2434                },
2435            )),
2436        }];
2437
2438        assert!(
2439            client.mutate_row(req).await.is_err(),
2440            "expected granularity mismatch for timestamp_micros=1"
2441        );
2442
2443        let _ = client.remove_user(&uaid).await;
2444    }
2445
2446    /// run a gauntlet of testing. These are a bit linear because they need
2447    /// to run in sequence.
2448    #[actix_rt::test]
2449    async fn run_gauntlet() -> DbResult<()> {
2450        let client = new_client()?;
2451
2452        let connected_at = ms_since_epoch();
2453
2454        let uaid = Uuid::parse_str(TEST_USER).unwrap();
2455        let chid = Uuid::parse_str(TEST_CHID).unwrap();
2456        let topic_chid = Uuid::parse_str(TOPIC_CHID).unwrap();
2457
2458        let node_id = "test_node".to_owned();
2459
2460        // purge the user record if it exists.
2461        let _ = client.remove_user(&uaid).await;
2462
2463        let test_user = User {
2464            uaid,
2465            router_type: "webpush".to_owned(),
2466            connected_at,
2467            router_data: None,
2468            node_id: Some(node_id.clone()),
2469            ..Default::default()
2470        };
2471
2472        // purge the old user (if present)
2473        // in case a prior test failed for whatever reason.
2474        let _ = client.remove_user(&uaid).await;
2475
2476        // can we add the user?
2477        client.add_user(&test_user).await?;
2478        let fetched = client.get_user(&uaid).await?;
2479        assert!(fetched.is_some());
2480        let fetched = fetched.unwrap();
2481        assert_eq!(fetched.router_type, "webpush".to_owned());
2482
2483        // Simulate a connected_at occuring before the following writes
2484        let connected_at = ms_since_epoch();
2485
2486        // can we add channels?
2487        client.add_channel(&uaid, &chid).await?;
2488        let channels = client.get_channels(&uaid).await?;
2489        assert!(channels.contains(&chid));
2490
2491        // can we add lots of channels?
2492        let mut new_channels: HashSet<Uuid> = HashSet::new();
2493        new_channels.insert(chid);
2494        for _ in 1..10 {
2495            new_channels.insert(uuid::Uuid::new_v4());
2496        }
2497        let chid_to_remove = uuid::Uuid::new_v4();
2498        new_channels.insert(chid_to_remove);
2499        client.add_channels(&uaid, new_channels.clone()).await?;
2500        let channels = client.get_channels(&uaid).await?;
2501        assert_eq!(channels, new_channels);
2502
2503        // can we remove a channel?
2504        assert!(client.remove_channel(&uaid, &chid_to_remove).await?);
2505        assert!(!client.remove_channel(&uaid, &chid_to_remove).await?);
2506        new_channels.remove(&chid_to_remove);
2507        let channels = client.get_channels(&uaid).await?;
2508        assert_eq!(channels, new_channels);
2509
2510        // now ensure that we can update a user that's after the time we set
2511        // prior. first ensure that we can't update a user that's before the
2512        // time we set prior to the last write
2513        let mut updated = User {
2514            connected_at,
2515            ..test_user.clone()
2516        };
2517        let result = client.update_user(&mut updated).await;
2518        assert!(result.is_ok());
2519        assert!(!result.unwrap());
2520
2521        // Make sure that the `connected_at` wasn't modified
2522        let fetched2 = client.get_user(&fetched.uaid).await?.unwrap();
2523        assert_eq!(fetched.connected_at, fetched2.connected_at);
2524
2525        // and make sure we can update a record with a later connected_at time.
2526        let mut updated = User {
2527            connected_at: fetched.connected_at + 300,
2528            ..fetched2
2529        };
2530        let result = client.update_user(&mut updated).await;
2531        assert!(result.is_ok());
2532        assert!(result.unwrap());
2533        assert_ne!(
2534            fetched2.connected_at,
2535            client.get_user(&uaid).await?.unwrap().connected_at
2536        );
2537
2538        // can we increment the storage for the user?
2539        client
2540            .increment_storage(
2541                &fetched.uaid,
2542                SystemTime::now()
2543                    .duration_since(SystemTime::UNIX_EPOCH)
2544                    .unwrap()
2545                    .as_secs(),
2546            )
2547            .await?;
2548
2549        let test_data = "An_encrypted_pile_of_crap".to_owned();
2550        let timestamp = now();
2551        let sort_key = now();
2552        // Can we store a message?
2553        let test_notification = crate::db::Notification {
2554            channel_id: chid,
2555            version: "test".to_owned(),
2556            ttl: 300,
2557            timestamp,
2558            data: Some(test_data.clone()),
2559            sortkey_timestamp: Some(sort_key),
2560            ..Default::default()
2561        };
2562        let res = client.save_message(&uaid, test_notification.clone()).await;
2563        assert!(res.is_ok());
2564
2565        let mut fetched = client.fetch_timestamp_messages(&uaid, None, 999).await?;
2566        assert_ne!(fetched.messages.len(), 0);
2567        let fm = fetched.messages.pop().unwrap();
2568        assert_eq!(fm.channel_id, test_notification.channel_id);
2569        assert_eq!(fm.data, Some(test_data));
2570
2571        // Grab all 1 of the messages that were submmited within the past 10 seconds.
2572        let fetched = client
2573            .fetch_timestamp_messages(&uaid, Some(timestamp - 10), 999)
2574            .await?;
2575        assert_ne!(fetched.messages.len(), 0);
2576
2577        // Try grabbing a message for 10 seconds from now.
2578        let fetched = client
2579            .fetch_timestamp_messages(&uaid, Some(timestamp + 10), 999)
2580            .await?;
2581        assert_eq!(fetched.messages.len(), 0);
2582
2583        // can we clean up our toys?
2584        assert!(
2585            client
2586                .remove_message(&uaid, &test_notification.chidmessageid())
2587                .await
2588                .is_ok()
2589        );
2590
2591        assert!(client.remove_channel(&uaid, &chid).await.is_ok());
2592
2593        // Now, can we do all that with topic messages
2594        client.add_channel(&uaid, &topic_chid).await?;
2595        let test_data = "An_encrypted_pile_of_crap_with_a_topic".to_owned();
2596        let timestamp = now();
2597        let sort_key = now();
2598        // Can we store a message?
2599        let test_notification = crate::db::Notification {
2600            channel_id: topic_chid,
2601            version: "test".to_owned(),
2602            ttl: 300,
2603            topic: Some("topic".to_owned()),
2604            timestamp,
2605            data: Some(test_data.clone()),
2606            sortkey_timestamp: Some(sort_key),
2607            ..Default::default()
2608        };
2609        assert!(
2610            client
2611                .save_message(&uaid, test_notification.clone())
2612                .await
2613                .is_ok()
2614        );
2615
2616        let mut fetched = client.fetch_topic_messages(&uaid, 999).await?;
2617        assert_ne!(fetched.messages.len(), 0);
2618        let fm = fetched.messages.pop().unwrap();
2619        assert_eq!(fm.channel_id, test_notification.channel_id);
2620        assert_eq!(fm.data, Some(test_data));
2621
2622        // Grab the message that was submmited.
2623        let fetched = client.fetch_topic_messages(&uaid, 999).await?;
2624        assert_ne!(fetched.messages.len(), 0);
2625
2626        // can we clean up our toys?
2627        assert!(
2628            client
2629                .remove_message(&uaid, &test_notification.chidmessageid())
2630                .await
2631                .is_ok()
2632        );
2633
2634        assert!(client.remove_channel(&uaid, &topic_chid).await.is_ok());
2635
2636        let msgs = client
2637            .fetch_timestamp_messages(&uaid, None, 999)
2638            .await?
2639            .messages;
2640        assert!(msgs.is_empty());
2641
2642        let fetched = client.get_user(&uaid).await?.unwrap();
2643        assert!(
2644            client
2645                .remove_node_id(&uaid, &node_id, connected_at, &fetched.version)
2646                .await
2647                .is_ok()
2648        );
2649        // did we remove it?
2650        let fetched = client.get_user(&uaid).await?.unwrap();
2651        assert_eq!(fetched.node_id, None);
2652
2653        assert!(client.remove_user(&uaid).await.is_ok());
2654
2655        assert!(client.get_user(&uaid).await?.is_none());
2656
2657        Ok(())
2658    }
2659
2660    #[actix_rt::test]
2661    async fn read_cells_family_id() -> DbResult<()> {
2662        let client = new_client().unwrap();
2663        let uaid = gen_test_uaid();
2664        client.remove_user(&uaid).await.unwrap();
2665
2666        let qualifier = "foo".to_owned();
2667
2668        let row_key = uaid.simple().to_string();
2669        let mut row = Row::new(row_key.clone());
2670        row.cells.insert(
2671            ROUTER_FAMILY.to_owned(),
2672            vec![cell::Cell {
2673                qualifier: qualifier.to_owned(),
2674                value: "bar".as_bytes().to_vec(),
2675                ..Default::default()
2676            }],
2677        );
2678        client.write_row(row).await.unwrap();
2679        let req = client.read_row_request(&row_key);
2680        let Some(row) = client.read_row(req).await.unwrap() else {
2681            panic!("Expected row");
2682        };
2683        assert_eq!(row.cells.len(), 1);
2684        assert_eq!(row.cells.keys().next().unwrap(), qualifier.as_str());
2685        client.remove_user(&uaid).await
2686    }
2687
2688    #[actix_rt::test]
2689    async fn add_user_existing() {
2690        let client = new_client().unwrap();
2691        let uaid = gen_test_uaid();
2692        let user = User {
2693            uaid,
2694            ..Default::default()
2695        };
2696        client.remove_user(&uaid).await.unwrap();
2697
2698        client.add_user(&user).await.unwrap();
2699        let err = client.add_user(&user).await.unwrap_err();
2700        assert!(matches!(err, DbError::Conditional));
2701    }
2702
2703    #[actix_rt::test]
2704    async fn version_check() {
2705        let client = new_client().unwrap();
2706        let uaid = gen_test_uaid();
2707        let user = User {
2708            uaid,
2709            ..Default::default()
2710        };
2711        client.remove_user(&uaid).await.unwrap();
2712
2713        client.add_user(&user).await.unwrap();
2714        let mut user = client.get_user(&uaid).await.unwrap().unwrap();
2715        assert!(client.update_user(&mut user.clone()).await.unwrap());
2716
2717        let fetched = client.get_user(&uaid).await.unwrap().unwrap();
2718        assert_ne!(user.version, fetched.version);
2719        // should now fail w/ a stale version
2720        assert!(!client.update_user(&mut user).await.unwrap());
2721
2722        client.remove_user(&uaid).await.unwrap();
2723    }
2724
2725    #[actix_rt::test]
2726    async fn lingering_chid_record() {
2727        let client = new_client().unwrap();
2728        let uaid = gen_test_uaid();
2729        let chid = Uuid::parse_str(TEST_CHID).unwrap();
2730        let user = User {
2731            uaid,
2732            ..Default::default()
2733        };
2734        client.remove_user(&uaid).await.unwrap();
2735
2736        // add_channel doesn't check for the existence of a user
2737        client.add_channel(&uaid, &chid).await.unwrap();
2738
2739        // w/ chid records in the router row, get_user should treat
2740        // this as the user not existing
2741        assert!(client.get_user(&uaid).await.unwrap().is_none());
2742
2743        client.add_user(&user).await.unwrap();
2744        // get_user should have also cleaned up the chids
2745        assert!(client.get_channels(&uaid).await.unwrap().is_empty());
2746
2747        client.remove_user(&uaid).await.unwrap();
2748    }
2749
2750    #[actix_rt::test]
2751    async fn lingering_current_timestamp() {
2752        let client = new_client().unwrap();
2753        let uaid = gen_test_uaid();
2754        client.remove_user(&uaid).await.unwrap();
2755
2756        client
2757            .increment_storage(&uaid, crate::util::sec_since_epoch())
2758            .await
2759            .unwrap();
2760        assert!(client.get_user(&uaid).await.unwrap().is_none());
2761
2762        client.remove_user(&uaid).await.unwrap();
2763    }
2764
2765    #[actix_rt::test]
2766    async fn lingering_chid_w_version_record() {
2767        let client = new_client().unwrap();
2768        let uaid = gen_test_uaid();
2769        let chid = Uuid::parse_str(TEST_CHID).unwrap();
2770        client.remove_user(&uaid).await.unwrap();
2771
2772        client.add_channel(&uaid, &chid).await.unwrap();
2773        assert!(client.remove_channel(&uaid, &chid).await.unwrap());
2774        assert!(client.get_user(&uaid).await.unwrap().is_none());
2775
2776        client.remove_user(&uaid).await.unwrap();
2777    }
2778
2779    #[actix_rt::test]
2780    async fn channel_and_current_timestamp_ttl_updates() {
2781        let client = new_client().unwrap();
2782        let uaid = gen_test_uaid();
2783        let chid = Uuid::parse_str(TEST_CHID).unwrap();
2784        client.remove_user(&uaid).await.unwrap();
2785
2786        // Setup a user with some channels and a current_timestamp
2787        let user = User {
2788            uaid,
2789            ..Default::default()
2790        };
2791        client.add_user(&user).await.unwrap();
2792
2793        client.add_channel(&uaid, &chid).await.unwrap();
2794        client
2795            .add_channel(&uaid, &uuid::Uuid::new_v4())
2796            .await
2797            .unwrap();
2798
2799        client
2800            .increment_storage(
2801                &uaid,
2802                SystemTime::now()
2803                    .duration_since(SystemTime::UNIX_EPOCH)
2804                    .unwrap()
2805                    .as_secs(),
2806            )
2807            .await
2808            .unwrap();
2809
2810        let req = client.read_row_request(&uaid.as_simple().to_string());
2811        let Some(mut row) = client.read_row(req).await.unwrap() else {
2812            panic!("Expected row");
2813        };
2814
2815        // Ensure the initial cell expiry (timestamp) of all the cells
2816        // in the row has been updated
2817        let ca_expiry = row.take_required_cell("connected_at").unwrap().timestamp;
2818        for mut cells in row.cells.into_values() {
2819            let Some(cell) = cells.pop() else {
2820                continue;
2821            };
2822            assert!(
2823                cell.timestamp >= ca_expiry,
2824                "{} cell timestamp should >= connected_at's",
2825                cell.qualifier
2826            );
2827        }
2828
2829        let mut user = client.get_user(&uaid).await.unwrap().unwrap();
2830
2831        // Quick nap to make sure that the ca_expiry values are different.
2832        tokio::time::sleep(Duration::from_secs_f32(0.2)).await;
2833        client.update_user(&mut user).await.unwrap();
2834
2835        // Ensure update_user updated the expiry (timestamp) of every cell in the row
2836        let req = client.read_row_request(&uaid.as_simple().to_string());
2837        let Some(mut row) = client.read_row(req).await.unwrap() else {
2838            panic!("Expected row");
2839        };
2840
2841        let ca_expiry2 = row.take_required_cell("connected_at").unwrap().timestamp;
2842
2843        assert!(ca_expiry2 > ca_expiry);
2844
2845        for mut cells in row.cells.into_values() {
2846            let Some(cell) = cells.pop() else {
2847                continue;
2848            };
2849            assert!(
2850                cell.timestamp >= ca_expiry2,
2851                "{} cell timestamp expiry should exceed connected_at's",
2852                cell.qualifier
2853            );
2854        }
2855
2856        client.remove_user(&uaid).await.unwrap();
2857    }
2858
2859    /// A router record whose core `connected_at` cell has expired (its TTL
2860    /// timestamp is in the past) is still served in full: reads don't apply a
2861    /// `max_age` policy the router family doesn't have, so a dormant device
2862    /// keeps working until server-side GC actually reclaims the row. The
2863    /// expiry is counted instead (see `report_expired_cells`).
2864    #[actix_rt::test]
2865    async fn expired_router_record_is_still_served() {
2866        let client = new_client().unwrap();
2867        let uaid = gen_test_uaid();
2868        let fresh_chid = Uuid::parse_str(TEST_CHID).unwrap();
2869        let expired_chid = Uuid::new_v4();
2870        client.remove_user(&uaid).await.unwrap();
2871
2872        // Hand-write a partially-expired record: the core user cells have an
2873        // expiry in the past, while `version` and one `chid:` cell still have
2874        // future expiries. This happens when e.g. a user has not checked in
2875        // recently but notifications were added for them.
2876        let row_key = uaid.simple().to_string();
2877        let past = SystemTime::now() - Duration::from_secs(10);
2878        let future = SystemTime::now() + Duration::from_secs(3600);
2879        let mut row = Row::new(row_key.clone());
2880        let mut cells = vec![
2881            cell::Cell {
2882                qualifier: "connected_at".to_owned(),
2883                value: 0u64.to_be_bytes().to_vec(),
2884                timestamp: past,
2885                ..Default::default()
2886            },
2887            cell::Cell {
2888                qualifier: "router_type".to_owned(),
2889                value: "apns".as_bytes().to_vec(),
2890                timestamp: past,
2891                ..Default::default()
2892            },
2893            cell::Cell {
2894                qualifier: "record_version".to_owned(),
2895                value: USER_RECORD_VERSION.to_be_bytes().to_vec(),
2896                timestamp: past,
2897                ..Default::default()
2898            },
2899            cell::Cell {
2900                qualifier: "version".to_owned(),
2901                value: Uuid::new_v4().into_bytes().to_vec(),
2902                timestamp: future,
2903                ..Default::default()
2904            },
2905        ];
2906        cells.extend(channels_to_cells(
2907            Cow::Owned(HashSet::from([fresh_chid])),
2908            future,
2909        ));
2910        cells.extend(channels_to_cells(
2911            Cow::Owned(HashSet::from([expired_chid])),
2912            past,
2913        ));
2914        row.add_cells(ROUTER_FAMILY, cells);
2915        client.write_row(row).await.unwrap();
2916
2917        let user = client
2918            .get_user(&uaid)
2919            .await
2920            .unwrap()
2921            .expect("an expired record still reads as a user");
2922        assert_eq!(user.router_type, "apns");
2923        // Channels are served whether or not their own cells have expired.
2924        assert_eq!(
2925            user.priv_channels,
2926            HashSet::from([fresh_chid, expired_chid])
2927        );
2928        assert_eq!(
2929            client.get_channels(&uaid).await.unwrap(),
2930            HashSet::from([fresh_chid, expired_chid])
2931        );
2932
2933        // and the row is left alone: reclaiming it is server-side GC's job
2934        let req = client.read_row_request(&row_key);
2935        assert!(client.read_row(req).await.unwrap().is_some());
2936
2937        client.remove_user(&uaid).await.unwrap();
2938    }
2939}