1 files changed, 36 insertions, 2 deletions

diff --git a/components/net/image_cache.rs b/components/net/image_cache.rs
index 81d9a09bf39..68d14f4c10c 100644
--- a/components/net/image_cache.rs
+++ b/components/net/image_cache.rs
@@ -2,8 +2,10 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 
-use std::collections::HashMap;
+use std::cell::{LazyCell, RefCell};
 use std::collections::hash_map::Entry::{Occupied, Vacant};
+use std::collections::{HashMap, HashSet};
+use std::ffi::c_void;
 use std::sync::{Arc, Mutex};
 use std::{mem, thread};
 
@@ -11,6 +13,8 @@ use compositing_traits::{CrossProcessCompositorApi, SerializableImageData};
 use imsz::imsz_from_reader;
 use ipc_channel::ipc::IpcSharedMemory;
 use log::{debug, warn};
+use malloc_size_of::{MallocSizeOf as MallocSizeOfTrait, MallocSizeOfOps};
+use malloc_size_of_derive::MallocSizeOf;
 use net_traits::image_cache::{
     ImageCache, ImageCacheResult, ImageOrMetadataAvailable, ImageResponder, ImageResponse,
     PendingImageId, UsePlaceholder,
@@ -18,6 +22,8 @@ use net_traits::image_cache::{
 use net_traits::request::CorsSettings;
 use net_traits::{FetchMetadata, FetchResponseMsg, FilteredMetadata, NetworkError};
 use pixels::{CorsStatus, Image, ImageMetadata, PixelFormat, load_from_memory};
+use profile_traits::mem::{Report, ReportKind};
+use profile_traits::path;
 use servo_config::pref;
 use servo_url::{ImmutableOrigin, ServoUrl};
 use webrender_api::units::DeviceIntSize;
@@ -98,6 +104,7 @@ type ImageKey = (ServoUrl, ImmutableOrigin, Option<CorsSettings>);
 
 // Represents all the currently pending loads/decodings. For
 // performance reasons, loads are indexed by a dedicated load key.
+#[derive(MallocSizeOf)]
 struct AllPendingLoads {
     // The loads, indexed by a load key. Used during most operations,
     // for performance reasons.
@@ -180,6 +187,7 @@ enum CacheResult<'a> {
 /// Images that fail to load (due to network or decode
 /// failure) are still stored here, so that they aren't
 /// fetched again.
+#[derive(MallocSizeOf)]
 struct CompletedLoad {
     image_response: ImageResponse,
     id: PendingImageId,
@@ -197,9 +205,10 @@ struct DecoderMsg {
     image: Option<Image>,
 }
 
+#[derive(MallocSizeOf)]
 enum ImageBytes {
     InProgress(Vec<u8>),
-    Complete(Arc<Vec<u8>>),
+    Complete(#[conditional_malloc_size_of] Arc<Vec<u8>>),
 }
 
 impl ImageBytes {
@@ -234,6 +243,7 @@ impl ImageBytes {
 // A key used to communicate during loading.
 type LoadKey = PendingImageId;
 
+#[derive(MallocSizeOf)]
 struct LoadKeyGenerator {
     counter: u64,
 }
@@ -257,6 +267,7 @@ enum LoadResult {
 
 /// Represents an image that is either being loaded
 /// by the resource thread, or decoded by a worker thread.
+#[derive(MallocSizeOf)]
 struct PendingLoad {
     /// The bytes loaded so far. Reset to an empty vector once loading
     /// is complete and the buffer has been transmitted to the decoder.
@@ -315,6 +326,7 @@ impl PendingLoad {
 // ======================================================================
 // Image cache implementation.
 // ======================================================================
+#[derive(MallocSizeOf)]
 struct ImageCacheStore {
     // Images that are loading over network, or decoding.
     pending_loads: AllPendingLoads,
@@ -323,12 +335,14 @@ struct ImageCacheStore {
     completed_loads: HashMap<ImageKey, CompletedLoad>,
 
     // The placeholder image used when an image fails to load
+    #[conditional_malloc_size_of]
     placeholder_image: Arc<Image>,
 
     // The URL used for the placeholder image
     placeholder_url: ServoUrl,
 
     // Cross-process compositor API instance.
+    #[ignore_malloc_size_of = "Channel from another crate"]
     compositor_api: CrossProcessCompositorApi,
 }
 
@@ -439,6 +453,22 @@ impl ImageCache for ImageCacheImpl {
         }
     }
 
+    fn memory_report(&self, prefix: &str) -> Report {
+        let seen_pointer =
+            move |ptr| SEEN_POINTERS.with(|pointers| !pointers.borrow_mut().insert(ptr));
+        let mut ops = MallocSizeOfOps::new(
+            servo_allocator::usable_size,
+            None,
+            Some(Box::new(seen_pointer)),
+        );
+        let size = self.store.lock().unwrap().size_of(&mut ops);
+        Report {
+            path: path![prefix, "image-cache"],
+            kind: ReportKind::ExplicitSystemHeapSize,
+            size,
+        }
+    }
+
     fn get_image(
         &self,
         url: ServoUrl,
@@ -641,3 +671,7 @@ impl ImageCacheImpl {
         warn!("Couldn't find cached entry for listener {:?}", id);
     }
 }
+
+thread_local!(static SEEN_POINTERS: LazyCell<RefCell<HashSet<*const c_void>>> = const {
+    LazyCell::new(|| RefCell::new(HashSet::new()))
+});