diff --git a/components/ads-client/integration-tests/tests/http_cache.rs b/components/ads-client/integration-tests/tests/http_cache.rs index 8651c07712f..3e646122501 100644 --- a/components/ads-client/integration-tests/tests/http_cache.rs +++ b/components/ads-client/integration-tests/tests/http_cache.rs @@ -6,7 +6,8 @@ use std::hash::{Hash, Hasher}; use std::time::Duration; -use ads_client::http_cache::{ByteSize, CacheOutcome, CachePolicy, HttpCache}; +use ads_client::database::bytesize::ByteSize; +use ads_client::http_cache::{CacheOutcome, CachePolicy, HttpCache}; use mockito::mock; use viaduct::{Client, ClientSettings, Request}; diff --git a/components/ads-client/src/ads_store.rs b/components/ads-client/src/ads_store.rs new file mode 100644 index 00000000000..34c44331ad3 --- /dev/null +++ b/components/ads-client/src/ads_store.rs @@ -0,0 +1,150 @@ +pub mod builder; +pub mod connection_initializer; +pub mod store; + +use crate::{ + ads_store::{builder::AdsStoreBuilder, store::AdsStoreHolder}, + database::bytesize::ByteSize, +}; +use std::path::Path; + +/// Identification of placement sent and returned from MARS (eg: `mock_spoc_1`) +#[derive(Debug, Hash, PartialEq, Eq, Clone)] +pub struct PlacementId(String); + +impl PlacementId { + pub fn new(s: &str) -> PlacementId { + PlacementId(s.to_string()) + } + pub fn into_inner(self) -> String { + self.0 + } +} + +impl AsRef for PlacementId { + fn as_ref(&self) -> &str { + &self.0 + } +} + +pub struct AdsStore { + #[allow(dead_code)] + max_size: ByteSize, + holder: AdsStoreHolder, +} + +impl AdsStore { + pub fn builder>(db_path: P) -> AdsStoreBuilder { + AdsStoreBuilder::new(db_path.as_ref()) + } + + pub fn clear(&self) -> Result<(), rusqlite::Error> { + self.holder.clear_all()?; + Ok(()) + } + + pub fn shutdown_db(self) -> Result<(), rusqlite::Error> { + self.holder.close() + } + + pub fn invalidate_by_id(&self, placement_id: &PlacementId) -> Result<(), rusqlite::Error> { + self.holder.invalidate_ad_by_id(placement_id)?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mars::ad_response::{AdCallbacks, AdImage, StorableAd, StorableAdType}; + use url::Url; + + #[test] + fn test_ads_store_creation() { + // Test that AdsStore can be created successfully with test config + let store: Result = AdsStore::builder("test_store.db").build(); + assert!(store.is_ok()); + } + + #[test] + fn test_clear_store() { + let store: AdsStore = AdsStore::builder("test_clear.db").build().unwrap(); + + // Create a test request and response + let base_url = mockito::server_url(); + let ad = AdImage { + url: "https://ads.fakeexample.org/example_ad_1".to_string(), + image_url: "https://ads.fakeexample.org/example_image_1".to_string(), + format: "billboard".to_string(), + block_key: "abc123".into(), + alt_text: Some("An ad for a puppy".to_string()), + callbacks: AdCallbacks { + click: Url::parse(&format!("{}/click/example_ad_1", base_url)).unwrap(), + impression: Url::parse(&format!("{}/impression/example_ad_1", base_url)).unwrap(), + report: Some(Url::parse(&format!("{}/report/example_ad_1", base_url)).unwrap()), + }, + }; + + let ad = StorableAd { + placement_id: PlacementId::new("mock_billboard_1"), + ad_type: StorableAdType::Image, + ad_body: serde_json::to_vec(&ad).unwrap(), + }; + + store.holder.store_ad(ad.clone()).unwrap(); + + // Verify it's cached + let retrieved = store.holder.lookup(&ad.placement_id).unwrap(); + assert!(retrieved.is_some()); + + // Clear the cache + store.clear().unwrap(); + + // Verify it's cleared + let retrieved_after_clear = store.holder.lookup(&ad.placement_id).unwrap(); + assert!(retrieved_after_clear.is_none()); + } + + #[test] + fn test_invalidate_by_hash() { + let store: AdsStore = AdsStore::builder("test_invalidate.db").build().unwrap(); + + // Create a test request and response + let base_url = mockito::server_url(); + let ad = AdImage { + url: "https://ads.fakeexample.org/example_ad_1".to_string(), + image_url: "https://ads.fakeexample.org/example_image_1".to_string(), + format: "billboard".to_string(), + block_key: "abc123".into(), + alt_text: Some("An ad for a puppy".to_string()), + callbacks: AdCallbacks { + click: Url::parse(&format!("{}/click/example_ad_1", base_url)).unwrap(), + impression: Url::parse(&format!("{}/impression/example_ad_1", base_url)).unwrap(), + report: Some(Url::parse(&format!("{}/report/example_ad_1", base_url)).unwrap()), + }, + }; + + let ad_1 = StorableAd { + placement_id: PlacementId::new("mock_billboard_1"), + ad_type: StorableAdType::Image, + ad_body: serde_json::to_vec(&ad).unwrap(), + }; + let ad_2 = StorableAd { + placement_id: PlacementId::new("mock_billboard_2"), + ad_type: StorableAdType::Image, + ad_body: serde_json::to_vec(&ad).unwrap(), + }; + + store.holder.store_ad(ad_1.clone()).unwrap(); + + store.holder.store_ad(ad_2.clone()).unwrap(); + + assert!(store.holder.lookup(&ad_1.placement_id).unwrap().is_some()); + assert!(store.holder.lookup(&ad_2.placement_id).unwrap().is_some()); + + store.invalidate_by_id(&ad_1.placement_id).unwrap(); + + assert!(store.holder.lookup(&ad_1.placement_id).unwrap().is_none()); + assert!(store.holder.lookup(&ad_2.placement_id).unwrap().is_some()); + } +} diff --git a/components/ads-client/src/ads_store/builder.rs b/components/ads-client/src/ads_store/builder.rs new file mode 100644 index 00000000000..f2f429cc4f1 --- /dev/null +++ b/components/ads-client/src/ads_store/builder.rs @@ -0,0 +1,169 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +use super::connection_initializer::HttpCacheConnectionInitializer; +use crate::ads_store::store::AdsStoreHolder; +use crate::ads_store::AdsStore; +use crate::database::bytesize::ByteSize; +use rusqlite::Connection; +use sql_support::open_database; +use std::path::PathBuf; + +const DEFAULT_MAX_SIZE: ByteSize = ByteSize::mib(10); +const MIN_STORE_SIZE: ByteSize = ByteSize::kib(1); +const MAX_STORE_SIZE: ByteSize = ByteSize::mib(100); + +#[derive(Debug, thiserror::Error)] +pub enum AdsStoreBuilderError { + #[error("Database path cannot be empty")] + EmptyDbPath, + #[error("Database error: {0}")] + Database(#[from] open_database::Error), + #[error( + "Maximum store size must be between {min_size} and {max_size}, got {size_bytes} bytes" + )] + InvalidMaxSize { + max_size: String, + min_size: String, + size_bytes: u64, + }, +} + +pub struct AdsStoreBuilder { + db_path: PathBuf, + max_size: Option, +} + +impl AdsStoreBuilder { + pub fn new(db_path: impl Into) -> Self { + Self { + db_path: db_path.into(), + max_size: None, + } + } + + pub fn max_size(mut self, max_size: ByteSize) -> Self { + self.max_size = Some(max_size); + self + } + + fn open_connection(&self) -> Result { + let initializer = HttpCacheConnectionInitializer {}; + let conn = if cfg!(test) { + open_database::open_memory_database(&initializer)? + } else { + open_database::open_database(&self.db_path, &initializer)? + }; + Ok(conn) + } + + fn validate(&self) -> Result<(), AdsStoreBuilderError> { + if self.db_path.to_string_lossy().trim().is_empty() { + return Err(AdsStoreBuilderError::EmptyDbPath); + } + + if let Some(max_size) = self.max_size { + if max_size < MIN_STORE_SIZE || max_size > MAX_STORE_SIZE { + return Err(AdsStoreBuilderError::InvalidMaxSize { + size_bytes: max_size.as_u64(), + min_size: MIN_STORE_SIZE.to_string(), + max_size: MAX_STORE_SIZE.to_string(), + }); + } + } + + Ok(()) + } + + pub fn build(&self) -> Result { + self.validate()?; + + let conn = self.open_connection()?; + let holder = AdsStoreHolder::new(conn); + let max_size = self.max_size.unwrap_or(DEFAULT_MAX_SIZE); + Ok(AdsStore { max_size, holder }) + } + + #[cfg(test)] + pub fn build_for_time_dependent_tests(&self) -> Result { + self.validate()?; + + let conn = self.open_connection()?; + let max_size = DEFAULT_MAX_SIZE; + let holder = AdsStoreHolder::new_with_test_clock(conn); + + Ok(AdsStore { max_size, holder }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_test_builder(path: &str) -> AdsStoreBuilder { + AdsStoreBuilder::new(path) + } + + #[test] + fn test_store_builder_with_defaults() { + let builder = make_test_builder("test.db"); + assert_eq!(builder.db_path, PathBuf::from("test.db")); + assert_eq!(builder.max_size, None); + assert!(builder.build().is_ok()); + } + + #[test] + fn test_cache_builder_valid_custom() { + let builder = make_test_builder("custom.db").max_size(ByteSize::b(1024)); + + assert_eq!(builder.db_path, PathBuf::from("custom.db")); + assert_eq!(builder.max_size, Some(ByteSize::b(1024))); + assert!(builder.build().is_ok()); + } + + #[test] + fn test_validation_empty_db_path() { + let result = make_test_builder(" ").build(); + assert!(matches!(result, Err(AdsStoreBuilderError::EmptyDbPath))); + } + + #[test] + fn test_validation_max_size_too_small() { + let result = make_test_builder("test.db") + .max_size(ByteSize::b(512)) + .build(); + assert!(matches!( + result, + Err(AdsStoreBuilderError::InvalidMaxSize { + size_bytes: 512, + min_size: _, + max_size: _, + }) + )); + } + + #[test] + fn test_validation_max_size_too_large() { + let result = make_test_builder("test.db") + .max_size(ByteSize::b(2 * 1024 * 1024 * 1024)) + .build(); + assert!(matches!( + result, + Err(AdsStoreBuilderError::InvalidMaxSize { + size_bytes: 2147483648, + min_size: _, + max_size: _, + }) + )); + } + + #[test] + fn test_validation_max_size_boundaries() { + let builder_min = make_test_builder("test.db").max_size(MIN_STORE_SIZE); + assert!(builder_min.build().is_ok()); + + let builder_max = make_test_builder("test.db").max_size(MAX_STORE_SIZE); + assert!(builder_max.build().is_ok()); + } +} diff --git a/components/ads-client/src/ads_store/connection_initializer.rs b/components/ads-client/src/ads_store/connection_initializer.rs new file mode 100644 index 00000000000..14b6d776461 --- /dev/null +++ b/components/ads-client/src/ads_store/connection_initializer.rs @@ -0,0 +1,86 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +use rusqlite::Connection; +use sql_support::open_database; +use std::time::Duration; + +pub struct HttpCacheConnectionInitializer {} + +impl open_database::ConnectionInitializer for HttpCacheConnectionInitializer { + const NAME: &'static str = "ads_cache"; + const END_VERSION: u32 = 1; + + fn prepare(&self, conn: &Connection, _db_empty: bool) -> open_database::Result<()> { + conn.execute_batch("PRAGMA journal_mode=wal;")?; + conn.busy_timeout(Duration::from_secs(5))?; + Ok(()) + } + + fn init(&self, tx: &rusqlite::Transaction<'_>) -> open_database::Result<()> { + const SCHEMA: &str = " + CREATE TABLE IF NOT EXISTS ads ( + stored_at INTEGER NOT NULL, + placement_id TEXT NOT NULL, + ad_type SMALLINT NOT NULL, + ad_body BLOB NOT NULL, + size_bytes INTEGER NOT NULL, + PRIMARY KEY (placement_id) + ); + CREATE INDEX IF NOT EXISTS idx_ads_stored_at ON ads(stored_at); + CREATE INDEX IF NOT EXISTS idx_ads_placement_id ON ads(placement_id); + "; + // If the schema fails to initialize, it might be corrupted or outdated so we drop the table and try again + if tx.execute_batch(SCHEMA).is_err() { + tx.execute_batch("DROP TABLE IF EXISTS ads")?; + tx.execute_batch(SCHEMA)?; + } + Ok(()) + } + + fn upgrade_from( + &self, + conn: &rusqlite::Transaction<'_>, + version: u32, + ) -> open_database::Result<()> { + match version { + 0 => self.init(conn), + _ => Err(open_database::Error::IncompatibleVersion(version)), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rusqlite::Connection; + use sql_support::open_database::ConnectionInitializer; + + #[test] + fn test_corrupted_schema_is_recreated() { + let mut conn = Connection::open_in_memory().unwrap(); + let initializer = HttpCacheConnectionInitializer {}; + + // Create a corrupted table with only one column + conn.execute_batch("CREATE TABLE ads (placement_id TEXT);") + .unwrap(); + + // Run init - should drop the corrupted table and recreate it properly + let tx = conn.transaction().unwrap(); + initializer.init(&tx).unwrap(); + tx.commit().unwrap(); + + // Verify the table was recreated with correct schema by checking column count + let column_count: i64 = conn + .query_row("SELECT COUNT(*) FROM pragma_table_info('ads')", [], |row| { + row.get(0) + }) + .unwrap(); + + assert!( + column_count > 1, + "Table should have more than 1 column after recreation" + ); + } +} diff --git a/components/ads-client/src/ads_store/store.rs b/components/ads-client/src/ads_store/store.rs new file mode 100644 index 00000000000..657defaf8b2 --- /dev/null +++ b/components/ads-client/src/ads_store/store.rs @@ -0,0 +1,359 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +use crate::database::bytesize::ByteSize; +use crate::database::clock::Clock; +use crate::{ + ads_store::PlacementId, + database::clock::CacheClock, + mars::ad_response::{StorableAd, StorableAdType}, +}; +use parking_lot::Mutex; +use rusqlite::{params, Connection, OptionalExtension, Result as SqliteResult}; +use std::sync::Arc; + +#[cfg(test)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum FaultKind { + None, + Lookup, + Store, + Trim, +} + +pub struct AdsStoreHolder { + conn: Mutex, + clock: Arc, + #[cfg(test)] + fault: parking_lot::Mutex, +} + +impl AdsStoreHolder { + pub fn new(conn: Connection) -> Self { + Self { + conn: Mutex::new(conn), + clock: Arc::new(CacheClock), + #[cfg(test)] + fault: parking_lot::Mutex::new(FaultKind::None), + } + } + + pub fn close(self) -> Result<(), rusqlite::Error> { + let conn = self.conn.into_inner(); + conn.close().map_err(|(_, err)| err) + } + + #[cfg(test)] + pub fn new_with_test_clock(conn: Connection) -> Self { + use crate::database::clock::TestClock; + + Self { + conn: Mutex::new(conn), + clock: Arc::new(TestClock::new(chrono::Utc::now().timestamp())), + #[cfg(test)] + fault: parking_lot::Mutex::new(FaultKind::None), + } + } + + #[cfg(test)] + pub fn get_clock(&self) -> &dyn Clock { + &*self.clock + } + + /// Removes all entries from cache. + pub fn clear_all(&self) -> SqliteResult { + let conn = self.conn.lock(); + let mut total = 0; + total += conn.execute("DELETE FROM ads", [])?; + Ok(total) + } + + /// Returns total size of the cache in bytes. + pub fn current_total_size_bytes(&self) -> SqliteResult { + let conn = self.conn.lock(); + let size_bytes_ads: u64 = + conn.query_row("SELECT COALESCE(SUM(size_bytes),0) FROM ads", [], |row| { + row.get(0) + })?; + Ok(ByteSize::b(size_bytes_ads)) + } + + pub fn lookup(&self, placement_id: &PlacementId) -> SqliteResult> { + #[cfg(test)] + if *self.fault.lock() == FaultKind::Lookup { + return Err(Self::forced_fault_error("forced lookup failure")); + } + let conn = self.conn.lock(); + conn.query_row( + "SELECT placement_id, ad_type, ad_body + FROM ads WHERE placement_id = ?1", + params![placement_id.as_ref()], + |row| { + let placement_id: String = row.get(0)?; + let ad_type: u8 = row.get(1)?; + let ad_body: Vec = row.get(2)?; + + let placement_id = PlacementId::new(&placement_id); + let ad_type = StorableAdType::try_from(ad_type).map_err(|e| { + rusqlite::Error::FromSqlConversionFailure( + 1, + rusqlite::types::Type::Integer, + e.into(), + ) + })?; + + Ok(StorableAd { + placement_id, + ad_type, + ad_body, + }) + }, + ) + .optional() + } + + /// Upsert an object into the store. + pub fn store_ad(&self, ad: StorableAd) -> SqliteResult<()> { + #[cfg(test)] + if *self.fault.lock() == FaultKind::Store { + return Err(Self::forced_fault_error("forced store failure")); + } + let placement_id_str: &str = ad.placement_id.as_ref(); + let size_bytes = ad.ad_body.len() as i64; + let now = self.clock.now_epoch_seconds(); + + let conn = self.conn.lock(); + conn.execute( + "INSERT INTO ads ( + stored_at, + placement_id, + ad_type, + ad_body, + size_bytes + ) + VALUES (?1, ?2, ?3, ?4, ?5) + ON CONFLICT(placement_id) DO UPDATE SET + stored_at=excluded.stored_at, + ad_type=excluded.ad_type, + ad_body=excluded.ad_body, + size_bytes=excluded.size_bytes", + params![ + now, + placement_id_str, + ad.ad_type.to_u8(), + ad.ad_body, + size_bytes, + ], + )?; + Ok(()) + } + + pub fn invalidate_ad_by_id(&self, placement_id: &PlacementId) -> SqliteResult { + let conn = self.conn.lock(); + conn.execute( + "DELETE FROM ads WHERE placement_id = ?1", + params![&placement_id.as_ref()], + ) + } + + pub fn trim_to_max_size(&self, max_size: &ByteSize) -> SqliteResult<()> { + #[cfg(test)] + if *self.fault.lock() == FaultKind::Trim { + return Err(Self::forced_fault_error("forced trim failure")); + } + loop { + let total = self.current_total_size_bytes()?; + if total.as_u64() <= max_size.as_u64() { + break; + } + let conn = self.conn.lock(); + conn.execute( + "DELETE FROM ads WHERE rowid IN ( + SELECT rowid FROM ads ORDER BY stored_at ASC LIMIT 1 + )", + [], + )?; + } + Ok(()) + } + + #[cfg(test)] + pub fn set_fault(&self, kind: FaultKind) { + *self.fault.lock() = kind; + } + + #[cfg(test)] + fn forced_fault_error(msg: &str) -> rusqlite::Error { + rusqlite::Error::SqliteFailure( + rusqlite::ffi::Error { + code: rusqlite::ErrorCode::InternalMalfunction, + extended_code: 0, + }, + Some(msg.to_string()), + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + ads_store::connection_initializer::HttpCacheConnectionInitializer, + mars::ad_response::{AdCallbacks, AdImage}, + }; + use sql_support::open_database; + use url::Url; + + // Create a sample ad for tests. The body defaults to an example serialized AdImage (if body is None). + fn create_test_raw_ad(placement_id: &str, body: Option>) -> StorableAd { + let base_url = mockito::server_url(); + let ad = AdImage { + url: "https://ads.fakeexample.org/example_ad_1".to_string(), + image_url: "https://ads.fakeexample.org/example_image_1".to_string(), + format: "billboard".to_string(), + block_key: "abc123".into(), + alt_text: Some("An ad for a puppy".to_string()), + callbacks: AdCallbacks { + click: Url::parse(&format!("{}/click/example_ad_1", base_url)).unwrap(), + impression: Url::parse(&format!("{}/impression/example_ad_1", base_url)).unwrap(), + report: Some(Url::parse(&format!("{}/report/example_ad_1", base_url)).unwrap()), + }, + }; + StorableAd { + placement_id: PlacementId::new(placement_id), + ad_type: StorableAdType::Image, + ad_body: body.unwrap_or(serde_json::to_vec(&ad).unwrap()), + } + } + + fn create_test_store() -> AdsStoreHolder { + let initializer = HttpCacheConnectionInitializer {}; + let conn = open_database::open_memory_database(&initializer) + .expect("failed to open memory cache db"); + AdsStoreHolder::new_with_test_clock(conn) + } + + #[test] + fn test_lookup_fault_injection() { + let store = create_test_store(); + store.set_fault(FaultKind::Lookup); + + let ad = create_test_raw_ad("mock_billboard_1", None); + let err = store.lookup(&ad.placement_id).unwrap_err(); + + match err { + rusqlite::Error::SqliteFailure(_, Some(msg)) => { + assert!(msg.contains("forced lookup failure")); + } + other => panic!("unexpected error: {other:?}"), + } + } + + #[test] + fn test_store_fault_injection() { + let store = create_test_store(); + store.set_fault(FaultKind::Store); + + let ad = create_test_raw_ad("mock_billboard_1", None); + + let err = store.store_ad(ad).unwrap_err(); + match err { + rusqlite::Error::SqliteFailure(_, Some(msg)) => { + assert!(msg.contains("forced store failure")); + } + other => panic!("unexpected error: {other:?}"), + } + } + + #[test] + fn test_trim_fault_injection() { + let store = create_test_store(); + store.set_fault(FaultKind::Trim); + + let ad = create_test_raw_ad("mock_billboard_1", None); + store.store_ad(ad).unwrap(); + + let err = store.trim_to_max_size(&ByteSize::b(1)).unwrap_err(); + match err { + rusqlite::Error::SqliteFailure(_, Some(msg)) => { + assert!(msg.contains("forced trim failure")); + } + other => panic!("unexpected error: {other:?}"), + } + } + + #[test] + fn test_store_and_retrieve_ads() { + let store = create_test_store(); + let ad = create_test_raw_ad("mock_billboard_1", None); + + store.store_ad(ad.clone()).unwrap(); + + let retrieved = store.lookup(&ad.placement_id).unwrap().unwrap(); + assert_eq!(retrieved.ad_body, ad.ad_body); + } + + #[test] + fn test_max_size_eviction_ads() { + let initializer = HttpCacheConnectionInitializer {}; + let conn = open_database::open_memory_database(&initializer) + .expect("failed to open memory cache db"); + let store = AdsStoreHolder::new(conn); + + for i in 0..5 { + let large_body = vec![0u8; 300]; + let ad = create_test_raw_ad(&format!("mock_billboard_{i}"), Some(large_body)); + store.store_ad(ad.clone()).unwrap(); + } + + store.trim_to_max_size(&ByteSize::kib(1)).unwrap(); + + let total_size = store.current_total_size_bytes().unwrap(); + assert!(total_size.as_u64() <= 1024); + + let first_placement_id = PlacementId::new("mock_billboard_0"); + let first_cached = store.lookup(&first_placement_id).unwrap(); + assert!(first_cached.is_none()); + } + + #[test] + fn test_clear_all_ads() { + let store = create_test_store(); + let ad_1 = create_test_raw_ad("mock_billboard_1", None); + + store.store_ad(ad_1.clone()).unwrap(); + + let ad_2 = create_test_raw_ad("mock_billboard_2", None); + store.store_ad(ad_2.clone()).unwrap(); + + assert!(store.lookup(&ad_1.placement_id).unwrap().is_some()); + assert!(store.lookup(&ad_2.placement_id).unwrap().is_some()); + + let deleted_count = store.clear_all().unwrap(); + assert_eq!(deleted_count, 2); + + assert!(store.lookup(&ad_1.placement_id).unwrap().is_none()); + assert!(store.lookup(&ad_2.placement_id).unwrap().is_none()); + } + + #[test] + fn test_invalidate_ad_by_placement_id() { + let store = create_test_store(); + + let ad_1 = create_test_raw_ad("mock_billboard_1", None); + let ad_2 = create_test_raw_ad("mock_billboard_2", None); + + store.store_ad(ad_1.clone()).unwrap(); + store.store_ad(ad_2.clone()).unwrap(); + + assert!(store.lookup(&ad_1.placement_id).unwrap().is_some()); + assert!(store.lookup(&ad_2.placement_id).unwrap().is_some()); + + let deleted = store.invalidate_ad_by_id(&ad_1.placement_id).unwrap(); + assert_eq!(deleted, 1); + + assert!(store.lookup(&ad_1.placement_id).unwrap().is_none()); + assert!(store.lookup(&ad_2.placement_id).unwrap().is_some()); + } +} diff --git a/components/ads-client/src/client.rs b/components/ads-client/src/client.rs index 1089948649b..a80519f3486 100644 --- a/components/ads-client/src/client.rs +++ b/components/ads-client/src/client.rs @@ -6,7 +6,9 @@ use std::collections::HashMap; use std::time::Duration; -use crate::http_cache::{ByteSize, CachePolicy, HttpCache}; +use crate::ads_store::AdsStore; +use crate::database::bytesize::ByteSize; +use crate::http_cache::{CachePolicy, HttpCache}; use crate::mars::ad_request::{AdPlacementRequest, AdRequestFlags}; use crate::mars::ad_response::{AdImage, AdResponse, AdResponseValue, AdSpoc, AdTile}; use crate::mars::error::{RecordClickError, RecordImpressionError, ReportAdError}; @@ -42,6 +44,7 @@ where client: MARSClient, context_id_provider: Box, telemetry: T, + ads_store: Option, } impl AdsClient @@ -85,12 +88,24 @@ where } }); + let ads_store = + client_config + .store_config + .and_then(|x| match AdsStore::builder(x.db_path).build() { + Ok(store) => Some(store), + Err(e) => { + telemetry.record(&e); + None + } + }); + let client = MARSClient::new(environment, http_cache, telemetry.clone()); telemetry.record(&ClientOperationEvent::New); Self { client, context_id_provider, telemetry: telemetry.clone(), + ads_store, } } @@ -105,8 +120,13 @@ where self.telemetry.shutdown(); // Shutdown DB - self.client.shutdown_db()?; + let r = self.client.shutdown_db(); + + if let Some(ads_store) = self.ads_store.take() { + ads_store.shutdown_db()?; + } + r?; Ok(()) } @@ -278,6 +298,7 @@ mod tests { use std::{assert_eq, assert_ne, sync::Arc}; use crate::{ + ads_store::builder::AdsStoreBuilder, ffi::telemetry::MozAdsTelemetryWrapper, mars::Environment, test_utils::{ @@ -301,6 +322,11 @@ mod tests { Box::new(DefaultContextIdCallback), )), telemetry, + ads_store: Some( + AdsStoreBuilder::new("test_store.db") + .build() + .expect("Simplest AdsStoreBuilder should be constructable"), + ), } } @@ -311,6 +337,7 @@ mod tests { context_id_provider: None, environment: Environment::Test, telemetry: MozAdsTelemetryWrapper::noop(), + store_config: None, }; let client = AdsClient::new(config); let context_id = client.get_context_id().unwrap(); @@ -415,6 +442,7 @@ mod tests { context_id_provider: Some(Box::new(FixedContextId)), environment: Environment::Test, telemetry: MozAdsTelemetryWrapper::noop(), + store_config: None, }; let client = AdsClient::new(config); @@ -534,6 +562,7 @@ mod tests { context_id_provider: None, environment: Environment::Test, telemetry: noop_telemetry, + store_config: None, }; let mut client = AdsClient::new(config); diff --git a/components/ads-client/src/client/config.rs b/components/ads-client/src/client/config.rs index 7c86c241418..2e96c34fa63 100644 --- a/components/ads-client/src/client/config.rs +++ b/components/ads-client/src/client/config.rs @@ -11,6 +11,7 @@ where T: Telemetry, { pub cache_config: Option, + pub store_config: Option, pub context_id_provider: Option>, pub environment: Environment, pub telemetry: T, @@ -22,3 +23,8 @@ pub struct AdsCacheConfig { pub default_cache_ttl_seconds: Option, pub max_size_mib: Option, } + +#[derive(Clone, Debug)] +pub struct AdsStoreConfig { + pub db_path: String, +} diff --git a/components/ads-client/src/database.rs b/components/ads-client/src/database.rs new file mode 100644 index 00000000000..9bf27276888 --- /dev/null +++ b/components/ads-client/src/database.rs @@ -0,0 +1,2 @@ +pub mod bytesize; +pub mod clock; diff --git a/components/ads-client/src/http_cache/bytesize.rs b/components/ads-client/src/database/bytesize.rs similarity index 100% rename from components/ads-client/src/http_cache/bytesize.rs rename to components/ads-client/src/database/bytesize.rs diff --git a/components/ads-client/src/http_cache/clock.rs b/components/ads-client/src/database/clock.rs similarity index 100% rename from components/ads-client/src/http_cache/clock.rs rename to components/ads-client/src/database/clock.rs diff --git a/components/ads-client/src/ffi.rs b/components/ads-client/src/ffi.rs index ba50352fdc1..e7435746c2e 100644 --- a/components/ads-client/src/ffi.rs +++ b/components/ads-client/src/ffi.rs @@ -8,7 +8,7 @@ pub mod telemetry; use std::sync::Arc; -use crate::client::config::{AdsCacheConfig, AdsClientConfig}; +use crate::client::config::{AdsCacheConfig, AdsClientConfig, AdsStoreConfig}; use crate::client::{AdsClient, ContextIdProvider}; use crate::ffi::telemetry::MozAdsTelemetryWrapper; use crate::http_cache::CachePolicy; @@ -107,6 +107,7 @@ struct MozAdsClientBuilderInner { context_id_provider: Option>, environment: Option, telemetry: Option>, + store_config: Option, } impl Default for MozAdsClientBuilder { @@ -137,6 +138,7 @@ impl MozAdsClientBuilder { .clone() .map(MozAdsTelemetryWrapper::new) .unwrap_or_else(MozAdsTelemetryWrapper::noop), + store_config: inner.store_config.clone().map(Into::into), }; let client = AdsClient::new(client_config); MozAdsClient { @@ -149,6 +151,11 @@ impl MozAdsClientBuilder { self } + pub fn store_config(self: Arc, store_config: MozAdsStoreConfig) -> Arc { + self.0.lock().store_config = Some(store_config); + self + } + pub fn context_id_provider( self: Arc, provider: Arc, @@ -186,6 +193,11 @@ pub struct MozAdsCacheConfig { pub max_size_mib: Option, } +#[derive(Clone, uniffi::Record)] +pub struct MozAdsStoreConfig { + pub db_path: String, +} + #[derive(Debug, PartialEq, uniffi::Record)] pub struct MozAdsContentCategory { pub categories: Vec, @@ -452,6 +464,14 @@ impl From for AdsCacheConfig { } } +impl From for AdsStoreConfig { + fn from(config: MozAdsStoreConfig) -> Self { + Self { + db_path: config.db_path, + } + } +} + impl From<&MozAdsPlacementRequest> for AdPlacementRequest { fn from(request: &MozAdsPlacementRequest) -> Self { Self { diff --git a/components/ads-client/src/ffi/telemetry.rs b/components/ads-client/src/ffi/telemetry.rs index 02a6fee2e46..63bbd9eaad1 100644 --- a/components/ads-client/src/ffi/telemetry.rs +++ b/components/ads-client/src/ffi/telemetry.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use parking_lot::RwLock; +use crate::ads_store::builder::AdsStoreBuilderError; use crate::client::error::RequestAdsError; use crate::client::ClientOperationEvent; use crate::http_cache::{CacheOutcome, HttpCacheBuilderError}; @@ -101,6 +102,19 @@ impl Telemetry for MozAdsTelemetryWrapper { }); return; } + if let Some(cache_builder_error) = event.downcast_ref::() { + inner.record_build_cache_error( + match cache_builder_error { + AdsStoreBuilderError::EmptyDbPath => "store_empty_db_path".to_string(), + AdsStoreBuilderError::Database(_) => "store_database_error".to_string(), + AdsStoreBuilderError::InvalidMaxSize { .. } => { + "store_invalid_max_size".to_string() + } + }, + format!("{}", cache_builder_error), + ); + return; + } if let Some(cache_builder_error) = event.downcast_ref::() { inner.record_build_cache_error( match cache_builder_error { diff --git a/components/ads-client/src/http_cache.rs b/components/ads-client/src/http_cache.rs index 81b056bb1df..e09da1c828d 100644 --- a/components/ads-client/src/http_cache.rs +++ b/components/ads-client/src/http_cache.rs @@ -3,9 +3,7 @@ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ mod builder; -mod bytesize; mod cache_control; -mod clock; mod connection_initializer; mod outcome; mod request_hash; @@ -18,12 +16,12 @@ use self::{ store::HttpCacheStore, strategy::{CacheFirst, NetworkFirst}, }; +use crate::database::bytesize::ByteSize; use std::hash::Hash; use viaduct::{Client, Request, Response}; pub use self::builder::HttpCacheBuilderError; -pub use self::bytesize::ByteSize; pub use self::outcome::CacheOutcome; pub use self::request_hash::RequestHash; use std::path::Path; diff --git a/components/ads-client/src/http_cache/builder.rs b/components/ads-client/src/http_cache/builder.rs index 97f92f004b0..5af4a610616 100644 --- a/components/ads-client/src/http_cache/builder.rs +++ b/components/ads-client/src/http_cache/builder.rs @@ -2,11 +2,10 @@ * License, v. 2.0. If a copy of the MPL was not distributed with this * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ -use crate::http_cache::HttpCache; - -use super::bytesize::ByteSize; use super::connection_initializer::HttpCacheConnectionInitializer; use super::store::HttpCacheStore; +use crate::database::bytesize::ByteSize; +use crate::http_cache::HttpCache; use rusqlite::Connection; use sql_support::open_database; use std::path::PathBuf; diff --git a/components/ads-client/src/http_cache/store.rs b/components/ads-client/src/http_cache/store.rs index a80f5a956ca..ae1bcf36bb4 100644 --- a/components/ads-client/src/http_cache/store.rs +++ b/components/ads-client/src/http_cache/store.rs @@ -4,10 +4,9 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; -use crate::http_cache::{ - clock::{CacheClock, Clock}, - request_hash::RequestHash, - ByteSize, +use crate::{ + database::clock::{CacheClock, Clock}, + http_cache::{request_hash::RequestHash, ByteSize}, }; use parking_lot::Mutex; use rusqlite::{params, Connection, OptionalExtension, Result as SqliteResult}; @@ -47,7 +46,7 @@ impl HttpCacheStore { #[cfg(test)] pub fn new_with_test_clock(conn: Connection) -> Self { - use crate::http_cache::clock::TestClock; + use crate::database::clock::TestClock; Self { conn: Mutex::new(conn), diff --git a/components/ads-client/src/lib.rs b/components/ads-client/src/lib.rs index 87cecb2a5f8..848975fa9fe 100644 --- a/components/ads-client/src/lib.rs +++ b/components/ads-client/src/lib.rs @@ -15,7 +15,9 @@ use client::AdsClient; use error_support::error; use http_cache::CachePolicy; use mars::ad_request::{AdPlacementRequest, AdRequestFlags}; +pub mod ads_store; mod client; +pub mod database; mod ffi; pub mod http_cache; mod mars; diff --git a/components/ads-client/src/mars.rs b/components/ads-client/src/mars.rs index 4519e0bf7e7..f26cba6de0e 100644 --- a/components/ads-client/src/mars.rs +++ b/components/ads-client/src/mars.rs @@ -259,7 +259,7 @@ mod tests { let cache = HttpCache::builder("test_fetch_ads_cache_hit_skips_network.db") .default_ttl(std::time::Duration::from_secs(300)) - .max_size(crate::http_cache::ByteSize::mib(1)) + .max_size(crate::database::bytesize::ByteSize::mib(1)) .build() .unwrap(); let client = make_test_client(Some(cache)); @@ -295,7 +295,7 @@ mod tests { viaduct_dev::init_backend_dev(); let cache = HttpCache::builder("test_record_click.db") .default_ttl(std::time::Duration::from_secs(300)) - .max_size(crate::http_cache::ByteSize::mib(1)) + .max_size(crate::database::bytesize::ByteSize::mib(1)) .build() .unwrap(); @@ -314,7 +314,7 @@ mod tests { viaduct_dev::init_backend_dev(); let cache = HttpCache::builder("test_record_impression.db") .default_ttl(std::time::Duration::from_secs(300)) - .max_size(crate::http_cache::ByteSize::mib(1)) + .max_size(crate::database::bytesize::ByteSize::mib(1)) .build() .unwrap(); diff --git a/components/ads-client/src/mars/ad_response.rs b/components/ads-client/src/mars/ad_response.rs index 5e057b7606b..8b5196e139a 100644 --- a/components/ads-client/src/mars/ad_response.rs +++ b/components/ads-client/src/mars/ad_response.rs @@ -3,6 +3,7 @@ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ +use crate::ads_store::PlacementId; use crate::http_cache::RequestHash; use crate::telemetry::Telemetry; use serde::de::DeserializeOwned; @@ -141,6 +142,42 @@ pub struct AdTile { pub url: String, } +#[derive(Debug, Clone)] +pub struct StorableAd { + pub placement_id: PlacementId, + pub ad_type: StorableAdType, + pub ad_body: Vec, +} + +#[derive(Clone, Copy, Debug)] +pub enum StorableAdType { + Image, + Spoc, + Tile, +} + +impl StorableAdType { + pub fn to_u8(self) -> u8 { + match self { + StorableAdType::Image => 0, + StorableAdType::Spoc => 1, + StorableAdType::Tile => 2, + } + } +} + +impl TryFrom for StorableAdType { + type Error = String; + fn try_from(value: u8) -> Result { + Ok(match value { + 0 => StorableAdType::Image, + 1 => StorableAdType::Spoc, + 2 => StorableAdType::Tile, + x => return Err(format!("Invalid variant for StorableAdType: {x}")), + }) + } +} + #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] pub struct SpocFrequencyCaps { pub cap_key: String,