Add a basic caching mechanism for ImageKeys. (#37369)

This creates a new method in shared/compositing/lib to generate image
keys that are send over the webview. This does not immediately return
the keys but goes over the constellation to receive the keys from the
IOCompositor. To make this more efficient, we now cache the keys in
image_cache in a simple FIFO order. The old blocking method stays intact
for now but got renamed to make the blocking clear.
The blocking calls that are left are in:
- `components/canvas/canvas_data.rs`
- `components/script/dom/htmlmediaelement.rs`

Testing: WPT tests should cover this as this doesn't change any
functionality.
Fixes: Was mentioned in
https://github.com/servo/servo/issues/37161#issuecomment-2915750051 and
part of https://github.com/servo/servo/issues/37086

---------

Signed-off-by: Narfinger <Narfinger@users.noreply.github.com>
Signed-off-by: gterzian <2792687+gterzian@users.noreply.github.com>
Co-authored-by: gterzian <2792687+gterzian@users.noreply.github.com>

components/net/image_cache.rs

@@ -2,8 +2,9 @@
  * 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::cmp::min;
 use std::collections::hash_map::Entry::{Occupied, Vacant};
+use std::collections::{HashMap, VecDeque};
 use std::sync::{Arc, Mutex};
 use std::{mem, thread};
 
@@ -11,7 +12,7 @@ use base::id::PipelineId;
 use compositing_traits::{CrossProcessCompositorApi, ImageUpdate, SerializableImageData};
 use imsz::imsz_from_reader;
 use ipc_channel::ipc::{IpcSender, IpcSharedMemory};
-use log::{debug, warn};
+use log::{debug, error, warn};
 use malloc_size_of::{MallocSizeOf as MallocSizeOfTrait, MallocSizeOfOps};
 use malloc_size_of_derive::MallocSizeOf;
 use mime::Mime;
@@ -29,7 +30,9 @@ use resvg::{tiny_skia, usvg};
 use servo_config::pref;
 use servo_url::{ImmutableOrigin, ServoUrl};
 use webrender_api::units::DeviceIntSize;
-use webrender_api::{ImageDescriptor, ImageDescriptorFlags, ImageFormat};
+use webrender_api::{
+    ImageDescriptor, ImageDescriptorFlags, ImageFormat, ImageKey as WebRenderImageKey,
+};
 
 use crate::resource_thread::CoreResourceThreadPool;
 
@@ -95,6 +98,8 @@ fn decode_bytes_sync(
     DecoderMsg { key, image }
 }
 
+/// This will block on getting an ImageKey
+/// but that is ok because it is done once upon start-up of a script-thread.
 fn get_placeholder_image(
     compositor_api: &CrossProcessCompositorApi,
     data: &[u8],
@@ -102,11 +107,18 @@ fn get_placeholder_image(
     let mut image = load_from_memory(data, CorsStatus::Unsafe)
         .or_else(|| load_from_memory(FALLBACK_RIPPY, CorsStatus::Unsafe))
         .expect("load fallback image failed");
-    set_webrender_image_key(compositor_api, &mut image);
+    let image_key = compositor_api
+        .generate_image_key_blocking()
+        .expect("Could not generate image key");
+    set_webrender_image_key(compositor_api, &mut image, image_key);
     Arc::new(image)
 }
 
-fn set_webrender_image_key(compositor_api: &CrossProcessCompositorApi, image: &mut RasterImage) {
+fn set_webrender_image_key(
+    compositor_api: &CrossProcessCompositorApi,
+    image: &mut RasterImage,
+    image_key: WebRenderImageKey,
+) {
     if image.id.is_some() {
         return;
     }
@@ -146,11 +158,9 @@ fn set_webrender_image_key(compositor_api: &CrossProcessCompositorApi, image: &m
         offset: 0,
         flags,
     };
-    if let Some(image_key) = compositor_api.generate_image_key() {
-        let data = SerializableImageData::Raw(IpcSharedMemory::from_bytes(&bytes));
-        compositor_api.add_image(image_key, descriptor, data);
-        image.id = Some(image_key);
-    }
+    let data = SerializableImageData::Raw(IpcSharedMemory::from_bytes(&bytes));
+    compositor_api.add_image(image_key, descriptor, data);
+    image.id = Some(image_key);
 }
 
 // ======================================================================
@@ -404,6 +414,51 @@ struct RasterizationTask {
     result: Option<RasterImage>,
 }
 
+/// Used for storing images that do not have a `WebRenderImageKey` yet.
+#[derive(Debug, MallocSizeOf)]
+enum PendingKey {
+    RasterImage((LoadKey, RasterImage)),
+    Svg((LoadKey, RasterImage, DeviceIntSize)),
+}
+
+/// The state of the `WebRenderImageKey`` cache
+#[derive(Debug, MallocSizeOf)]
+enum KeyCacheState {
+    /// We already requested a batch of keys.
+    PendingBatch,
+    /// We have some keys in the cache.
+    Ready(Vec<WebRenderImageKey>),
+}
+
+impl KeyCacheState {
+    fn size(&self) -> usize {
+        match self {
+            KeyCacheState::PendingBatch => 0,
+            KeyCacheState::Ready(items) => items.len(),
+        }
+    }
+}
+
+/// As getting new keys takes a round trip over the constellation, we keep a small cache of them.
+/// Additionally, this cache will store image resources that do not have a key yet because those
+/// are needed to complete the load.
+#[derive(MallocSizeOf)]
+struct KeyCache {
+    /// A cache of `WebRenderImageKey`.
+    cache: KeyCacheState,
+    /// These images are loaded but have no key assigned to yet.
+    images_pending_keys: VecDeque<PendingKey>,
+}
+
+impl KeyCache {
+    fn new() -> Self {
+        KeyCache {
+            cache: KeyCacheState::Ready(Vec::new()),
+            images_pending_keys: VecDeque::new(),
+        }
+    }
+}
+
 /// ## Image cache implementation.
 #[derive(MallocSizeOf)]
 struct ImageCacheStore {
@@ -433,33 +488,126 @@ struct ImageCacheStore {
     /// Cross-process compositor API instance.
     #[ignore_malloc_size_of = "Channel from another crate"]
     compositor_api: CrossProcessCompositorApi,
+
+    // The PipelineId will initially be None because the constructed cache is not associated
+    // with any pipeline yet. This will happen later by way of `create_new_image_cache`.
+    pipeline_id: Option<PipelineId>,
+
+    /// Main struct to handle the cache of `WebRenderImageKey` and
+    /// images that do not have a key yet.
+    key_cache: KeyCache,
 }
 
 impl ImageCacheStore {
-    // Change state of a url from pending -> loaded.
-    fn complete_load(&mut self, key: LoadKey, mut load_result: LoadResult) {
+    /// Finishes loading the image by setting the WebRenderImageKey and calling `compete_load` or `complete_load_svg`.
+    fn set_key_and_finish_load(&mut self, pending_image: PendingKey, image_key: WebRenderImageKey) {
+        match pending_image {
+            PendingKey::RasterImage((pending_id, mut raster_image)) => {
+                set_webrender_image_key(&self.compositor_api, &mut raster_image, image_key);
+                self.complete_load(pending_id, LoadResult::LoadedRasterImage(raster_image));
+            },
+            PendingKey::Svg((pending_id, mut raster_image, requested_size)) => {
+                set_webrender_image_key(&self.compositor_api, &mut raster_image, image_key);
+                self.complete_load_svg(raster_image, pending_id, requested_size);
+            },
+        }
+    }
+
+    /// If a key is available the image will be immediately loaded, otherwise it will load then the next batch of
+    /// keys is received. Only call this if the image does not have a `LoadKey` yet.
+    fn load_image_with_keycache(&mut self, pending_image: PendingKey) {
+        if let Some(pipeline_id) = self.pipeline_id {
+            match self.key_cache.cache {
+                KeyCacheState::PendingBatch => {
+                    self.key_cache.images_pending_keys.push_back(pending_image);
+                },
+                KeyCacheState::Ready(ref mut cache) => match cache.pop() {
+                    Some(image_key) => {
+                        self.set_key_and_finish_load(pending_image, image_key);
+                    },
+                    None => {
+                        self.key_cache.images_pending_keys.push_back(pending_image);
+                        self.compositor_api.generate_image_key_async(pipeline_id);
+                        self.key_cache.cache = KeyCacheState::PendingBatch
+                    },
+                },
+            }
+        } else {
+            error!("No pipeline id for this image key cache.");
+        }
+    }
+
+    /// Insert received keys into the cache and complete the loading of images.
+    fn insert_keys_and_load_images(&mut self, image_keys: Vec<WebRenderImageKey>) {
+        if let KeyCacheState::PendingBatch = self.key_cache.cache {
+            self.key_cache.cache = KeyCacheState::Ready(image_keys);
+            let len = min(
+                self.key_cache.cache.size(),
+                self.key_cache.images_pending_keys.len(),
+            );
+            let images = self
+                .key_cache
+                .images_pending_keys
+                .drain(0..len)
+                .collect::<Vec<PendingKey>>();
+            for key in images {
+                self.load_image_with_keycache(key);
+            }
+            if !self.key_cache.images_pending_keys.is_empty() {
+                self.compositor_api
+                    .generate_image_key_async(self.pipeline_id.unwrap());
+                self.key_cache.cache = KeyCacheState::PendingBatch
+            }
+        } else {
+            unreachable!("A batch was received while we didn't request one")
+        }
+    }
+
+    /// Complete the loading the of the rasterized svg image. This needs the `RasterImage` to
+    /// already have a `WebRenderImageKey`.
+    fn complete_load_svg(
+        &mut self,
+        rasterized_image: RasterImage,
+        pending_image_id: PendingImageId,
+        requested_size: DeviceIntSize,
+    ) {
+        let listeners = {
+            self.rasterized_vector_images
+                .get_mut(&(pending_image_id, requested_size))
+                .map(|task| {
+                    task.result = Some(rasterized_image);
+                    std::mem::take(&mut task.listeners)
+                })
+                .unwrap_or_default()
+        };
+
+        for (pipeline_id, sender) in listeners {
+            let _ = sender.send(ImageCacheResponseMessage::VectorImageRasterizationComplete(
+                RasterizationCompleteResponse {
+                    pipeline_id,
+                    image_id: pending_image_id,
+                    requested_size,
+                },
+            ));
+        }
+    }
+
+    /// The rest of complete load. This requires that images have a valid `WebRenderImageKey`.
+    fn complete_load(&mut self, key: LoadKey, load_result: LoadResult) {
         debug!("Completed decoding for {:?}", load_result);
         let pending_load = match self.pending_loads.remove(&key) {
             Some(load) => load,
             None => return,
         };
 
-        match load_result {
-            LoadResult::LoadedRasterImage(ref mut raster_image) => {
-                set_webrender_image_key(&self.compositor_api, raster_image)
-            },
-            LoadResult::LoadedVectorImage(ref vector_image) => {
-                self.vector_images.insert(key, vector_image.clone());
-            },
-            LoadResult::PlaceholderLoaded(..) | LoadResult::None => {},
-        }
-
         let url = pending_load.final_url.clone();
         let image_response = match load_result {
             LoadResult::LoadedRasterImage(raster_image) => {
+                assert!(raster_image.id.is_some());
                 ImageResponse::Loaded(Image::Raster(Arc::new(raster_image)), url.unwrap())
             },
             LoadResult::LoadedVectorImage(vector_image) => {
+                self.vector_images.insert(key, vector_image.clone());
                 let natural_dimensions = vector_image.svg_tree.size().to_int_size();
                 let metadata = ImageMetadata {
                     width: natural_dimensions.width(),
@@ -523,7 +671,10 @@ impl ImageCacheStore {
     fn handle_decoder(&mut self, msg: DecoderMsg) {
         let image = match msg.image {
             None => LoadResult::None,
-            Some(DecodedImage::Raster(raster_image)) => LoadResult::LoadedRasterImage(raster_image),
+            Some(DecodedImage::Raster(raster_image)) => {
+                self.load_image_with_keycache(PendingKey::RasterImage((msg.key, raster_image)));
+                return;
+            },
             Some(DecodedImage::Vector(vector_image_data)) => {
                 LoadResult::LoadedVectorImage(vector_image_data)
             },
@@ -559,7 +710,9 @@ impl ImageCache for ImageCacheImpl {
                 rasterized_vector_images: HashMap::new(),
                 placeholder_image: get_placeholder_image(&compositor_api, &rippy_data),
                 placeholder_url: ServoUrl::parse("chrome://resources/rippy.png").unwrap(),
-                compositor_api,
+                compositor_api: compositor_api.clone(),
+                pipeline_id: None,
+                key_cache: KeyCache::new(),
             })),
             thread_pool: Arc::new(CoreResourceThreadPool::new(
                 thread_count,
@@ -623,7 +776,7 @@ impl ImageCache for ImageCacheImpl {
             }
         }
 
-        let decoded = {
+        let (key, decoded) = {
             let result = store
                 .pending_loads
                 .get_cached(url.clone(), origin.clone(), cors_setting);
@@ -631,11 +784,14 @@ impl ImageCache for ImageCacheImpl {
                 CacheResult::Hit(key, pl) => match (&pl.result, &pl.metadata) {
                     (&Some(Ok(_)), _) => {
                         debug!("Sync decoding {} ({:?})", url, key);
-                        decode_bytes_sync(
+                        (
                             key,
-                            pl.bytes.as_slice(),
-                            pl.cors_status,
-                            pl.content_type.clone(),
+                            decode_bytes_sync(
+                                key,
+                                pl.bytes.as_slice(),
+                                pl.cors_status,
+                                pl.content_type.clone(),
+                            ),
                         )
                     },
                     (&None, Some(meta)) => {
@@ -674,7 +830,8 @@ impl ImageCache for ImageCacheImpl {
                     is_placeholder,
                 })
             },
-            _ => ImageCacheResult::LoadError,
+            // Note: this happens if we are pending a batch of image keys.
+            _ => ImageCacheResult::Pending(key),
         }
     }
 
@@ -767,7 +924,7 @@ impl ImageCache for ImageCacheImpl {
                 height: tinyskia_requested_size.height(),
             };
 
-            let mut rasterized_image = RasterImage {
+            let rasterized_image = RasterImage {
                 metadata: ImageMetadata {
                     width: tinyskia_requested_size.width(),
                     height: tinyskia_requested_size.height(),
@@ -779,28 +936,12 @@ impl ImageCache for ImageCacheImpl {
                 cors_status: vector_image.cors_status,
             };
 
-            let listeners = {
-                let mut store = store.lock().unwrap();
-                set_webrender_image_key(&store.compositor_api, &mut rasterized_image);
-                store
-                    .rasterized_vector_images
-                    .get_mut(&(image_id, requested_size))
-                    .map(|task| {
-                        task.result = Some(rasterized_image);
-                        std::mem::take(&mut task.listeners)
-                    })
-                    .unwrap_or_default()
-            };
-
-            for (pipeline_id, sender) in listeners {
-                let _ = sender.send(ImageCacheResponseMessage::VectorImageRasterizationComplete(
-                    RasterizationCompleteResponse {
-                        pipeline_id,
-                        image_id,
-                        requested_size,
-                    },
-                ));
-            }
+            let mut store = store.lock().unwrap();
+            store.load_image_with_keycache(PendingKey::Svg((
+                image_id,
+                rasterized_image,
+                requested_size,
+            )));
         });
 
         None
@@ -905,6 +1046,7 @@ impl ImageCache for ImageCacheImpl {
 
     fn create_new_image_cache(
         &self,
+        pipeline_id: Option<PipelineId>,
         compositor_api: CrossProcessCompositorApi,
     ) -> Arc<dyn ImageCache> {
         let store = self.store.lock().unwrap();
@@ -919,10 +1061,17 @@ impl ImageCache for ImageCacheImpl {
                 compositor_api,
                 vector_images: HashMap::new(),
                 rasterized_vector_images: HashMap::new(),
+                key_cache: KeyCache::new(),
+                pipeline_id,
             })),
             thread_pool: self.thread_pool.clone(),
         })
     }
+
+    fn fill_key_cache_with_batch_of_keys(&self, image_keys: Vec<WebRenderImageKey>) {
+        let mut store = self.store.lock().unwrap();
+        store.insert_keys_and_load_images(image_keys);
+    }
 }
 
 impl Drop for ImageCacheStore {