diff --git a/components/ads-client/Cargo.toml b/components/ads-client/Cargo.toml index 15eac5443a..15062fc915 100644 --- a/components/ads-client/Cargo.toml +++ b/components/ads-client/Cargo.toml @@ -15,6 +15,7 @@ context_id = { path = "../context_id" } error-support = { path = "../support/error" } parking_lot = "0.12" rusqlite = { version = "0.37.0", features = [ + "array", "functions", "bundled", "serde_json", diff --git a/components/ads-client/android/src/main/java/mozilla/appservices/adsclient/AdsClientTelemetry.kt b/components/ads-client/android/src/main/java/mozilla/appservices/adsclient/AdsClientTelemetry.kt index 856cdd4920..46a8ca730b 100644 --- a/components/ads-client/android/src/main/java/mozilla/appservices/adsclient/AdsClientTelemetry.kt +++ b/components/ads-client/android/src/main/java/mozilla/appservices/adsclient/AdsClientTelemetry.kt @@ -31,4 +31,12 @@ class AdsClientTelemetry : MozAdsTelemetry { override fun recordHttpCacheOutcome(label: String, value: String) { AdsClient.httpCacheOutcome[label].set(value) } + + override fun recordBuildImpressionLogError(label: String, value: String) { + AdsClient.impressionLogError[label].set(value) + } + + override fun recordImpressionLogOutcome(label: String, value: String) { + AdsClient.impressionLogOutcome[label].set(value) + } } diff --git a/components/ads-client/src/client.rs b/components/ads-client/src/client.rs index 5c818e2a4c..2090d76f7e 100644 --- a/components/ads-client/src/client.rs +++ b/components/ads-client/src/client.rs @@ -7,8 +7,11 @@ use std::collections::HashMap; use std::time::Duration; use crate::http_cache::{ByteSize, CachePolicy, HttpCache}; +use crate::impression_log::{ImpressionCappingPolicy, ImpressionLog}; use crate::mars::ad_request::{AdPlacementRequest, AdRequestFlags}; -use crate::mars::ad_response::{AdImage, AdResponse, AdResponseValue, AdSpoc, AdTile}; +use crate::mars::ad_response::{ + pop_query_param_from_url, AdImage, AdResponse, AdResponseValue, AdSpoc, AdTile, +}; use crate::mars::error::{RecordClickError, RecordImpressionError, ReportAdError}; use crate::mars::{MARSClient, ReportReason}; use crate::telemetry::Telemetry; @@ -85,7 +88,20 @@ where } }); - let client = MARSClient::new(environment, http_cache, telemetry.clone()); + let impression_log = + client_config + .impression_log_config + .and_then(|impression_log_config| { + match ImpressionLog::builder(impression_log_config.db_path).build() { + Ok(cache) => Some(cache), + Err(e) => { + telemetry.record(&e); + None + } + } + }); + + let client = MARSClient::new(environment, http_cache, impression_log, telemetry.clone()); telemetry.record(&ClientOperationEvent::New); Self { client, @@ -121,7 +137,7 @@ where pub fn record_impression( &self, - impression_url: Url, + mut impression_url: Url, ohttp: bool, ) -> Result<(), RecordImpressionError> { // TODO: Re-enable cache invalidation behind a Nimbus experiment. @@ -131,30 +147,9 @@ where // let _ = self.client.invalidate_cache_by_hash(&request_hash); // } - // TODO: Add count call with _cap_key for impression capping logic - let impression_url = if let Some((_, _cap_key)) = impression_url - .query_pairs() - .find(|(key, _)| key == "cap_key") - { - let mut new_url = impression_url.clone(); - new_url - .query_pairs_mut() - .clear() - .extend_pairs( - impression_url - .query_pairs() - .collect::>() - .iter() - .filter(|(key, _)| key != "cap_key"), - ) - .finish(); - new_url - } else { - impression_url - }; - + let cap_key = pop_query_param_from_url(&mut impression_url, "cap_key"); self.client - .record_impression(impression_url, ohttp) + .record_impression(impression_url, ohttp, cap_key.as_deref()) .inspect_err(|e| { self.telemetry.record(e); }) @@ -185,10 +180,17 @@ where ad_placement_requests: Vec, flags: AdRequestFlags, options: Option, + impression_capping_policy: Option, ohttp: bool, ) -> Result, RequestAdsError> { let response = self - .request_ads::(ad_placement_requests, flags, options, ohttp) + .request_ads::( + ad_placement_requests, + flags, + options, + impression_capping_policy, + ohttp, + ) .inspect_err(|e| { self.telemetry.record(e); })?; @@ -201,9 +203,16 @@ where ad_placement_requests: Vec, flags: AdRequestFlags, options: Option, + impression_capping_policy: Option, ohttp: bool, ) -> Result>, RequestAdsError> { - let result = self.request_ads::(ad_placement_requests, flags, options, ohttp); + let result = self.request_ads::( + ad_placement_requests, + flags, + options, + impression_capping_policy, + ohttp, + ); result .inspect_err(|e| { self.telemetry.record(e); @@ -219,9 +228,16 @@ where ad_placement_requests: Vec, flags: AdRequestFlags, options: Option, + impression_capping_policy: Option, ohttp: bool, ) -> Result, RequestAdsError> { - let result = self.request_ads::(ad_placement_requests, flags, options, ohttp); + let result = self.request_ads::( + ad_placement_requests, + flags, + options, + impression_capping_policy, + ohttp, + ); result .inspect_err(|e| { self.telemetry.record(e); @@ -237,6 +253,7 @@ where placements: Vec, flags: AdRequestFlags, options: Option, + impression_capping_policy: Option, ohttp: bool, ) -> Result, RequestAdsError> where @@ -244,9 +261,15 @@ where { let context_id = self.get_context_id()?; let cache_policy = options.unwrap_or_default(); - let (mut response, request_hash) = - self.client - .fetch_ads::(context_id, flags, placements, cache_policy, ohttp)?; + let impression_capping_policy = impression_capping_policy.unwrap_or_default(); + let (mut response, request_hash) = self.client.fetch_ads::( + context_id, + flags, + placements, + cache_policy, + impression_capping_policy, + ohttp, + )?; response.enrich_callbacks(&request_hash); Ok(response) } @@ -295,6 +318,7 @@ mod tests { cache_config: None, context_id_provider: None, environment: Environment::Test, + impression_log_config: None, telemetry: MozAdsTelemetryWrapper::noop(), }; let client = AdsClient::new(config); @@ -313,13 +337,19 @@ mod tests { .with_body(serde_json::to_string(&expected_response.data).unwrap()) .create(); - let mars_client = MARSClient::new(Environment::Test, None, MozAdsTelemetryWrapper::noop()); + let mars_client = MARSClient::new( + Environment::Test, + None, + None, + MozAdsTelemetryWrapper::noop(), + ); let ads_client = new_with_mars_client(mars_client); let result = ads_client.request_image_ads( make_happy_placement_requests(), AdRequestFlags::default(), None, + None, false, ); assert!(result.is_ok()); @@ -337,13 +367,19 @@ mod tests { .with_body(serde_json::to_string(&expected_response.data).unwrap()) .create(); - let mars_client = MARSClient::new(Environment::Test, None, MozAdsTelemetryWrapper::noop()); + let mars_client = MARSClient::new( + Environment::Test, + None, + None, + MozAdsTelemetryWrapper::noop(), + ); let ads_client = new_with_mars_client(mars_client); let result = ads_client.request_spoc_ads( make_happy_placement_requests(), AdRequestFlags::default(), None, + None, false, ); assert!(result.is_ok()); @@ -361,13 +397,19 @@ mod tests { .with_body(serde_json::to_string(&expected_response.data).unwrap()) .create(); - let mars_client = MARSClient::new(Environment::Test, None, MozAdsTelemetryWrapper::noop()); + let mars_client = MARSClient::new( + Environment::Test, + None, + None, + MozAdsTelemetryWrapper::noop(), + ); let ads_client = new_with_mars_client(mars_client); let result = ads_client.request_tile_ads( make_happy_placement_requests(), AdRequestFlags::default(), None, + None, false, ); assert!(result.is_ok()); @@ -399,6 +441,7 @@ mod tests { cache_config: None, context_id_provider: Some(Box::new(FixedContextId)), environment: Environment::Test, + impression_log_config: None, telemetry: MozAdsTelemetryWrapper::noop(), }; let client = AdsClient::new(config); @@ -409,6 +452,7 @@ mod tests { make_happy_placement_requests(), AdRequestFlags::default(), None, + None, false, ); assert!(result.is_ok()); @@ -418,7 +462,12 @@ mod tests { #[test] fn test_record_impression_removes_cap_key() { viaduct_dev::init_backend_dev(); - let mars_client = MARSClient::new(Environment::Test, None, MozAdsTelemetryWrapper::noop()); + let mars_client = MARSClient::new( + Environment::Test, + None, + None, + MozAdsTelemetryWrapper::noop(), + ); let ads_client = new_with_mars_client(mars_client); let base_url = mockito::server_url(); @@ -452,6 +501,7 @@ mod tests { let mars_client = MARSClient::new( Environment::Test, Some(cache), + None, MozAdsTelemetryWrapper::noop(), ); let ads_client = new_with_mars_client(mars_client); @@ -470,6 +520,7 @@ mod tests { make_happy_placement_requests(), AdRequestFlags::default(), None, + None, false, ) .unwrap(); @@ -484,6 +535,7 @@ mod tests { make_happy_placement_requests(), AdRequestFlags::default(), None, + None, false, ) .unwrap(); @@ -495,6 +547,7 @@ mod tests { make_happy_placement_requests(), AdRequestFlags::default(), Some(CachePolicy::default()), + None, false, ) .unwrap(); diff --git a/components/ads-client/src/client/config.rs b/components/ads-client/src/client/config.rs index 7c86c24141..937595f0ae 100644 --- a/components/ads-client/src/client/config.rs +++ b/components/ads-client/src/client/config.rs @@ -13,6 +13,7 @@ where pub cache_config: Option, pub context_id_provider: Option>, pub environment: Environment, + pub impression_log_config: Option, 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 ImpressionLogConfig { + pub db_path: String, +} diff --git a/components/ads-client/src/clock.rs b/components/ads-client/src/clock.rs new file mode 100644 index 0000000000..4ad821fc5b --- /dev/null +++ b/components/ads-client/src/clock.rs @@ -0,0 +1,35 @@ +/* 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/. */ + +pub trait Clock: Send + Sync + 'static { + fn now_epoch_seconds(&self) -> i64; + #[cfg(test)] + fn advance(&self, secs: i64); +} + +#[cfg(test)] +pub struct TestClock { + now: std::sync::atomic::AtomicI64, +} + +#[cfg(test)] +impl TestClock { + pub fn new(start: i64) -> Self { + Self { + now: std::sync::atomic::AtomicI64::new(start), + } + } +} + +#[cfg(test)] +impl Clock for TestClock { + fn now_epoch_seconds(&self) -> i64 { + self.now.load(std::sync::atomic::Ordering::Relaxed) + } + + fn advance(&self, secs: i64) { + self.now + .fetch_add(secs, std::sync::atomic::Ordering::Relaxed); + } +} diff --git a/components/ads-client/src/ffi.rs b/components/ads-client/src/ffi.rs index ba50352fdc..0b5cbdf7d4 100644 --- a/components/ads-client/src/ffi.rs +++ b/components/ads-client/src/ffi.rs @@ -8,10 +8,11 @@ pub mod telemetry; use std::sync::Arc; -use crate::client::config::{AdsCacheConfig, AdsClientConfig}; +use crate::client::config::{AdsCacheConfig, AdsClientConfig, ImpressionLogConfig}; use crate::client::{AdsClient, ContextIdProvider}; use crate::ffi::telemetry::MozAdsTelemetryWrapper; use crate::http_cache::CachePolicy; +use crate::impression_log::ImpressionCappingPolicy; use crate::mars::ad_request::{ AdContentCategory, AdPlacementRequest, AdRequestFlags, IABContentTaxonomy, }; @@ -58,6 +59,7 @@ impl From for Box { #[derive(Default, uniffi::Record)] pub struct MozAdsRequestOptions { pub cache_policy: Option, + pub impression_capping_policy: Option, #[uniffi(default)] pub flags: HashMap, #[uniffi(default = false)] @@ -106,6 +108,7 @@ struct MozAdsClientBuilderInner { cache_config: Option, context_id_provider: Option>, environment: Option, + impression_log_config: Option, telemetry: Option>, } @@ -132,6 +135,7 @@ impl MozAdsClientBuilder { .map(MozAdsContextIdProviderWrapper::new) .map(Into::into), environment: inner.environment.unwrap_or_default().into(), + impression_log_config: inner.impression_log_config.clone().map(Into::into), telemetry: inner .telemetry .clone() @@ -186,6 +190,11 @@ pub struct MozAdsCacheConfig { pub max_size_mib: Option, } +#[derive(Clone, uniffi::Record)] +pub struct MozAdsImpressionLogConfig { + pub db_path: String, +} + #[derive(Debug, PartialEq, uniffi::Record)] pub struct MozAdsContentCategory { pub categories: Vec, @@ -232,6 +241,13 @@ pub enum MozAdsCacheMode { NetworkFirst, } +#[derive(Clone, Copy, Debug, Default, uniffi::Enum)] +pub enum MozAdsImpressionCappingPolicy { + #[default] + TelemetryOnly, + ImpressionCapEnforced, +} + #[derive(Debug, PartialEq, uniffi::Record)] pub struct MozAdsImage { pub alt_text: Option, @@ -421,6 +437,17 @@ impl From for CachePolicy { } } +impl From for ImpressionCappingPolicy { + fn from(policy: MozAdsImpressionCappingPolicy) -> Self { + match policy { + MozAdsImpressionCappingPolicy::TelemetryOnly => ImpressionCappingPolicy::TelemetryOnly, + MozAdsImpressionCappingPolicy::ImpressionCapEnforced => { + ImpressionCappingPolicy::ImpressionCapEnforced + } + } + } +} + impl From<&MozAdsIABContent> for AdContentCategory { fn from(content: &MozAdsIABContent) -> Self { Self { @@ -436,12 +463,21 @@ impl From<&MozAdsRequestOptions> for AdRequestFlags { } } -impl From for CachePolicy { - fn from(options: MozAdsRequestOptions) -> Self { +impl From<&MozAdsRequestOptions> for CachePolicy { + fn from(options: &MozAdsRequestOptions) -> Self { options.cache_policy.map(Into::into).unwrap_or_default() } } +impl From<&MozAdsRequestOptions> for ImpressionCappingPolicy { + fn from(options: &MozAdsRequestOptions) -> Self { + options + .impression_capping_policy + .map(Into::into) + .unwrap_or_default() + } +} + impl From for AdsCacheConfig { fn from(config: MozAdsCacheConfig) -> Self { Self { @@ -452,6 +488,14 @@ impl From for AdsCacheConfig { } } +impl From for ImpressionLogConfig { + fn from(config: MozAdsImpressionLogConfig) -> 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 70972bf66f..7b905ea28d 100644 --- a/components/ads-client/src/ffi/telemetry.rs +++ b/components/ads-client/src/ffi/telemetry.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use crate::client::error::RequestAdsError; use crate::client::ClientOperationEvent; use crate::http_cache::{CacheOutcome, HttpCacheBuilderError}; +use crate::impression_log::{ImpressionLogBuilderError, ImpressionLogOutcome}; use crate::mars::error::{RecordClickError, RecordImpressionError, ReportAdError}; use crate::telemetry::Telemetry; @@ -19,6 +20,8 @@ pub trait MozAdsTelemetry: Send + Sync { fn record_client_operation_total(&self, label: String); fn record_deserialization_error(&self, label: String, value: String); fn record_http_cache_outcome(&self, label: String, value: String); + fn record_build_impression_log_error(&self, label: String, value: String); + fn record_impression_log_outcome(&self, label: String, value: String); } pub struct NoopMozAdsTelemetry; @@ -29,6 +32,8 @@ impl MozAdsTelemetry for NoopMozAdsTelemetry { fn record_client_operation_total(&self, _label: String) {} fn record_deserialization_error(&self, _label: String, _value: String) {} fn record_http_cache_outcome(&self, _label: String, _value: String) {} + fn record_build_impression_log_error(&self, _label: String, _value: String) {} + fn record_impression_log_outcome(&self, _label: String, _value: String) {} } #[derive(Clone)] @@ -125,6 +130,47 @@ impl Telemetry for MozAdsTelemetryWrapper { ); return; } + if let Some(impression_log_builder_error) = + event.downcast_ref::() + { + self.inner.record_build_impression_log_error( + match impression_log_builder_error { + ImpressionLogBuilderError::EmptyDbPath => "empty_db_path".to_string(), + ImpressionLogBuilderError::Database(_) => "database_error".to_string(), + }, + format!("{}", impression_log_builder_error), + ); + return; + } + if let Some(impression_log_outcome) = event.downcast_ref::() { + self.inner.record_impression_log_outcome( + match impression_log_outcome { + ImpressionLogOutcome::RetainImpressionsFailed(_) => { + "retain_impressions_failed".to_string() + } + ImpressionLogOutcome::RecordImpressionFailed(_) => { + "record_impression_failed".to_string() + } + ImpressionLogOutcome::CountImpressionsFailed(_) => { + "count_impressions_failed".to_string() + } + ImpressionLogOutcome::ImpressionCapHit => "impression_cap_hit".to_string(), + ImpressionLogOutcome::ImpressionCapEnforced => { + "impression_cap_enforced".to_string() + } + ImpressionLogOutcome::ImpressionCapNotEnforced => { + "impression_cap_not_enforced".to_string() + } + }, + match impression_log_outcome { + ImpressionLogOutcome::RetainImpressionsFailed(e) + | ImpressionLogOutcome::RecordImpressionFailed(e) + | ImpressionLogOutcome::CountImpressionsFailed(e) => e.to_string(), + _ => "".to_string(), + }, + ); + return; + } eprintln!("Unsupported telemetry event type: {:?}", event.type_id()); #[cfg(test)] panic!("Unsupported telemetry event type: {:?}", event.type_id()); diff --git a/components/ads-client/src/impression_log.rs b/components/ads-client/src/impression_log.rs new file mode 100644 index 0000000000..0cc1ebeb80 --- /dev/null +++ b/components/ads-client/src/impression_log.rs @@ -0,0 +1,56 @@ +/* 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/. */ + +mod builder; +mod clock; +mod connection_initializer; +mod outcome; +mod store; + +use std::collections::HashMap; +use std::path::Path; + +use self::builder::ImpressionLogBuilder; +use self::store::ImpressionLogStore; + +pub use self::builder::ImpressionLogBuilderError; +pub use self::outcome::ImpressionLogOutcome; + +#[derive(Clone, Copy, Debug, Default)] +pub enum ImpressionCappingPolicy { + #[default] + TelemetryOnly, + ImpressionCapEnforced, +} + +pub struct ImpressionLog { + store: ImpressionLogStore, +} + +impl ImpressionLog { + pub fn builder>(db_path: P) -> ImpressionLogBuilder { + ImpressionLogBuilder::new(db_path.as_ref()) + } + + pub fn record_impression(&self, cap_key: &str) -> Result<(), rusqlite::Error> { + self.store.record_impression(cap_key)?; + Ok(()) + } + + pub fn count_impressions( + &self, + cap_keys: impl IntoIterator, + ) -> Result, rusqlite::Error> { + let counts = self.store.count_impressions(cap_keys)?; + Ok(counts) + } + + pub fn retain_impressions( + &self, + cap_keys: impl IntoIterator, + ) -> Result<(), rusqlite::Error> { + self.store.retain_impressions(cap_keys)?; + Ok(()) + } +} diff --git a/components/ads-client/src/impression_log/builder.rs b/components/ads-client/src/impression_log/builder.rs new file mode 100644 index 0000000000..b277b53ec6 --- /dev/null +++ b/components/ads-client/src/impression_log/builder.rs @@ -0,0 +1,59 @@ +/* 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 std::path::PathBuf; + +use crate::impression_log::connection_initializer::ImpressionLogConnectionInitializer; +use crate::impression_log::store::ImpressionLogStore; +use crate::impression_log::ImpressionLog; + +use rusqlite::Connection; +use sql_support::open_database; + +#[derive(Debug, thiserror::Error)] +pub enum ImpressionLogBuilderError { + #[error("Database path cannot be empty")] + EmptyDbPath, + #[error("Database error: {0}")] + Database(#[from] open_database::Error), +} + +pub struct ImpressionLogBuilder { + db_path: PathBuf, +} + +impl ImpressionLogBuilder { + pub fn new(db_path: impl Into) -> Self { + Self { + db_path: db_path.into(), + } + } + + fn validate(&self) -> Result<(), ImpressionLogBuilderError> { + if self.db_path.to_string_lossy().trim().is_empty() { + return Err(ImpressionLogBuilderError::EmptyDbPath); + } + + Ok(()) + } + + fn open_connection(&self) -> Result { + let initializer = ImpressionLogConnectionInitializer {}; + let conn = if cfg!(test) { + open_database::open_memory_database(&initializer)? + } else { + open_database::open_database(&self.db_path, &initializer)? + }; + Ok(conn) + } + + pub fn build(&self) -> Result { + self.validate()?; + + let conn = self.open_connection()?; + let store = ImpressionLogStore::new(conn); + + Ok(ImpressionLog { store }) + } +} diff --git a/components/ads-client/src/impression_log/clock.rs b/components/ads-client/src/impression_log/clock.rs new file mode 100644 index 0000000000..3d5e521d7a --- /dev/null +++ b/components/ads-client/src/impression_log/clock.rs @@ -0,0 +1,21 @@ +/* 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::clock::Clock; + +pub struct ImpressionLogClock; + +impl Clock for ImpressionLogClock { + fn now_epoch_seconds(&self) -> i64 { + chrono::Utc::now().timestamp() + } + + #[cfg(test)] + fn advance(&self, _secs: i64) { + panic!( + "You cannot advance a non-test clock. + Be sure to build the log or store with the test clock for time-dependent tests." + ) + } +} diff --git a/components/ads-client/src/impression_log/connection_initializer.rs b/components/ads-client/src/impression_log/connection_initializer.rs new file mode 100644 index 0000000000..1406084465 --- /dev/null +++ b/components/ads-client/src/impression_log/connection_initializer.rs @@ -0,0 +1,84 @@ +/* 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::{vtab::array, Connection}; +use sql_support::open_database; +use std::time::Duration; + +pub struct ImpressionLogConnectionInitializer {} + +impl open_database::ConnectionInitializer for ImpressionLogConnectionInitializer { + const NAME: &'static str = "impression_log"; + const END_VERSION: u32 = 1; + + fn prepare(&self, conn: &Connection, _db_empty: bool) -> open_database::Result<()> { + conn.execute_batch("PRAGMA journal_mode=wal;")?; + array::load_module(conn)?; + 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 impression_log ( + cap_key TEXT NOT NULL, + recorded_at INTEGER NOT NULL + ); + CREATE UNIQUE INDEX IF NOT EXISTS idx_impression_log_pk ON impression_log(cap_key, recorded_at); + "; + // 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 impression_log")?; + 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 = ImpressionLogConnectionInitializer {}; + + // Create a corrupted table missing needed index columns + conn.execute_batch("CREATE TABLE impression_log (cap_key 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('impression_log')", + [], + |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/impression_log/outcome.rs b/components/ads-client/src/impression_log/outcome.rs new file mode 100644 index 0000000000..3cfa8ad6eb --- /dev/null +++ b/components/ads-client/src/impression_log/outcome.rs @@ -0,0 +1,16 @@ +/* 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/. +*/ + +#[derive(Debug)] +pub enum ImpressionLogOutcome { + // DB errors + RecordImpressionFailed(rusqlite::Error), // Failed to record impression to log + CountImpressionsFailed(rusqlite::Error), // Failed to get counts of impressions from log + RetainImpressionsFailed(rusqlite::Error), // Failed to clear impressions from log + // Simple events + ImpressionCapHit, // Impression limit reached for a cap_key + ImpressionCapEnforced, // Ad filtered from results due to CappingPolicy + ImpressionCapNotEnforced, // Ad remained in results due to CappingPolicy +} diff --git a/components/ads-client/src/impression_log/store.rs b/components/ads-client/src/impression_log/store.rs new file mode 100644 index 0000000000..75abdf8773 --- /dev/null +++ b/components/ads-client/src/impression_log/store.rs @@ -0,0 +1,275 @@ +/* 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 std::collections::HashMap; +use std::{rc::Rc, sync::Arc}; + +use parking_lot::Mutex; +use rusqlite::{params, types::Value, Connection, Result as SqliteResult, Row}; +use sql_support::ConnExt; + +use crate::clock::Clock; +use crate::impression_log::clock::ImpressionLogClock; + +const SECONDS_IN_DAY: i64 = 60 * 60 * 24; + +#[cfg(test)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum FaultKind { + None, + RecordImpression, + CountImpressions, + RetainImpressions, +} + +pub struct ImpressionLogStore { + conn: Mutex, + clock: Arc, + #[cfg(test)] + fault: Mutex, +} + +fn as_sql_values(raw_values: impl IntoIterator) -> Rc> { + Rc::new( + raw_values + .into_iter() + .map(|s| Value::Text(s.to_string())) + .collect(), + ) +} + +impl ImpressionLogStore { + /// Create new store from connection + pub fn new(conn: Connection) -> Self { + Self { + conn: Mutex::new(conn), + clock: Arc::new(ImpressionLogClock), + #[cfg(test)] + fault: parking_lot::Mutex::new(FaultKind::None), + } + } + + /// Add impression to log. + pub fn record_impression(&self, cap_key: &str) -> SqliteResult { + #[cfg(test)] + if *self.fault.lock() == FaultKind::RecordImpression { + return Err(Self::forced_fault_error("forced record_impression failure")); + } + + let conn = self.conn.lock(); + conn.execute( + "INSERT INTO impression_log (cap_key, recorded_at) + VALUES (?1, ?2) + ON CONFLICT (cap_key, recorded_at) DO NOTHING;", + params![cap_key, self.clock.now_epoch_seconds()], + ) + } + + /// Counts impressions in log. + pub fn count_impressions( + &self, + cap_keys: impl IntoIterator, + ) -> SqliteResult> { + #[cfg(test)] + if *self.fault.lock() == FaultKind::CountImpressions { + return Err(Self::forced_fault_error("forced count_impressions failure")); + } + + let conn = self.conn.lock(); + conn.query_rows_into( + "SELECT value, COUNT(*) + FROM impression_log + INNER JOIN rarray(?1) ON value = cap_key + WHERE recorded_at > ?2 + GROUP BY value;", + params![ + as_sql_values(cap_keys), + self.clock.now_epoch_seconds() - SECONDS_IN_DAY, + ], + |row: &Row<'_>| Ok((row.get::<_, String>(0)?, row.get::<_, u32>(1)?)), + ) + } + + /// Removes other impressions from log. + pub fn retain_impressions( + &self, + cap_keys: impl IntoIterator, + ) -> SqliteResult { + #[cfg(test)] + if *self.fault.lock() == FaultKind::RetainImpressions { + return Err(Self::forced_fault_error( + "forced retain_impressions failure", + )); + } + + let conn = self.conn.lock(); + conn.execute( + "DELETE FROM impression_log + WHERE rowid NOT IN ( + SELECT impression_log.rowid + FROM impression_log + INNER JOIN rarray(?1) ON value = cap_key + );", + params![as_sql_values(cap_keys)], + ) + } + + #[cfg(test)] + pub fn new_with_test_clock(conn: Connection) -> Self { + use crate::clock::TestClock; + + Self { + conn: Mutex::new(conn), + clock: Arc::new(TestClock::new(chrono::Utc::now().timestamp())), + fault: parking_lot::Mutex::new(FaultKind::None), + } + } + + #[cfg(test)] + 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::impression_log::connection_initializer::ImpressionLogConnectionInitializer; + use sql_support::open_database; + + fn create_test_store() -> ImpressionLogStore { + let initializer = ImpressionLogConnectionInitializer {}; + let conn = open_database::open_memory_database(&initializer) + .expect("failed to open memory cache db"); + ImpressionLogStore::new_with_test_clock(conn) + } + + #[test] + fn test_record_impression_fault_injection() { + let store = create_test_store(); + store.set_fault(FaultKind::RecordImpression); + + let err = store.record_impression("test").unwrap_err(); + + match err { + rusqlite::Error::SqliteFailure(_, Some(msg)) => { + assert!(msg.contains("forced record_impression failure")); + } + other => panic!("unexpected error: {other:?}"), + } + } + + #[test] + fn test_count_impressions_fault_injection() { + let store = create_test_store(); + store.set_fault(FaultKind::CountImpressions); + + let err = store.count_impressions(["test_cap_key"]).unwrap_err(); + + match err { + rusqlite::Error::SqliteFailure(_, Some(msg)) => { + assert!(msg.contains("forced count_impressions failure")); + } + other => panic!("unexpected error: {other:?}"), + } + } + + #[test] + fn test_retain_impressions_fault_injection() { + let store = create_test_store(); + store.set_fault(FaultKind::RetainImpressions); + + let err = store.retain_impressions(["test_cap_key"]).unwrap_err(); + + match err { + rusqlite::Error::SqliteFailure(_, Some(msg)) => { + assert!(msg.contains("forced retain_impressions failure")); + } + other => panic!("unexpected error: {other:?}"), + } + } + + #[test] + fn test_impression_roundtrip_simple() { + let store = create_test_store(); + + store.record_impression("test_cap_key").unwrap(); + + assert_eq!( + store.count_impressions(["test_cap_key"]).unwrap(), + HashMap::from([("test_cap_key".into(), 1)]) + ) + } + + #[test] + fn test_impression_roundtrip_multiple() { + let store = create_test_store(); + + store.record_impression("test_cap_key1").unwrap(); + store.record_impression("test_cap_key2").unwrap(); + + assert_eq!( + store + .count_impressions(["test_cap_key1", "test_cap_key2"]) + .unwrap(), + HashMap::from([("test_cap_key1".into(), 1), ("test_cap_key2".into(), 1)]) + ) + } + + #[test] + fn test_impression_roundtrip_duplicate() { + let store = create_test_store(); + + store.record_impression("test_cap_key").unwrap(); + store.record_impression("test_cap_key").unwrap(); + + assert_eq!( + store.count_impressions(["test_cap_key"]).unwrap(), + HashMap::from([("test_cap_key".into(), 1)]) + ) + } + + #[test] + fn test_impression_roundtrip_over_time() { + let store = create_test_store(); + + store.record_impression("test_cap_key").unwrap(); + store.clock.advance(SECONDS_IN_DAY); + store.record_impression("test_cap_key").unwrap(); + store.clock.advance(1); + store.record_impression("test_cap_key").unwrap(); + + assert_eq!( + store.count_impressions(["test_cap_key"]).unwrap(), + HashMap::from([("test_cap_key".into(), 2)]) + ) + } + + #[test] + fn test_impression_cleanup() { + let store = create_test_store(); + + store.record_impression("test_cap_key1").unwrap(); + store.record_impression("test_cap_key2").unwrap(); + store.retain_impressions(["test_cap_key1"]).unwrap(); + + assert_eq!( + store + .count_impressions(["test_cap_key1", "test_cap_key2"]) + .unwrap(), + HashMap::from([("test_cap_key1".into(), 1)]) + ) + } +} diff --git a/components/ads-client/src/lib.rs b/components/ads-client/src/lib.rs index 2088ee16cb..ea54acf407 100644 --- a/components/ads-client/src/lib.rs +++ b/components/ads-client/src/lib.rs @@ -13,11 +13,14 @@ use url::Url as AdsClientUrl; use client::AdsClient; use http_cache::CachePolicy; +use impression_log::ImpressionCappingPolicy; use mars::ad_request::{AdPlacementRequest, AdRequestFlags}; mod client; +mod clock; mod ffi; pub mod http_cache; +pub mod impression_log; mod mars; pub mod telemetry; @@ -114,9 +117,16 @@ impl MozAdsClient { let options = options.unwrap_or_default(); let flags = AdRequestFlags::from(&options); let ohttp = options.ohttp; - let cache_policy: CachePolicy = options.into(); + let cache_policy = CachePolicy::from(&options); + let impression_capping_policy = ImpressionCappingPolicy::from(&options); let response = inner - .request_image_ads(requests, flags, Some(cache_policy), ohttp) + .request_image_ads( + requests, + flags, + Some(cache_policy), + Some(impression_capping_policy), + ohttp, + ) .map_err(ComponentError::RequestAds)?; Ok(response.into_iter().map(|(k, v)| (k, v.into())).collect()) } @@ -133,9 +143,16 @@ impl MozAdsClient { let options = options.unwrap_or_default(); let flags = AdRequestFlags::from(&options); let ohttp = options.ohttp; - let cache_policy: CachePolicy = options.into(); + let cache_policy = CachePolicy::from(&options); + let impression_capping_policy = ImpressionCappingPolicy::from(&options); let response = inner - .request_spoc_ads(requests, flags, Some(cache_policy), ohttp) + .request_spoc_ads( + requests, + flags, + Some(cache_policy), + Some(impression_capping_policy), + ohttp, + ) .map_err(ComponentError::RequestAds)?; Ok(response .into_iter() @@ -155,9 +172,16 @@ impl MozAdsClient { let options = options.unwrap_or_default(); let flags = AdRequestFlags::from(&options); let ohttp = options.ohttp; - let cache_policy: CachePolicy = options.into(); + let cache_policy = CachePolicy::from(&options); + let impression_capping_policy = ImpressionCappingPolicy::from(&options); let response = inner - .request_tile_ads(requests, flags, Some(cache_policy), ohttp) + .request_tile_ads( + requests, + flags, + Some(cache_policy), + Some(impression_capping_policy), + ohttp, + ) .map_err(ComponentError::RequestAds)?; Ok(response.into_iter().map(|(k, v)| (k, v.into())).collect()) } diff --git a/components/ads-client/src/mars.rs b/components/ads-client/src/mars.rs index 59e044212b..147b313629 100644 --- a/components/ads-client/src/mars.rs +++ b/components/ads-client/src/mars.rs @@ -5,6 +5,7 @@ pub mod ad_request; pub mod ad_response; +mod capping; pub mod environment; pub mod error; mod preflight; @@ -17,6 +18,7 @@ pub use report_reason::ReportReason; use self::{ ad_request::{AdPlacementRequest, AdRequest, AdRequestFlags}, ad_response::{AdResponse, AdResponseValue}, + capping::MARSCapping, error::{ CallbackRequestError, FetchAdsError, RecordClickError, RecordImpressionError, ReportAdError, }, @@ -25,8 +27,9 @@ use self::{ }; use crate::{ http_cache::{HttpCache, RequestHash}, + impression_log::ImpressionLog, telemetry::Telemetry, - CachePolicy, + CachePolicy, ImpressionCappingPolicy, }; use url::Url; use viaduct::{Headers, Request}; @@ -38,18 +41,26 @@ where environment: Environment, telemetry: T, transport: MARSTransport, + capping: MARSCapping, } impl MARSClient where T: Clone + Telemetry, { - pub fn new(environment: Environment, http_cache: Option, telemetry: T) -> Self { + pub fn new( + environment: Environment, + http_cache: Option, + impression_log: Option, + telemetry: T, + ) -> Self { let transport = MARSTransport::new(http_cache, telemetry.clone()); + let capping = MARSCapping::new(impression_log, telemetry.clone()); Self { environment, telemetry, transport, + capping, } } @@ -63,6 +74,7 @@ where flags: AdRequestFlags, placements: Vec, cache_policy: CachePolicy, + impression_capping_policy: ImpressionCappingPolicy, ohttp: bool, ) -> Result<(AdResponse, RequestHash), FetchAdsError> where @@ -80,7 +92,10 @@ where let response = self.transport.send(ad_request, &cache_policy, ohttp)?; let ads = AdResponse::::parse(response.json()?, &self.telemetry)?; - Ok((ads, request_hash)) + let filtered_ads = self + .capping + .apply_impression_capping(ads, &impression_capping_policy); + Ok((filtered_ads, request_hash)) } // TODO: Remove this allow(dead_code) when cache invalidation is re-enabled behind Nimbus experiment @@ -100,7 +115,11 @@ where &self, callback: Url, ohttp: bool, + cap_key: Option<&str>, ) -> Result<(), RecordImpressionError> { + if let Some(cap_key) = cap_key { + self.capping.record_impression(cap_key); + } Ok(self.make_callback_request(callback, ohttp)?) } @@ -151,10 +170,14 @@ mod tests { }; use mockito::mock; - fn make_test_client(http_cache: Option) -> MARSClient { + fn make_test_client( + http_cache: Option, + impression_log: Option, + ) -> MARSClient { MARSClient::new( Environment::Test, http_cache, + impression_log, MozAdsTelemetryWrapper::noop(), ) } @@ -165,13 +188,13 @@ mod tests { let m = mock("GET", "/impression_callback_url") .with_status(200) .create(); - let client = make_test_client(None); + let client = make_test_client(None, None); let url = Url::parse(&format!( "{}/impression_callback_url", &mockito::server_url() )) .unwrap(); - let result = client.record_impression(url, false); + let result = client.record_impression(url, false, None); assert!(result.is_ok()); m.assert(); } @@ -181,7 +204,7 @@ mod tests { viaduct_dev::init_backend_dev(); let m = mock("GET", "/click_callback_url").with_status(200).create(); - let client = make_test_client(None); + let client = make_test_client(None, None); let url = Url::parse(&format!("{}/click_callback_url", &mockito::server_url())).unwrap(); let result = client.record_click(url, false); assert!(result.is_ok()); @@ -199,7 +222,7 @@ mod tests { .with_status(200) .create(); - let client = make_test_client(None); + let client = make_test_client(None, None); let url = Url::parse(&format!( "{}/report_ad_callback_url", &mockito::server_url() @@ -222,13 +245,14 @@ mod tests { .with_body(serde_json::to_string(&expected_response.data).unwrap()) .create(); - let client = make_test_client(None); + let client = make_test_client(None, None); let result = client.fetch_ads::( TEST_CONTEXT_ID.to_string(), AdRequestFlags::default(), make_happy_placement_requests(), CachePolicy::default(), + ImpressionCappingPolicy::default(), false, ); assert!(result.is_ok()); @@ -253,7 +277,7 @@ mod tests { .max_size(crate::http_cache::ByteSize::mib(1)) .build() .unwrap(); - let client = make_test_client(Some(cache)); + let client = make_test_client(Some(cache), None); // First call should be a miss then warm the cache let (response1, _) = client @@ -262,6 +286,7 @@ mod tests { AdRequestFlags::default(), make_happy_placement_requests(), CachePolicy::default(), + ImpressionCappingPolicy::default(), false, ) .unwrap(); @@ -274,6 +299,7 @@ mod tests { AdRequestFlags::default(), make_happy_placement_requests(), CachePolicy::default(), + ImpressionCappingPolicy::default(), false, ) .unwrap(); @@ -290,7 +316,7 @@ mod tests { .build() .unwrap(); - let client = make_test_client(Some(cache)); + let client = make_test_client(Some(cache), None); let callback_url = Url::parse(&format!("{}/click", mockito::server_url())).unwrap(); let m = mock("GET", "/click").with_status(200).create(); @@ -309,12 +335,12 @@ mod tests { .build() .unwrap(); - let client = make_test_client(Some(cache)); + let client = make_test_client(Some(cache), None); let callback_url = Url::parse(&format!("{}/impression", mockito::server_url())).unwrap(); let m = mock("GET", "/impression").with_status(200).create(); - let result = client.record_impression(callback_url, false); + let result = client.record_impression(callback_url, false, None); assert!(result.is_ok()); m.assert(); } diff --git a/components/ads-client/src/mars/ad_response.rs b/components/ads-client/src/mars/ad_response.rs index 5e057b7606..351f298c2b 100644 --- a/components/ads-client/src/mars/ad_response.rs +++ b/components/ads-client/src/mars/ad_response.rs @@ -47,7 +47,7 @@ impl AdResponse { let hash_str = request_hash.to_string(); for (placement_id, ads) in self.data.iter_mut() { for (position, ad) in ads.iter_mut().enumerate() { - let cap_key = ad.cap_key(); + let cap_key = ad.cap_pair().map(|(cap_key, _)| cap_key.to_owned()); let callbacks = ad.callbacks_mut(); callbacks .click @@ -82,15 +82,13 @@ impl AdResponse { } } -// TODO: Remove this allow(dead_code) when cache invalidation is re-enabled behind Nimbus experiment -#[allow(dead_code)] -pub fn pop_request_hash_from_url(url: &mut Url) -> Option { +pub fn pop_query_param_from_url(url: &mut Url, query_param: &str) -> Option { let mut request_hash = None; let mut query = url::form_urlencoded::Serializer::new(String::new()); for (key, value) in url.query_pairs() { - if key == "request_hash" { - request_hash = Some(RequestHash::from(value.as_ref())); + if key == query_param { + request_hash = Some(value.into_owned()); } else { query.append_pair(&key, &value); } @@ -102,6 +100,7 @@ pub fn pop_request_hash_from_url(url: &mut Url) -> Option { } else { url.set_query(Some(&query_string)); } + request_hash } @@ -163,7 +162,7 @@ pub struct AdCallbacks { pub trait AdResponseValue: DeserializeOwned { fn callbacks_mut(&mut self) -> &mut AdCallbacks; - fn cap_key(&self) -> Option { + fn cap_pair(&self) -> Option<(&str, &u32)> { None } } @@ -179,8 +178,8 @@ impl AdResponseValue for AdSpoc { &mut self.callbacks } - fn cap_key(&self) -> Option { - Some(self.caps.cap_key.clone()) + fn cap_pair(&self) -> Option<(&str, &u32)> { + Some((&self.caps.cap_key, &self.caps.day)) } } @@ -645,21 +644,21 @@ mod tests { } #[test] - fn test_pop_request_hash_from_url() { + fn test_pop_query_param_from_url() { let mut url_with_hash = Url::parse("https://example.com/callback?request_hash=abc123def456&other=param") .unwrap(); - let extracted = pop_request_hash_from_url(&mut url_with_hash); - assert_eq!(extracted, Some(RequestHash::from("abc123def456"))); + let extracted = pop_query_param_from_url(&mut url_with_hash, "request_hash"); + assert_eq!(extracted, Some("abc123def456".into())); assert_eq!(url_with_hash.query(), Some("other=param")); let mut url_without_hash = Url::parse("https://example.com/callback?other=param").unwrap(); - let extracted_none = pop_request_hash_from_url(&mut url_without_hash); + let extracted_none = pop_query_param_from_url(&mut url_without_hash, "request_hash"); assert_eq!(extracted_none, None); assert_eq!(url_without_hash.query(), Some("other=param")); let mut url_no_query = Url::parse("https://example.com/callback").unwrap(); - let extracted_empty = pop_request_hash_from_url(&mut url_no_query); + let extracted_empty = pop_query_param_from_url(&mut url_no_query, "request_hash"); assert_eq!(extracted_empty, None); assert_eq!(url_no_query.query(), None); } diff --git a/components/ads-client/src/mars/capping.rs b/components/ads-client/src/mars/capping.rs new file mode 100644 index 0000000000..ab4d3b65e4 --- /dev/null +++ b/components/ads-client/src/mars/capping.rs @@ -0,0 +1,159 @@ +/* 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 std::collections::{HashMap, HashSet}; + +use crate::{ + impression_log::{ImpressionLog, ImpressionLogOutcome}, + mars::ad_response::{AdResponse, AdResponseValue}, + telemetry::Telemetry, + ImpressionCappingPolicy, +}; + +pub struct MARSCapping { + impression_log: Option, + telemetry: T, +} + +impl MARSCapping { + pub fn new(impression_log: Option, telemetry: T) -> Self { + Self { + impression_log, + telemetry, + } + } + + pub fn record_impression(&self, cap_key: &str) { + if let Some(impression_log) = &self.impression_log { + if let Err(e) = impression_log.record_impression(cap_key) { + self.telemetry + .record(&ImpressionLogOutcome::RecordImpressionFailed(e)); + } + }; + } + + pub fn apply_impression_capping( + &self, + mut ads: AdResponse, + impression_capping_policy: &ImpressionCappingPolicy, + ) -> AdResponse { + if let Some(impression_log) = &self.impression_log { + let caps: HashMap<&str, &u32> = ads + .data + .iter() + .flat_map(|(_, placement_ads)| placement_ads.iter().flat_map(|a| a.cap_pair())) + .collect(); + + let counts = match impression_log.count_impressions(caps.keys()) { + Ok(counts) => counts, + Err(e) => { + self.telemetry + .record(&ImpressionLogOutcome::CountImpressionsFailed(e)); + + // Skip unnecessary work if DB access failed + return ads; + } + }; + + if let Err(e) = impression_log.retain_impressions(caps.keys()) { + self.telemetry + .record(&ImpressionLogOutcome::RetainImpressionsFailed(e)); + }; + + let cap_keys_to_filter: HashSet = caps + .iter() + .flat_map(|(&cap_key, max_impressions)| { + if counts.get(cap_key).unwrap_or(&0) >= max_impressions { + self.telemetry + .record(&ImpressionLogOutcome::ImpressionCapHit); + match impression_capping_policy { + ImpressionCappingPolicy::TelemetryOnly => { + self.telemetry + .record(&ImpressionLogOutcome::ImpressionCapNotEnforced); + None + } + ImpressionCappingPolicy::ImpressionCapEnforced => { + self.telemetry + .record(&ImpressionLogOutcome::ImpressionCapEnforced); + Some(cap_key.to_owned()) + } + } + } else { + None + } + }) + .collect(); + + if !cap_keys_to_filter.is_empty() { + ads.data.iter_mut().for_each(|(_, placement_ads)| { + placement_ads.retain(|a| { + if let Some((cap_key, _)) = a.cap_pair() { + !cap_keys_to_filter.contains(cap_key) + } else { + true + } + }); + }); + } + }; + + ads + } +} + +#[cfg(test)] +mod tests { + use url::Url; + + use crate::ffi::telemetry::MozAdsTelemetryWrapper; + use crate::impression_log::ImpressionCappingPolicy; + use crate::mars::ad_response::{AdCallbacks, AdSpoc, SpocFrequencyCaps, SpocRanking}; + + use super::*; + + #[test] + fn test_no_impression_log_does_not_error() { + let capping = MARSCapping::new(None, MozAdsTelemetryWrapper::noop()); + + let spoc = AdSpoc { + block_key: "test_block_key".into(), + callbacks: AdCallbacks { + click: Url::parse("https://example.com/test_click").unwrap(), + impression: Url::parse("https://example.com/test_impression").unwrap(), + report: None, + }, + caps: SpocFrequencyCaps { + cap_key: "test_cap_key".into(), + day: 10, + }, + domain: "example.com".into(), + excerpt: "test_excerpt".into(), + format: "test_format".into(), + image_url: "https://example.com/test_image".into(), + ranking: SpocRanking { + priority: 0, + personalization_models: None, + item_score: 0.0, + }, + sponsor: "test_sponsor".into(), + sponsored_by_override: None, + title: "test_title".into(), + url: "https://example.com/test_url".into(), + }; + + capping.record_impression("test_cap_key"); + capping.apply_impression_capping( + AdResponse:: { + data: HashMap::from([("".into(), vec![spoc.clone()])]), + }, + &ImpressionCappingPolicy::TelemetryOnly, + ); + capping.apply_impression_capping( + AdResponse:: { + data: HashMap::from([("".into(), vec![spoc.clone()])]), + }, + &ImpressionCappingPolicy::TelemetryOnly, + ); + } +}