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
47pub type RowKey = String;
49
50pub type Qualifier = String;
55pub type FamilyId = String;
56
57const ROUTER_FAMILY: &str = "router";
58const MESSAGE_FAMILY: &str = "message"; const MESSAGE_TOPIC_FAMILY: &str = "message_topic";
60#[cfg(feature = "reliable_report")]
61const RELIABLE_LOG_FAMILY: &str = "reliability";
62#[cfg(feature = "reliable_report")]
63pub const RELIABLE_LOG_TTL: TimeDelta = TimeDelta::days(60);
66
67pub(crate) const RETRY_COUNT: usize = 2;
74
75const MAX_MESSAGE_LEN: usize = 1 << 28;
77
78const BIGTABLE_DATA_SCOPES: &[&str] = &["https://www.googleapis.com/auth/bigtable.data"];
80
81#[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 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 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 true
119 } else {
120 false
121 }
122 }
123
124 pub fn record_success(&self) {
126 self.consecutive_failures.store(0, Ordering::Relaxed);
127 }
128
129 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 Self::new(5, 30)
146 }
147}
148
149struct 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)]
166pub struct BigTableClientImpl {
168 pub(crate) settings: BigTableDbSettings,
169 metrics: Arc<StatsdClient>,
171 pool: BigTablePool,
173 metadata: MetadataMap,
174 circuit_breaker: Arc<CircuitBreaker>,
176}
177
178fn max_versions_filter() -> bigtable::RowFilter {
182 bigtable::RowFilter {
183 filter: Some(bigtable::row_filter::Filter::CellsPerColumnLimitFilter(1)),
184 }
185}
186
187fn 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
207fn router_gc_policy_filter() -> bigtable::RowFilter {
215 max_versions_filter()
216}
217
218fn message_gc_policy_filter() -> Result<Vec<bigtable::RowFilter>, error::BigTableError> {
221 Ok(vec![max_versions_filter(), expiry_filter()?])
222}
223
224fn family_filter(regex: String) -> bigtable::RowFilter {
226 bigtable::RowFilter {
227 filter: Some(bigtable::row_filter::Filter::FamilyNameRegexFilter(regex)),
228 }
229}
230
231fn 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 vec.extend("\\x00".as_bytes());
244 continue;
245 }
246 vec.push(b'\\');
247 }
248 vec.push(b);
249 }
250 vec
251}
252
253fn 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
274fn 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
284fn 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
293fn 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
324fn 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
344fn 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 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
444fn retryable_read_err(status: &Status) -> bool {
461 if retryable_transient_err(status) {
462 return true;
463 }
464 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 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 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 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 let (socket, _) = listener.accept().await.unwrap();
686 let service = service_fn(|_request: HyperRequest<Incoming>| async move {
687 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 let (socket, _) = listener.accept().await.unwrap();
721 let service = service_fn(|_request: HyperRequest<Incoming>| async move {
722 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
934fn 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 error::BigTableError::PreSendTimeout => Some(RetryMetric {
962 error: "PreSendTimeout",
963 code: None,
964 }),
965 error::BigTableError::Auth(_) => Some(RetryMetric {
968 error: "Auth",
969 code: None,
970 }),
971 _ => None,
972 }
973}
974
975fn 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
994fn 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
1011fn 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 #[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 #[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 let odd: HashMap<_, _> = [("app_id".to_owned(), json!(7))].into();
1065 assert_eq!(app_id_tag(Some(&odd)), "none");
1066 }
1067}
1068
1069impl 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 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 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 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 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 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 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 if let Some(retry_metric) = take_pending_retry_metric(&attempt_retry_metric) {
1215 retry_metric.send(&metrics);
1216 }
1217 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 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 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 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 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 async fn write_row(&self, row: row::Row) -> Result<(), error::BigTableError> {
1384 let mut req = self.mutate_row_request(&row.row_key);
1385 req.mutations = self.get_mutations(row.cells)?;
1390 self.mutate_row(req).await?;
1391 Ok(())
1392 }
1393
1394 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 timestamp_micros: (timestamp.as_millis() * 1000) as i64,
1416 value: cell.value,
1417 },
1418 )),
1419 });
1420 }
1421 }
1422 Ok(mutations)
1423 }
1424
1425 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 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 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 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 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: {:?}", ¬if);
1575 Ok(notif)
1576 }
1577
1578 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 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 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 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 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 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 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 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 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 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 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 async fn add_channels(&self, uaid: &Uuid, channels: HashSet<Uuid>) -> DbResult<()> {
1877 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 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 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 let column = format!("chid:{}", channel_id.as_hyphenated());
1934 let mut mutations = self.get_delete_mutations(ROUTER_FAMILY, &[column.as_ref()], None)?;
1935
1936 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 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 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 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 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 async fn save_messages(&self, uaid: &Uuid, messages: Vec<Notification>) -> DbResult<()> {
2085 for message in messages {
2087 self.save_message(uaid, message).await?;
2088 }
2089 Ok(())
2090 }
2091
2092 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 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 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 Ok(FetchMessageResponse {
2196 messages,
2197 timestamp: None,
2198 })
2199 }
2200
2201 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 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 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 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 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 self.read_rows_with_policy(req, RpcClass::Point, BreakerPolicy::Ignore)
2286 .await?;
2287 Ok(true)
2288 }
2289
2290 async fn router_table_exists(&self) -> DbResult<bool> {
2293 Ok(true)
2294 }
2295
2296 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 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 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 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 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 #[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 #[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 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 let _ = client.remove_user(&uaid).await;
2475
2476 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 let connected_at = ms_since_epoch();
2485
2486 client.add_channel(&uaid, &chid).await?;
2488 let channels = client.get_channels(&uaid).await?;
2489 assert!(channels.contains(&chid));
2490
2491 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 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 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 let fetched2 = client.get_user(&fetched.uaid).await?.unwrap();
2523 assert_eq!(fetched.connected_at, fetched2.connected_at);
2524
2525 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 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 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 let fetched = client
2573 .fetch_timestamp_messages(&uaid, Some(timestamp - 10), 999)
2574 .await?;
2575 assert_ne!(fetched.messages.len(), 0);
2576
2577 let fetched = client
2579 .fetch_timestamp_messages(&uaid, Some(timestamp + 10), 999)
2580 .await?;
2581 assert_eq!(fetched.messages.len(), 0);
2582
2583 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 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 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 let fetched = client.fetch_topic_messages(&uaid, 999).await?;
2624 assert_ne!(fetched.messages.len(), 0);
2625
2626 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 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 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 client.add_channel(&uaid, &chid).await.unwrap();
2738
2739 assert!(client.get_user(&uaid).await.unwrap().is_none());
2742
2743 client.add_user(&user).await.unwrap();
2744 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 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 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 tokio::time::sleep(Duration::from_secs_f32(0.2)).await;
2833 client.update_user(&mut user).await.unwrap();
2834
2835 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 #[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 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 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 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}