|
| 1 | +use std::sync::Arc; |
| 2 | + |
| 3 | +use rayon::prelude::*; |
| 4 | +use rspack_cacheable::{cacheable, with::AsPreset}; |
| 5 | +use rspack_error::Result; |
| 6 | +use rspack_sources::BoxSource; |
| 7 | +use rustc_hash::FxHashMap; |
| 8 | + |
| 9 | +use super::{ |
| 10 | + super::{codec::CacheCodec, storage::Storage}, |
| 11 | + Occasion, |
| 12 | +}; |
| 13 | +use crate::RayonConsumer; |
| 14 | + |
| 15 | +pub const SCOPE: &str = "occasion_minimize"; |
| 16 | + |
| 17 | +#[cacheable] |
| 18 | +struct Entry { |
| 19 | + #[cacheable(with=AsPreset)] |
| 20 | + pub source: BoxSource, |
| 21 | + pub extracted_comments: Option<ExtractedCommentsEntry>, |
| 22 | +} |
| 23 | + |
| 24 | +#[cacheable] |
| 25 | +struct ExtractedCommentsEntry { |
| 26 | + #[cacheable(with=AsPreset)] |
| 27 | + pub source: BoxSource, |
| 28 | + pub comments_file_name: String, |
| 29 | +} |
| 30 | + |
| 31 | +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] |
| 32 | +pub struct MinimizeCacheKey(u64); |
| 33 | + |
| 34 | +impl MinimizeCacheKey { |
| 35 | + pub fn new(hash: u64) -> Self { |
| 36 | + Self(hash) |
| 37 | + } |
| 38 | + |
| 39 | + fn to_bytes(self) -> Vec<u8> { |
| 40 | + self.0.to_ne_bytes().to_vec() |
| 41 | + } |
| 42 | + |
| 43 | + fn from_bytes(bytes: &[u8]) -> Option<Self> { |
| 44 | + <[u8; 8]>::try_from(bytes) |
| 45 | + .ok() |
| 46 | + .map(|b| Self(u64::from_ne_bytes(b))) |
| 47 | + } |
| 48 | +} |
| 49 | + |
| 50 | +#[derive(Debug, Default)] |
| 51 | +pub struct MinimizePersistentCacheArtifact { |
| 52 | + entries: FxHashMap<MinimizeCacheKey, CachedMinimizeEntry>, |
| 53 | + /// Keys of entries that were added during this build and need to be persisted. |
| 54 | + dirty_keys: Vec<MinimizeCacheKey>, |
| 55 | +} |
| 56 | + |
| 57 | +#[derive(Debug, Clone)] |
| 58 | +pub struct CachedMinimizeEntry { |
| 59 | + pub source: BoxSource, |
| 60 | + pub extracted_comments: Option<CachedExtractedComments>, |
| 61 | +} |
| 62 | + |
| 63 | +#[derive(Debug, Clone)] |
| 64 | +pub struct CachedExtractedComments { |
| 65 | + pub source: BoxSource, |
| 66 | + pub comments_file_name: String, |
| 67 | +} |
| 68 | + |
| 69 | +impl MinimizePersistentCacheArtifact { |
| 70 | + pub fn get(&self, key: MinimizeCacheKey) -> Option<&CachedMinimizeEntry> { |
| 71 | + self.entries.get(&key) |
| 72 | + } |
| 73 | + |
| 74 | + pub fn insert(&mut self, key: MinimizeCacheKey, entry: CachedMinimizeEntry) { |
| 75 | + self.dirty_keys.push(key); |
| 76 | + self.entries.insert(key, entry); |
| 77 | + } |
| 78 | +} |
| 79 | + |
| 80 | +#[derive(Debug)] |
| 81 | +pub struct MinimizeOccasion { |
| 82 | + codec: Arc<CacheCodec>, |
| 83 | +} |
| 84 | + |
| 85 | +impl MinimizeOccasion { |
| 86 | + pub fn new(codec: Arc<CacheCodec>) -> Self { |
| 87 | + Self { codec } |
| 88 | + } |
| 89 | +} |
| 90 | + |
| 91 | +#[async_trait::async_trait] |
| 92 | +impl Occasion for MinimizeOccasion { |
| 93 | + type Artifact = MinimizePersistentCacheArtifact; |
| 94 | + |
| 95 | + #[tracing::instrument(name = "Cache::Occasion::Minimize::reset", skip_all)] |
| 96 | + fn reset(&self, storage: &mut dyn Storage) { |
| 97 | + storage.reset(SCOPE); |
| 98 | + } |
| 99 | + |
| 100 | + #[tracing::instrument(name = "Cache::Occasion::Minimize::save", skip_all)] |
| 101 | + fn save(&self, storage: &mut dyn Storage, artifact: &MinimizePersistentCacheArtifact) { |
| 102 | + // Only persist entries that were added during this build. |
| 103 | + artifact |
| 104 | + .dirty_keys |
| 105 | + .par_iter() |
| 106 | + .filter_map(|key| { |
| 107 | + let entry = artifact.entries.get(key)?; |
| 108 | + let storage_entry = Entry { |
| 109 | + source: entry.source.clone(), |
| 110 | + extracted_comments: entry |
| 111 | + .extracted_comments |
| 112 | + .as_ref() |
| 113 | + .map(|ec| ExtractedCommentsEntry { |
| 114 | + source: ec.source.clone(), |
| 115 | + comments_file_name: ec.comments_file_name.clone(), |
| 116 | + }), |
| 117 | + }; |
| 118 | + match self.codec.encode(&storage_entry) { |
| 119 | + Ok(bytes) => Some((key.to_bytes(), bytes)), |
| 120 | + Err(err) => { |
| 121 | + tracing::warn!("minimize persistent cache encode failed: {:?}", err); |
| 122 | + None |
| 123 | + } |
| 124 | + } |
| 125 | + }) |
| 126 | + .consume(|(key, bytes)| { |
| 127 | + storage.set(SCOPE, key, bytes); |
| 128 | + }); |
| 129 | + |
| 130 | + tracing::debug!( |
| 131 | + "saved {} minimize persistent cache entries", |
| 132 | + artifact.dirty_keys.len() |
| 133 | + ); |
| 134 | + } |
| 135 | + |
| 136 | + #[tracing::instrument(name = "Cache::Occasion::Minimize::recovery", skip_all)] |
| 137 | + async fn recovery(&self, storage: &dyn Storage) -> Result<MinimizePersistentCacheArtifact> { |
| 138 | + let items = storage.load(SCOPE).await?; |
| 139 | + let mut entries = FxHashMap::default(); |
| 140 | + entries.reserve(items.len()); |
| 141 | + |
| 142 | + for (key, value) in items { |
| 143 | + let Some(key) = MinimizeCacheKey::from_bytes(&key) else { |
| 144 | + tracing::warn!("minimize persistent cache key has invalid length"); |
| 145 | + continue; |
| 146 | + }; |
| 147 | + match self.codec.decode::<Entry>(&value) { |
| 148 | + Ok(entry) => { |
| 149 | + entries.insert( |
| 150 | + key, |
| 151 | + CachedMinimizeEntry { |
| 152 | + source: entry.source, |
| 153 | + extracted_comments: entry.extracted_comments.map(|ec| CachedExtractedComments { |
| 154 | + source: ec.source, |
| 155 | + comments_file_name: ec.comments_file_name, |
| 156 | + }), |
| 157 | + }, |
| 158 | + ); |
| 159 | + } |
| 160 | + Err(err) => { |
| 161 | + tracing::warn!("minimize persistent cache decode failed: {:?}", err); |
| 162 | + } |
| 163 | + } |
| 164 | + } |
| 165 | + |
| 166 | + tracing::debug!( |
| 167 | + "recovered {} minimize persistent cache entries", |
| 168 | + entries.len() |
| 169 | + ); |
| 170 | + Ok(MinimizePersistentCacheArtifact { |
| 171 | + entries, |
| 172 | + dirty_keys: Vec::new(), |
| 173 | + }) |
| 174 | + } |
| 175 | +} |
0 commit comments