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

CHANGELOG.md

@@ -4,6 +4,10 @@
 
 ## ✨ What's Changed ✨
 
+### Ads-Client
+
+- Add `AdsClient::shutdown_client()`, 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

@@ -41,7 +41,7 @@ where
 {
     client: MARSClient<T>,
     context_id_provider: Box<dyn ContextIdProvider>,
-    telemetry: T,
+    telemetry: Option<T>,
 }
 
 impl<T> AdsClient<T>
@@ -90,14 +90,27 @@ where
         Self {
             client,
             context_id_provider,
-            telemetry: telemetry.clone(),
+            telemetry: Some(telemetry.clone()),
         }
     }
 
     pub fn clear_cache(&self) -> Result<(), rusqlite::Error> {
         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> {
+        // Shutdown DB
+        self.client.shutdown_db()?;
+
+        // Drop telemetry recursively
+        self.telemetry = None;
+        self.client.drop_telemetry();
+
+        Ok(())
+    }
+
     pub fn get_context_id(&self) -> context_id::ApiResult<String> {
         self.context_id_provider.context_id()
     }
@@ -112,10 +125,14 @@ where
         self.client
             .record_click(click_url, ohttp)
             .inspect_err(|e| {
-                self.telemetry.record(e);
+                if let Some(telemetry) = &self.telemetry {
+                    telemetry.record(e);
+                }
             })
             .inspect(|_| {
-                self.telemetry.record(&ClientOperationEvent::RecordClick);
+                if let Some(telemetry) = &self.telemetry {
+                    telemetry.record(&ClientOperationEvent::RecordClick);
+                }
             })
     }
 
@@ -156,11 +173,14 @@ where
         self.client
             .record_impression(impression_url, ohttp)
             .inspect_err(|e| {
-                self.telemetry.record(e);
+                if let Some(telemetry) = &self.telemetry {
+                    telemetry.record(e);
+                }
             })
             .inspect(|_| {
-                self.telemetry
-                    .record(&ClientOperationEvent::RecordImpression);
+                if let Some(telemetry) = &self.telemetry {
+                    telemetry.record(&ClientOperationEvent::RecordImpression);
+                }
             })
     }
 
@@ -173,10 +193,14 @@ where
         self.client
             .report_ad(report_url, reason, ohttp)
             .inspect_err(|e| {
-                self.telemetry.record(e);
+                if let Some(telemetry) = &self.telemetry {
+                    telemetry.record(e);
+                }
             })
             .inspect(|_| {
-                self.telemetry.record(&ClientOperationEvent::ReportAd);
+                if let Some(telemetry) = &self.telemetry {
+                    telemetry.record(&ClientOperationEvent::ReportAd);
+                }
             })
     }
 
@@ -190,9 +214,13 @@ where
         let response = self
             .request_ads::<AdImage>(ad_placement_requests, flags, options, ohttp)
             .inspect_err(|e| {
-                self.telemetry.record(e);
+                if let Some(telemetry) = &self.telemetry {
+                    telemetry.record(e);
+                }
             })?;
-        self.telemetry.record(&ClientOperationEvent::RequestAds);
+        if let Some(telemetry) = &self.telemetry {
+            telemetry.record(&ClientOperationEvent::RequestAds);
+        }
         Ok(response.take_first())
     }
 
@@ -206,10 +234,14 @@ where
         let result = self.request_ads::<AdSpoc>(ad_placement_requests, flags, options, ohttp);
         result
             .inspect_err(|e| {
-                self.telemetry.record(e);
+                if let Some(telemetry) = &self.telemetry {
+                    telemetry.record(e);
+                }
             })
             .map(|response| {
-                self.telemetry.record(&ClientOperationEvent::RequestAds);
+                if let Some(telemetry) = &self.telemetry {
+                    telemetry.record(&ClientOperationEvent::RequestAds);
+                }
                 response.data
             })
     }
@@ -224,10 +256,14 @@ where
         let result = self.request_ads::<AdTile>(ad_placement_requests, flags, options, ohttp);
         result
             .inspect_err(|e| {
-                self.telemetry.record(e);
+                if let Some(telemetry) = &self.telemetry {
+                    telemetry.record(e);
+                }
             })
             .map(|response| {
-                self.telemetry.record(&ClientOperationEvent::RequestAds);
+                if let Some(telemetry) = &self.telemetry {
+                    telemetry.record(&ClientOperationEvent::RequestAds);
+                }
                 response.take_first()
             })
     }
@@ -263,6 +299,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 +315,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 +324,7 @@ mod tests {
                 false,
                 Box::new(DefaultContextIdCallback),
             )),
-            telemetry: MozAdsTelemetryWrapper::noop(),
+            telemetry,
         }
     }
 
@@ -502,4 +541,39 @@ 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());
+        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());
+        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

@@ -46,6 +46,11 @@ impl MozAdsTelemetryWrapper {
             inner: Arc::new(NoopMozAdsTelemetry),
         }
     }
+
+    #[cfg(test)]
+    pub fn clone_inner_arc(&self) -> Arc<dyn MozAdsTelemetry> {
+        self.inner.clone()
+    }
 }
 
 impl Telemetry for MozAdsTelemetryWrapper {

components/ads-client/src/http_cache.rs

@@ -60,6 +60,10 @@ impl HttpCache {
         Ok(())
     }
 
+    pub fn shutdown(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,15 @@ impl MozAdsClient {
             })
     }
 
+    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

@@ -36,7 +36,7 @@ where
     T: Clone + Telemetry,
 {
     environment: Environment,
-    telemetry: T,
+    telemetry: Option<T>,
     transport: MARSTransport<T>,
 }
 
@@ -48,7 +48,7 @@ where
         let transport = MARSTransport::new(http_cache, telemetry.clone());
         Self {
             environment,
-            telemetry,
+            telemetry: Some(telemetry),
             transport,
         }
     }
@@ -57,6 +57,15 @@ where
         self.transport.clear_cache()
     }
 
+    pub fn shutdown_db(&mut self) -> Result<(), rusqlite::Error> {
+        self.transport.shutdown_db()
+    }
+
+    pub fn drop_telemetry(&mut self) {
+        self.telemetry = None;
+        self.transport.drop_telemetry();
+    }
+
     pub fn fetch_ads<A>(
         &self,
         context_id: String,
@@ -79,7 +88,7 @@ where
         }
 
         let response = self.transport.send(ad_request, &cache_policy, ohttp)?;
-        let ads = AdResponse::<A>::parse(response.json()?, &self.telemetry)?;
+        let ads = AdResponse::<A>::parse(response.json()?, self.telemetry.as_ref())?;
         Ok((ads, request_hash))
     }
 
@@ -138,6 +147,11 @@ where
         }
         self.transport.fire(request, ohttp).map_err(Into::into)
     }
+
+    #[cfg(test)]
+    pub fn get_telemetry(&self) -> Option<T> {
+        self.telemetry.clone()
+    }
 }
 
 #[cfg(test)]

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

@@ -18,7 +18,7 @@ pub struct AdResponse<A: AdResponseValue> {
 impl<A: AdResponseValue> AdResponse<A> {
     pub fn parse<T: Telemetry>(
         data: serde_json::Value,
-        telemetry: &T,
+        telemetry: Option<&T>,
     ) -> Result<AdResponse<A>, serde_json::Error> {
         let raw: HashMap<String, serde_json::Value> = serde_json::from_value(data)?;
         let mut result = HashMap::new();
@@ -30,7 +30,9 @@ impl<A: AdResponseValue> AdResponse<A> {
                     match serde_json::from_value::<A>(item.clone()) {
                         Ok(ad) => ads.push(ad),
                         Err(e) => {
-                            telemetry.record(&e);
+                            if let Some(telemetry) = telemetry {
+                                telemetry.record(&e);
+                            }
                         }
                     }
                 }
@@ -335,7 +337,8 @@ mod tests {
         });
 
         let parsed =
-            AdResponse::<AdImage>::parse(raw_ad_response, &MozAdsTelemetryWrapper::noop()).unwrap();
+            AdResponse::<AdImage>::parse(raw_ad_response, Some(&MozAdsTelemetryWrapper::noop()))
+                .unwrap();
 
         let expected = AdResponse {
             data: HashMap::from([(
@@ -372,7 +375,8 @@ mod tests {
         });
 
         let parsed =
-            AdResponse::<AdImage>::parse(raw_ad_response, &MozAdsTelemetryWrapper::noop()).unwrap();
+            AdResponse::<AdImage>::parse(raw_ad_response, Some(&MozAdsTelemetryWrapper::noop()))
+                .unwrap();
 
         let expected = AdResponse {
             data: HashMap::from([]),