[AC-152] Fixes moz ads telemetry UniFFI callback leak

CHANGELOG.md

@@ -4,6 +4,10 @@
 
 ## ✨ What's Changed ✨
 
+### Ads-Client
+
+- Add `AdsClient::shutdown()`, which closes the database connection early so it happens before Firefox Desktop's late-write shutdown barrier rather than during GC. In addition, this drops all held UniFFI callbacks (held in the MozAdsTelemetryWrapper) to avoid crash on a Firefox Desktop quit.
+
 ### Autofill
 
 - Add `Store::shutdown()`, which closes the database connection early so it happens before Firefox Desktop's late-write shutdown barrier rather than during GC. Operations after shutdown return `DatabaseClosed`. ([Bug 2050036](https://bugzilla.mozilla.org/show_bug.cgi?id=2050036))

components/ads-client/src/client.rs

@@ -98,6 +98,18 @@ where
         self.client.clear_cache()
     }
 
+    // Shutdown the db connection and drop references to telemetry callbacks.
+    // Should be used only when dropping the ads client, this may be extended to drop more things.
+    pub fn shutdown_client(&mut self) -> Result<(), rusqlite::Error> {
+        // Drop telemetry (within the telemetry wrapper)
+        self.telemetry.shutdown();
+
+        // Shutdown DB
+        self.client.shutdown_db()?;
+
+        Ok(())
+    }
+
     pub fn get_context_id(&self) -> context_id::ApiResult<String> {
         self.context_id_provider.context_id()
     }
@@ -263,6 +275,8 @@ pub enum ClientOperationEvent {
 
 #[cfg(test)]
 mod tests {
+    use std::{assert_eq, assert_ne, sync::Arc};
+
     use crate::{
         ffi::telemetry::MozAdsTelemetryWrapper,
         mars::Environment,
@@ -277,6 +291,7 @@ mod tests {
     fn new_with_mars_client(
         client: MARSClient<MozAdsTelemetryWrapper>,
     ) -> AdsClient<MozAdsTelemetryWrapper> {
+        let telemetry = client.get_telemetry();
         AdsClient {
             client,
             context_id_provider: Box::new(ContextIDComponent::new(
@@ -285,7 +300,7 @@ mod tests {
                 false,
                 Box::new(DefaultContextIdCallback),
             )),
-            telemetry: MozAdsTelemetryWrapper::noop(),
+            telemetry,
         }
     }
 
@@ -502,4 +517,47 @@ mod tests {
         m1.assert();
         m2.assert();
     }
+
+    #[test]
+    fn test_shutdown_telemetry() {
+        viaduct_dev::init_backend_dev();
+
+        // test with client created from config
+        let noop_telemetry = MozAdsTelemetryWrapper::noop();
+        let weak_reference = Arc::downgrade(
+            &noop_telemetry
+                .clone_inner_arc()
+                .expect("Inner telemetry should be Some before dropping"),
+        );
+        let config = AdsClientConfig {
+            cache_config: None,
+            context_id_provider: None,
+            environment: Environment::Test,
+            telemetry: noop_telemetry,
+        };
+        let mut client = AdsClient::new(config);
+
+        // weak ref will show 0 strong references when the Arc<dyn MozAdsTelemetry> is gone.
+        assert_ne!(weak_reference.strong_count(), 0);
+        client.shutdown_client().unwrap();
+        assert_eq!(weak_reference.strong_count(), 0);
+
+        // test also with internal function from_mars
+        let noop_telemetry = MozAdsTelemetryWrapper::noop();
+        let weak_reference = Arc::downgrade(
+            &noop_telemetry
+                .clone_inner_arc()
+                .expect("Inner telemetry should be Some before dropping"),
+        );
+        let cache = HttpCache::builder("test_shutdown_telemetry")
+            .build()
+            .unwrap();
+        let mars_client = MARSClient::new(Environment::Test, Some(cache), noop_telemetry);
+        let mut client = new_with_mars_client(mars_client);
+
+        // weak ref will show 0 strong references when the Arc<dyn MozAdsTelemetry> is gone.
+        assert_ne!(weak_reference.strong_count(), 0);
+        client.shutdown_client().unwrap();
+        assert_eq!(weak_reference.strong_count(), 0);
+    }
 }

components/ads-client/src/ffi/telemetry.rs

@@ -6,6 +6,8 @@
 use std::any::Any;
 use std::sync::Arc;
 
+use parking_lot::RwLock;
+
 use crate::client::error::RequestAdsError;
 use crate::client::ClientOperationEvent;
 use crate::http_cache::{CacheOutcome, HttpCacheBuilderError};
@@ -33,25 +35,42 @@ impl MozAdsTelemetry for NoopMozAdsTelemetry {
 
 #[derive(Clone)]
 pub struct MozAdsTelemetryWrapper {
-    inner: Arc<dyn MozAdsTelemetry>,
+    inner: Arc<RwLock<Option<Arc<dyn MozAdsTelemetry>>>>,
 }
 
 impl MozAdsTelemetryWrapper {
     pub fn new(inner: Arc<dyn MozAdsTelemetry>) -> Self {
-        Self { inner }
+        Self {
+            inner: Arc::new(RwLock::new(Some(inner))),
+        }
     }
 
     pub fn noop() -> Self {
         Self {
-            inner: Arc::new(NoopMozAdsTelemetry),
+            inner: Arc::new(RwLock::new(Some(Arc::new(NoopMozAdsTelemetry)))),
         }
     }
+
+    #[cfg(test)]
+    pub fn clone_inner_arc(&self) -> Option<Arc<dyn MozAdsTelemetry>> {
+        self.inner.read().clone()
+    }
 }
 
 impl Telemetry for MozAdsTelemetryWrapper {
+    // MozAdsTelemetry has hanging uniffi callbacks which need to be explicitly dropped before closing.
+    // This replaces it with a `None` internally, meaning future calls will be noops.
+    fn shutdown(&self) {
+        let _dropped = self.inner.write().take();
+    }
+
     fn record(&self, event: &dyn Any) {
+        let Some(inner) = self.inner.read().clone() else {
+            return;
+        };
+
         if let Some(cache_outcome) = event.downcast_ref::<CacheOutcome>() {
-            self.inner.record_http_cache_outcome(
+            inner.record_http_cache_outcome(
                 match cache_outcome {
                     CacheOutcome::Hit => "hit".to_string(),
                     CacheOutcome::LookupFailed(_) => "lookup_failed".to_string(),
@@ -73,7 +92,7 @@ impl Telemetry for MozAdsTelemetryWrapper {
             return;
         }
         if let Some(client_op) = event.downcast_ref::<ClientOperationEvent>() {
-            self.inner.record_client_operation_total(match client_op {
+            inner.record_client_operation_total(match client_op {
                 ClientOperationEvent::New => "new".to_string(),
                 ClientOperationEvent::RecordClick => "record_click".to_string(),
                 ClientOperationEvent::RecordImpression => "record_impression".to_string(),
@@ -83,7 +102,7 @@ impl Telemetry for MozAdsTelemetryWrapper {
             return;
         }
         if let Some(cache_builder_error) = event.downcast_ref::<HttpCacheBuilderError>() {
-            self.inner.record_build_cache_error(
+            inner.record_build_cache_error(
                 match cache_builder_error {
                     HttpCacheBuilderError::EmptyDbPath => "empty_db_path".to_string(),
                     HttpCacheBuilderError::Database(_) => "database_error".to_string(),
@@ -95,31 +114,29 @@ impl Telemetry for MozAdsTelemetryWrapper {
             return;
         }
         if let Some(record_click_error) = event.downcast_ref::<RecordClickError>() {
-            self.inner.record_client_error(
+            inner.record_client_error(
                 "record_click".to_string(),
                 format!("{}", record_click_error),
             );
             return;
         }
         if let Some(record_impression_error) = event.downcast_ref::<RecordImpressionError>() {
-            self.inner.record_client_error(
+            inner.record_client_error(
                 "record_impression".to_string(),
                 format!("{}", record_impression_error),
             );
             return;
         }
         if let Some(report_ad_error) = event.downcast_ref::<ReportAdError>() {
-            self.inner
-                .record_client_error("report_ad".to_string(), format!("{}", report_ad_error));
+            inner.record_client_error("report_ad".to_string(), format!("{}", report_ad_error));
             return;
         }
         if let Some(request_ads_error) = event.downcast_ref::<RequestAdsError>() {
-            self.inner
-                .record_client_error("request_ads".to_string(), format!("{}", request_ads_error));
+            inner.record_client_error("request_ads".to_string(), format!("{}", request_ads_error));
             return;
         }
         if let Some(json_error) = event.downcast_ref::<serde_json::Error>() {
-            self.inner.record_deserialization_error(
+            inner.record_deserialization_error(
                 "invalid_ad_item".to_string(),
                 format!("{}", json_error),
             );

components/ads-client/src/http_cache.rs

@@ -60,6 +60,10 @@ impl HttpCache {
         Ok(())
     }
 
+    pub fn shutdown_db(self) -> Result<(), rusqlite::Error> {
+        self.store.close()
+    }
+
     pub fn invalidate_by_hash(&self, request_hash: &RequestHash) -> Result<(), rusqlite::Error> {
         self.store.invalidate_by_hash(request_hash)?;
         Ok(())

components/ads-client/src/http_cache/store.rs

@@ -40,6 +40,11 @@ impl HttpCacheStore {
         }
     }
 
+    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::http_cache::clock::TestClock;

components/ads-client/src/lib.rs

@@ -12,9 +12,9 @@ use parking_lot::Mutex;
 use url::Url as AdsClientUrl;
 
 use client::AdsClient;
+use error_support::error;
 use http_cache::CachePolicy;
 use mars::ad_request::{AdPlacementRequest, AdRequestFlags};
-
 mod client;
 mod ffi;
 pub mod http_cache;
@@ -52,6 +52,18 @@ impl MozAdsClient {
             })
     }
 
+    // Allows the ads-client to unload some references and prepare for a safe shutdown.
+    // Other methods should not be called after this one.
+    #[uniffi::method()]
+    pub fn shutdown(&self) -> AdsClientApiResult<()> {
+        let mut inner = self.inner.lock();
+        if let Err(err) = inner.shutdown_client() {
+            // Log the error, but continue with shutdown.
+            error!("Failed to shutdown the ads client: {:?}", err);
+        }
+        Ok(())
+    }
+
     #[handle_error(ComponentError)]
     #[uniffi::method(default(options = None))]
     pub fn record_click(

components/ads-client/src/mars.rs

@@ -57,6 +57,10 @@ where
         self.transport.clear_cache()
     }
 
+    pub fn shutdown_db(&mut self) -> Result<(), rusqlite::Error> {
+        self.transport.shutdown_db()
+    }
+
     pub fn fetch_ads<A>(
         &self,
         context_id: String,
@@ -138,6 +142,11 @@ where
         }
         self.transport.fire(request, ohttp).map_err(Into::into)
     }
+
+    #[cfg(test)]
+    pub fn get_telemetry(&self) -> T {
+        self.telemetry.clone()
+    }
 }
 
 #[cfg(test)]

components/ads-client/src/mars/transport.rs

@@ -29,6 +29,13 @@ impl<T: Telemetry> MARSTransport<T> {
         }
     }
 
+    pub fn shutdown_db(&mut self) -> Result<(), rusqlite::Error> {
+        if let Some(cache) = self.http_cache.take() {
+            cache.shutdown_db()?;
+        }
+        Ok(())
+    }
+
     pub fn clear_cache(&self) -> Result<(), rusqlite::Error> {
         if let Some(cache) = &self.http_cache {
             cache.clear()?;

components/ads-client/src/telemetry.rs

@@ -7,4 +7,8 @@ use std::any::Any;
 
 pub trait Telemetry {
     fn record(&self, event: &dyn Any);
+
+    // Shuts down any telemetry structures. This should be called before dropping the struct implementing this trait.
+    // Future calls to `record` will not record anything.
+    fn shutdown(&self);
 }