Skip to main content

net/
disk_cache.rs

1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
4
5use std::collections::VecDeque;
6use std::path::PathBuf;
7use std::sync::Arc;
8use std::time::{SystemTime, UNIX_EPOCH};
9
10use log::error;
11use malloc_size_of_derive::MallocSizeOf;
12use rusqlite::Row;
13use sea_query::{ColumnDef, Expr, ExprTrait, Iden, OnConflict, Query, SqliteQueryBuilder, Table};
14use sea_query_rusqlite::RusqliteBinder;
15use servo_config::pref;
16use servo_url::ServoUrl;
17use tokio::sync::{Mutex as TokioMutex, RwLock as TokioRwLock};
18
19use crate::http_cache::{
20    CacheEntry, CacheKey, CachedResource, HttpCacheAssignment, MemoryCacheLifecycle,
21};
22
23#[derive(MallocSizeOf)]
24pub(crate) struct DiskCacheMetadata {
25    key: CacheKey,
26    /// The size of the serialization or the exact size from the cache metadata.
27    size: usize,
28}
29
30impl From<&Row<'_>> for DiskCacheMetadata {
31    fn from(row: &Row) -> Self {
32        let s: String = row.get_unwrap("key");
33        Self {
34            key: CacheKey::from_url(ServoUrl::parse(&s).unwrap()),
35            size: row.get_unwrap("size"),
36        }
37    }
38}
39
40/// This data structure will be per [`HttpCacheAssignment`]. Currently, we only store HttpCacheAssignment::Public and otherwise return zero.
41/// As this data structure stores the state of the cache, if other instances share the same sqlite database, special care has to be taken
42/// to ensure conistency.
43#[derive(MallocSizeOf)]
44struct DiskCacheInner {
45    entries: VecDeque<DiskCacheMetadata>,
46    size: usize,
47    #[ignore_malloc_size_of = "Find a better way"]
48    db: rusqlite::Connection,
49    cache_assignment: HttpCacheAssignment,
50}
51
52#[derive(MallocSizeOf)]
53/// A struct representing the disk cache.
54pub(crate) struct DiskCache {
55    path: PathBuf,
56    max_size: usize,
57    // the non constant data.
58    inner: TokioMutex<DiskCacheInner>,
59}
60
61/// Identifications for sea_query of the cache table.
62enum DiskCacheTable {
63    Table,
64    Key,
65    Data,
66    Size,
67    InsertionTimestamp,
68}
69
70// Mapping between Enum variant and its corresponding string value
71impl Iden for DiskCacheTable {
72    fn unquoted(&self) -> &str {
73        match self {
74            DiskCacheTable::Table => "disk_cache",
75            DiskCacheTable::Key => "key",
76            DiskCacheTable::Data => "data",
77            DiskCacheTable::Size => "size",
78            DiskCacheTable::InsertionTimestamp => "insertion_timestamp",
79        }
80    }
81}
82
83/// Get the storage path out of `network_http_disk_cache` preference and `temporary_storage` option.
84fn storage_dir() -> Option<PathBuf> {
85    let disk_storage_path = pref!(network_http_disk_cache);
86    match (
87        servo_config::opts::get().temporary_storage,
88        disk_storage_path.is_empty(),
89    ) {
90        (false, false) => Some(disk_storage_path.into()),
91        (true, true) => {
92            let tmp_dir = tempfile::tempdir().unwrap();
93            let mut path = tmp_dir.path().to_path_buf();
94            path.set_file_name("cache.sqlite3");
95            Some(path)
96        },
97        (true, false) => {
98            error!(
99                "Temporary storage cannot be set with explicit disk storage path. Disabling http_disk_cache"
100            );
101            None
102        },
103        (false, true) => None,
104    }
105}
106
107impl DiskCache {
108    /// Creates a new [`DiskCache`] if the preference if set.
109    /// Creates the sqlite table if it does not exist and starts the db connection.
110    /// TODO: Implement WAL and other sqlite pragma.
111    pub(crate) fn new(
112        cache_assignment: HttpCacheAssignment,
113    ) -> (Option<Arc<DiskCache>>, MemoryCacheLifecycle) {
114        // For private browsing we currently do not want to store any disk cache.
115        let disk_cache_path = storage_dir();
116
117        if let Some(disk_cache_path) = disk_cache_path &&
118            cache_assignment == HttpCacheAssignment::Public
119        {
120            let Ok(max_disk_cache_size) = pref!(network_http_disk_cache_size).try_into() else {
121                return (None, MemoryCacheLifecycle::empty());
122            };
123
124            let Ok(db) = rusqlite::Connection::open(&disk_cache_path) else {
125                error!("Could not open disk cache database");
126                return (None, MemoryCacheLifecycle::empty());
127            };
128
129            let _ = db.execute("PRAGMA journal_mode = WAL;", ());
130            let query = Table::create()
131                .table(DiskCacheTable::Table)
132                .if_not_exists()
133                .col(
134                    ColumnDef::new(DiskCacheTable::Key)
135                        .text()
136                        .not_null()
137                        .primary_key(),
138                )
139                .col(ColumnDef::new(DiskCacheTable::Data).blob().not_null())
140                .col(ColumnDef::new(DiskCacheTable::Size).integer().not_null())
141                .col(
142                    ColumnDef::new(DiskCacheTable::InsertionTimestamp)
143                        .integer()
144                        .not_null(),
145                )
146                .build(SqliteQueryBuilder);
147            if let Err(e) = db.execute(query.as_str(), ()) {
148                error!("Could not create table. DB Error {:?}", e);
149                return (None, MemoryCacheLifecycle::empty());
150            }
151
152            let (query, values) = Query::select()
153                .columns([DiskCacheTable::Key, DiskCacheTable::Size])
154                .from(DiskCacheTable::Table)
155                .build_rusqlite(SqliteQueryBuilder);
156
157            let (entries, size) = {
158                let Ok(mut st) = db.prepare(query.as_str()) else {
159                    error!("Could not get disk data");
160                    return (None, MemoryCacheLifecycle::empty());
161                };
162                let entries = st
163                    .query_map(&*values.as_params(), |row| Ok(DiskCacheMetadata::from(row)))
164                    .unwrap()
165                    .map(|entry| entry.unwrap())
166                    .collect::<VecDeque<_>>();
167
168                let size = entries.iter().map(|entry| entry.size).sum();
169                (entries, size)
170            };
171            let inner = DiskCacheInner {
172                entries,
173                size,
174                db,
175                cache_assignment,
176            };
177            let disk_cache_data = std::sync::Arc::new(DiskCache {
178                inner: TokioMutex::new(inner),
179                path: disk_cache_path,
180                max_size: max_disk_cache_size,
181            });
182
183            (
184                Some(disk_cache_data.clone()),
185                MemoryCacheLifecycle {
186                    disk_cache: Some(disk_cache_data),
187                },
188            )
189        } else {
190            (None, MemoryCacheLifecycle::empty())
191        }
192    }
193
194    /// Restores a cache entry from the disk if it exists.
195    /// Deletes the entry from the disk cache
196    #[servo_tracing::instrument(skip(self))]
197    pub(crate) async fn get(&self, key: CacheKey) -> Option<Arc<TokioRwLock<Vec<CachedResource>>>> {
198        let bytes = {
199            // we lock the metadata before we update the sqlite database so that
200            // the database and metadata are consistent when this lock is released.
201            let mut inner = self.inner.lock().await;
202            let (bytes, new_size) = {
203                let _span = profile_traits::trace_span!("query disk cache").entered();
204                let (query, query_values) = Query::select()
205                    .columns([DiskCacheTable::Data])
206                    .from(DiskCacheTable::Table)
207                    .and_where(Expr::col(DiskCacheTable::Key).eq(key.as_ref()))
208                    .build_rusqlite(SqliteQueryBuilder);
209                let (delete, delete_values) = Query::delete()
210                    .from_table(DiskCacheTable::Table)
211                    .and_where(Expr::col(DiskCacheTable::Key).eq(key.as_ref()))
212                    .build_rusqlite(SqliteQueryBuilder);
213
214                let mut st = inner.db.prepare(query.as_str()).ok()?;
215                let data: Vec<u8> = st
216                    .query_one(&*query_values.as_params(), |row| Ok(row.get_unwrap("data")))
217                    .ok()?;
218
219                if inner
220                    .db
221                    .execute(delete.as_str(), &*delete_values.as_params())
222                    .is_err()
223                {
224                    error!("Could not delete cached data from disk");
225                    return None;
226                }
227
228                (data, self.get_disk_cache_total_size(&inner.db))
229            };
230
231            {
232                // update the metadata
233                let entry_index = inner
234                    .entries
235                    .iter()
236                    .position(|metadata| metadata.key == key);
237                if let Some(entry_index) = entry_index {
238                    inner.entries.remove(entry_index);
239                }
240                if let Some(new_size) = new_size {
241                    inner.size = new_size;
242                } else {
243                    error!("Could not get disk cache size");
244                }
245            }
246            bytes
247        };
248        let _span = profile_traits::trace_span!("deserialize cache request").entered();
249        let Ok(value) = postcard::from_bytes(&bytes) else {
250            error!("Could not deserialize cached resource");
251            return None;
252        };
253        let deserialized_vec_cached_response = std::sync::Arc::new(TokioRwLock::new(value));
254
255        Some(deserialized_vec_cached_response)
256    }
257
258    /// Stores a [`CacheEntry`]` to disk.
259    #[servo_tracing::instrument(skip(self))]
260    pub(crate) async fn store(&self, key: CacheKey, entry: CacheEntry) {
261        let entry = entry.read().await;
262        let data_to_serialize: Vec<&CachedResource> = entry
263            .iter()
264            .filter(|cached_resource| cached_resource.is_done())
265            .collect();
266        let Ok(data) = postcard::to_stdvec(&*data_to_serialize) else {
267            error!("Could not deserialize value");
268            return;
269        };
270
271        {
272            let mut inner = self.inner.lock().await;
273            let data_size = data.len();
274
275            let timestamp = SystemTime::now()
276                .duration_since(UNIX_EPOCH)
277                .map(|duration| duration.as_secs())
278                .unwrap_or(0);
279            let (query, params) = Query::insert()
280                .into_table(DiskCacheTable::Table)
281                .columns([
282                    DiskCacheTable::Key,
283                    DiskCacheTable::Data,
284                    DiskCacheTable::Size,
285                    DiskCacheTable::InsertionTimestamp,
286                ])
287                .on_conflict(
288                    OnConflict::column(DiskCacheTable::Key)
289                        .update_columns([
290                            DiskCacheTable::Data,
291                            DiskCacheTable::Data,
292                            DiskCacheTable::Size,
293                            DiskCacheTable::InsertionTimestamp,
294                        ])
295                        .to_owned(),
296                )
297                .values_panic([
298                    key.as_ref().into(),
299                    data.into(),
300                    (data_size as u32).into(),
301                    timestamp.into(),
302                ])
303                .build_rusqlite(SqliteQueryBuilder);
304
305            if let Err(e) = inner.db.execute(query.as_str(), &*params.as_params()) {
306                error!("Could not insert cache data. Error {}", e);
307            }
308            inner.entries.push_back(DiskCacheMetadata {
309                key,
310                size: data_size,
311            });
312            if let Some(new_cache_size) = self.get_disk_cache_total_size(&inner.db) {
313                inner.size = new_cache_size;
314            }
315        }
316        self.delete_until_cache_size().await;
317    }
318
319    /// Deletes data from the cache until the size is <= max_size
320    #[servo_tracing::instrument(skip(self))]
321    async fn delete_until_cache_size(&self) {
322        let mut inner = self.inner.lock().await;
323        let mut keys_to_delete = vec![];
324        while self.max_size < inner.size {
325            if let Some(metadata) = inner.entries.pop_back() {
326                keys_to_delete.push(metadata.key);
327                inner.size -= metadata.size;
328            }
329        }
330
331        let keys_ref = keys_to_delete.iter().map(|key| key.as_ref());
332        let (query, values) = Query::delete()
333            .from_table(DiskCacheTable::Table)
334            .and_where(Expr::col(DiskCacheTable::Key).is_in(keys_ref))
335            .build_rusqlite(SqliteQueryBuilder);
336
337        if inner
338            .db
339            .execute(query.as_str(), &*values.as_params())
340            .is_err()
341        {
342            error!("Could not delete old disk cache entries");
343        }
344    }
345
346    /// Queries the current disk cache size from the sql database.
347    #[servo_tracing::instrument(skip(self))]
348    fn get_disk_cache_total_size(&self, conn: &rusqlite::Connection) -> Option<usize> {
349        let (size, size_values) = Query::select()
350            .expr(Expr::col(DiskCacheTable::Size).sum())
351            .from(DiskCacheTable::Table)
352            .build_rusqlite(SqliteQueryBuilder);
353        let Ok(mut st) = conn.prepare(size.as_str()) else {
354            return None;
355        };
356
357        // According to the sqlite documentation we will return NULL on an empty table.
358        let query_result =
359            st.query_one(&*size_values.as_params(), |row| Ok(row.get(0).unwrap_or(0)));
360        if let Err(query_result) = query_result {
361            error!("Could nto get new sum size {}", query_result);
362            None
363        } else {
364            query_result.ok()
365        }
366    }
367
368    /// Clears the disk cache.
369    /// Should only be called in sync context and will panic.
370    #[servo_tracing::instrument(skip(self))]
371    pub(crate) fn clear(&self) {
372        let mut inner = self.inner.blocking_lock();
373        let (query, params) = Query::delete()
374            .from_table(DiskCacheTable::Table)
375            .build_rusqlite(SqliteQueryBuilder);
376        if inner
377            .db
378            .execute(query.as_str(), &*params.as_params())
379            .is_err()
380        {
381            error!("Could not clear disk cache");
382        }
383        inner.entries.clear();
384        inner.size = 0;
385    }
386}