1use std::cell::RefCell;
3use std::sync::{Arc, atomic::AtomicU64, atomic::Ordering::Relaxed};
4
5use crate::stvec::{Header, StorageVec};
6use crate::{BytePageSize, size::PAGE_SIZES};
7
8const CLASSES: usize = PAGE_SIZES.len();
10const _: () = assert!(CLASSES == BytePageSize::Unset as usize);
11
12const DEFAULT_PAGES_CACHE: [usize; CLASSES] = [128, 64, 64, 32, 16, 8, 16, 2, 1];
16
17thread_local! {
18 static CACHE: RefCell<Cache> = RefCell::new(Cache::default());
20}
21
22struct Cache {
23 limits: [usize; CLASSES],
24 pages: [Vec<StorageVec>; CLASSES],
25 metrics: Arc<PageCacheMetrics>,
26}
27
28impl Default for Cache {
29 fn default() -> Self {
30 Self {
31 limits: DEFAULT_PAGES_CACHE,
32 pages: Default::default(),
33 metrics: PageCacheMetrics::new(),
34 }
35 }
36}
37
38impl Drop for Cache {
39 fn drop(&mut self) {
40 for (class, pages) in self.pages.iter().enumerate() {
41 self.metrics.remove_pages(class, pages.len());
42 }
43 }
44}
45
46impl Cache {
47 fn set_limit(&mut self, class: usize, limit: usize) {
49 self.limits[class] = limit;
50 let pages = &mut self.pages[class];
51 if pages.len() > limit {
52 self.metrics.remove_pages(class, pages.len() - limit);
53 pages.truncate(limit);
56 }
57 }
58}
59
60#[inline]
65fn with_cache<R>(f: impl FnOnce(&mut Cache) -> R) -> Option<R> {
66 CACHE
67 .try_with(|c| c.try_borrow_mut().ok().map(|mut cst| f(&mut cst)))
68 .ok()
69 .flatten()
70}
71
72#[inline]
74pub(crate) fn take_page(size: BytePageSize) -> Option<StorageVec> {
75 let class = size as usize;
76 with_cache(|cst| {
77 let page = cst.pages[class].pop();
78 if page.is_some() {
79 cst.metrics.hit(class);
80 } else {
81 cst.metrics.miss(class);
82 }
83 page
84 })
85 .flatten()
86}
87
88#[inline]
94pub(crate) fn take_page_or_prev(size: BytePageSize) -> Option<StorageVec> {
95 let class = size as usize;
96 with_cache(|cst| {
97 if let Some(page) = cst.pages[class].pop() {
98 cst.metrics.hit(class);
99 return Some(page);
100 }
101 cst.metrics.miss(class);
102
103 let prev = size.prev() as usize;
104 if prev != class
105 && let Some(page) = cst.pages[prev].pop()
106 {
107 cst.metrics.hit(prev);
108 cst.metrics.fallback(class);
109 return Some(page);
110 }
111 None
112 })
113 .flatten()
114}
115
116#[inline]
125pub(crate) unsafe fn put_page(ptr: *mut Header, size: BytePageSize) -> bool {
126 let class = size as usize;
127 with_cache(|cst| {
128 if cst.pages[class].len() < cst.limits[class] {
129 cst.pages[class].push(unsafe { StorageVec::recycle(ptr) });
132 cst.metrics.cached(class);
133 true
134 } else {
135 cst.metrics.freed(class);
136 false
137 }
138 })
139 .unwrap_or(false)
140}
141
142pub(crate) fn set_pages_cache(count: usize) {
144 with_cache(|cst| {
145 for class in 0..CLASSES {
146 cst.set_limit(class, count);
147 }
148 });
149}
150
151pub(crate) fn set_page_cache_size(size: BytePageSize, count: usize) {
153 if size != BytePageSize::Unset {
154 with_cache(|cst| cst.set_limit(size as usize, count));
155 }
156}
157
158#[cfg(test)]
159pub(crate) fn reset() {
164 CACHE.with(|c| {
165 let mut cst = c.borrow_mut();
166 drop(std::mem::take(&mut *cst));
167 });
168}
169
170#[cfg(test)]
171pub(crate) fn cached_pages(size: BytePageSize) -> usize {
173 with_cache(|cst| cst.pages[size as usize].len()).unwrap()
174}
175
176#[derive(Debug, Default)]
193pub struct PageCacheMetrics {
194 classes: [Class; CLASSES],
195}
196
197#[derive(Debug, Default)]
198struct Class {
199 hits: AtomicU64,
200 misses: AtomicU64,
201 fallbacks: AtomicU64,
202 cached: AtomicU64,
203 freed: AtomicU64,
204 added: AtomicU64,
207 removed: AtomicU64,
208}
209
210#[derive(Copy, Clone, Debug, PartialEq, Eq)]
212pub struct PageCacheSnapshot {
213 sizes: [PageCacheStats; CLASSES],
214}
215
216#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
218#[non_exhaustive]
219pub struct PageCacheStats {
220 pub size: BytePageSize,
222 pub hits: u64,
225 pub misses: u64,
230 pub fallbacks: u64,
233 pub cached: u64,
235 pub freed: u64,
237 pub pages: usize,
239 pub bytes: usize,
241}
242
243impl PageCacheMetrics {
244 pub fn new() -> Arc<Self> {
246 Arc::default()
247 }
248
249 pub fn current() -> Arc<Self> {
254 with_cache(|cst| cst.metrics.clone()).unwrap_or_default()
255 }
256
257 pub fn set_current(metrics: Arc<Self>) {
263 with_cache(|cst| {
264 for (class, pages) in cst.pages.iter().enumerate() {
265 cst.metrics.remove_pages(class, pages.len());
266 metrics.add_pages(class, pages.len());
267 }
268 cst.metrics = metrics;
269 });
270 }
271
272 pub fn snapshot(&self) -> PageCacheSnapshot {
274 let mut snapshot = PageCacheSnapshot::default();
275 for (stats, class) in snapshot.sizes.iter_mut().zip(&self.classes) {
276 stats.hits = class.hits.load(Relaxed);
277 stats.misses = class.misses.load(Relaxed);
278 stats.fallbacks = class.fallbacks.load(Relaxed);
279 stats.cached = class.cached.load(Relaxed);
280 stats.freed = class.freed.load(Relaxed);
281 let removed = class.removed.load(Relaxed);
282 let hits = stats.hits;
283 stats.pages =
285 (stats.cached + class.added.load(Relaxed)).saturating_sub(hits + removed) as usize;
286 stats.bytes = stats.pages * stats.size.alloc_size();
287 }
288 snapshot
289 }
290
291 pub(crate) fn hit(&self, class: usize) {
292 inc(&self.classes[class].hits, 1);
293 }
294
295 pub(crate) fn miss(&self, class: usize) {
296 inc(&self.classes[class].misses, 1);
297 }
298
299 pub(crate) fn fallback(&self, class: usize) {
300 inc(&self.classes[class].fallbacks, 1);
301 }
302
303 pub(crate) fn cached(&self, class: usize) {
304 inc(&self.classes[class].cached, 1);
305 }
306
307 pub(crate) fn freed(&self, class: usize) {
308 inc(&self.classes[class].freed, 1);
309 }
310
311 pub(crate) fn remove_pages(&self, class: usize, pages: usize) {
313 if pages != 0 {
314 inc(&self.classes[class].removed, pages as u64);
315 }
316 }
317
318 pub(crate) fn add_pages(&self, class: usize, pages: usize) {
319 if pages != 0 {
320 inc(&self.classes[class].added, pages as u64);
321 }
322 }
323}
324
325#[inline]
327fn inc(counter: &AtomicU64, n: u64) {
328 counter.store(counter.load(Relaxed).wrapping_add(n), Relaxed);
329}
330
331impl PageCacheSnapshot {
332 pub fn get(&self, size: BytePageSize) -> Option<&PageCacheStats> {
335 self.sizes.get(size as usize)
336 }
337
338 pub fn iter(&self) -> impl Iterator<Item = &PageCacheStats> {
340 self.sizes.iter()
341 }
342
343 pub fn total(&self) -> PageCacheStats {
345 self.sizes.iter().fold(
346 PageCacheStats {
347 size: BytePageSize::Unset,
348 ..PageCacheStats::default()
349 },
350 |mut total, s| {
351 total.hits += s.hits;
352 total.misses += s.misses;
353 total.fallbacks += s.fallbacks;
354 total.cached += s.cached;
355 total.freed += s.freed;
356 total.pages += s.pages;
357 total.bytes += s.bytes;
358 total
359 },
360 )
361 }
362}
363
364impl Default for PageCacheSnapshot {
365 fn default() -> Self {
366 Self {
367 sizes: PAGE_SIZES.map(|size| PageCacheStats {
368 size,
369 ..PageCacheStats::default()
370 }),
371 }
372 }
373}
374
375#[cfg(test)]
376mod tests {
377 use super::*;
378 use crate::*;
379
380 #[test]
385 fn cached() {
386 reset();
387
388 let mut st = StorageVec::sized(BytePageSize::Size8);
389 assert_eq!(st.page_size(), BytePageSize::Size8);
390
391 st.put_u8(b'h');
392 let addr = st.0;
393 drop(st);
394
395 let st = StorageVec::sized(BytePageSize::Size8);
396 assert_eq!(addr, st.0);
397 }
398
399 #[test]
405 fn cache_metrics() {
406 reset();
407 let metrics = PageCacheMetrics::current();
408 assert_eq!(metrics.snapshot(), PageCacheSnapshot::default());
409
410 let size = BytePageSize::Size8;
411 let st = StorageVec::sized(size);
412 let s = *metrics.snapshot().get(size).unwrap();
413 assert_eq!(
414 (s.size, s.hits, s.misses, s.cached, s.pages),
415 (size, 0, 1, 0, 0)
416 );
417
418 drop(st);
419 let s = *metrics.snapshot().get(size).unwrap();
420 assert_eq!(
421 (s.hits, s.misses, s.cached, s.freed, s.pages),
422 (0, 1, 1, 0, 1)
423 );
424 assert_eq!(s.bytes, 8 * 1024);
425
426 let st = StorageVec::sized(size);
427 let s = *metrics.snapshot().get(size).unwrap();
428 assert_eq!(
429 (s.hits, s.misses, s.cached, s.pages, s.bytes),
430 (1, 1, 1, 0, 0)
431 );
432
433 crate::set_page_cache_size(size, 0);
434 drop(st);
435 let s = *metrics.snapshot().get(size).unwrap();
436 assert_eq!((s.cached, s.freed, s.pages), (1, 1, 0));
437
438 drop(StorageVec::sized(BytePageSize::Size4));
440 drop(StorageVec::with_capacity(100));
441 let snapshot = metrics.snapshot();
442 let s = *snapshot.get(BytePageSize::Size4).unwrap();
443 assert_eq!((s.misses, s.cached, s.pages), (1, 1, 1));
444 assert!(snapshot.get(BytePageSize::Unset).is_none());
445 assert_eq!(snapshot.iter().count(), 9);
446 let total = snapshot.total();
447 assert_eq!(total.size, BytePageSize::Unset);
448 assert_eq!(
449 (total.hits, total.misses, total.cached, total.freed),
450 (1, 2, 2, 1)
451 );
452 assert_eq!((total.pages, total.bytes), (1, 4 * 1024));
453
454 reset();
456 assert_eq!(metrics.snapshot().total().pages, 0);
457 assert!(!Arc::ptr_eq(&metrics, &PageCacheMetrics::current()));
458 }
459
460 #[test]
466 fn cache_metrics_set_current() {
467 reset();
468 let size = BytePageSize::Size16;
469 let old = PageCacheMetrics::current();
470 drop(StorageVec::sized(size));
471 assert_eq!(old.snapshot().get(size).unwrap().pages, 1);
472
473 let metrics = PageCacheMetrics::new();
474 PageCacheMetrics::set_current(metrics.clone());
475 assert!(Arc::ptr_eq(&metrics, &PageCacheMetrics::current()));
476 assert_eq!(old.snapshot().get(size).unwrap().pages, 0);
477 assert_eq!(metrics.snapshot().get(size).unwrap().pages, 1);
478
479 drop(StorageVec::sized(size));
480 let s = *metrics.snapshot().get(size).unwrap();
481 assert_eq!((s.hits, s.misses, s.cached, s.pages), (1, 0, 1, 1));
482 assert_eq!(old.snapshot().get(size).unwrap().hits, 0);
483
484 let other = std::thread::spawn(PageCacheMetrics::current)
486 .join()
487 .unwrap();
488 assert!(!Arc::ptr_eq(&metrics, &other));
489 reset();
490 }
491
492 #[test]
497 fn default_cache_limit_per_page_size() {
498 assert_eq!(DEFAULT_PAGES_CACHE, [128, 64, 64, 32, 16, 8, 16, 2, 1]);
499 for size in crate::size::PAGE_SIZES {
500 reset();
501
502 let limit = DEFAULT_PAGES_CACHE[size as usize];
503 let pages: Vec<_> = (0..limit + 4).map(|_| StorageVec::sized(size)).collect();
504 drop(pages);
505 assert_eq!(cached_pages(size), limit, "{size:?}");
506 }
507 reset();
508 }
509
510 #[test]
516 #[allow(deprecated)]
517 fn pages_cache_size() {
518 reset();
519
520 crate::set_pages_cache(1);
521 drop((
522 StorageVec::sized(BytePageSize::Size4),
523 StorageVec::sized(BytePageSize::Size4),
524 ));
525 assert_eq!(cached_pages(BytePageSize::Size4), 1);
526
527 crate::set_pages_cache(0);
528 drop(StorageVec::sized(BytePageSize::Size8));
529 assert_eq!(cached_pages(BytePageSize::Size8), 0);
530
531 crate::set_pages_cache(16);
533 let st = StorageVec::sized(BytePageSize::Size16);
534 CACHE.with(|c| {
535 let _cst = c.borrow_mut();
536 drop(st);
537 });
538 assert_eq!(cached_pages(BytePageSize::Size16), 0);
539
540 CACHE.with(|c| {
542 let _cst = c.borrow_mut();
543 crate::set_pages_cache(3);
544 });
545 assert_eq!(cache_limits(), [16; CLASSES]);
546 }
547
548 fn cache_limits() -> [usize; CLASSES] {
549 CACHE.with(|c| c.borrow().limits)
550 }
551
552 #[test]
558 #[allow(deprecated)]
559 fn page_cache_size() {
560 reset();
561
562 crate::set_page_cache_size(BytePageSize::Size64, 1);
563 crate::set_page_cache_size(BytePageSize::Unset, 5);
564 let mut expected = DEFAULT_PAGES_CACHE;
565 expected[BytePageSize::Size64 as usize] = 1;
566 assert_eq!(cache_limits(), expected);
567
568 drop((
569 StorageVec::sized(BytePageSize::Size64),
570 StorageVec::sized(BytePageSize::Size64),
571 StorageVec::sized(BytePageSize::Size48),
572 StorageVec::sized(BytePageSize::Size48),
573 ));
574 assert_eq!(cached_pages(BytePageSize::Size64), 1);
575 assert_eq!(cached_pages(BytePageSize::Size48), 2);
576
577 CACHE.with(|c| {
579 let _cst = c.borrow_mut();
580 crate::set_page_cache_size(BytePageSize::Size48, 3);
581 });
582 assert_eq!(cache_limits(), expected);
583
584 crate::set_pages_cache(2);
585 assert_eq!(cache_limits(), [2; CLASSES]);
586 reset();
587 }
588
589 #[test]
595 #[allow(deprecated)]
596 fn lowered_limit_releases_pages() {
597 reset();
598 let metrics = PageCacheMetrics::current();
599 let pages = |size| metrics.snapshot().get(size).unwrap().pages;
600
601 let (s4, s64) = (BytePageSize::Size4, BytePageSize::Size64);
602 drop([0; 4].map(|_| StorageVec::sized(s64)));
603 drop([0; 4].map(|_| StorageVec::sized(s4)));
604 assert_eq!((cached_pages(s64), pages(s64)), (4, 4));
605
606 crate::set_page_cache_size(s64, 1);
607 assert_eq!((cached_pages(s64), pages(s64)), (1, 1));
608 assert_eq!((cached_pages(s4), pages(s4)), (4, 4));
609 crate::set_page_cache_size(s64, 0);
610 assert_eq!((cached_pages(s64), pages(s64)), (0, 0));
611
612 crate::set_pages_cache(1);
614 assert_eq!((cached_pages(s4), pages(s4)), (1, 1));
615 let total = metrics.snapshot().total();
616 assert_eq!((total.cached, total.freed, total.pages), (8, 0, 1));
617
618 crate::set_page_cache_size(s4, 8);
620 assert_eq!((cached_pages(s4), pages(s4)), (1, 1));
621 reset();
622 }
623
624 #[test]
630 fn bytes_mut_with_page_size() {
631 reset();
632
633 let mut buf = BytesMut::with_page_size(BytePageSize::Size8);
634 assert_eq!(buf.page_size(), BytePageSize::Size8);
635 assert_eq!(buf.capacity(), BytePageSize::Size8.capacity());
636 let ptr = buf.as_ptr();
637
638 buf.extend_from_slice(&[1; 1000]);
640 let head = buf.split_to(500);
641 drop(buf);
642 assert_eq!(cached_pages(BytePageSize::Size8), 0);
643 drop(head);
644 assert_eq!(cached_pages(BytePageSize::Size8), 1);
645
646 let buf = BytesMut::with_page_size(BytePageSize::Size8);
647 assert_eq!(buf.as_ptr(), ptr);
648 assert_eq!(buf.capacity(), BytePageSize::Size8.capacity());
649 assert_eq!(cached_pages(BytePageSize::Size8), 0);
650
651 let buf = BytesMut::with_page_size(BytePageSize::Unset);
653 assert_eq!(buf.page_size(), BytePageSize::Size64);
654 assert_eq!(buf.capacity(), BytePageSize::Size64.capacity());
655 drop(buf);
656 assert_eq!(cached_pages(BytePageSize::Size64), 1);
657
658 let buf = BytesMut::with_capacity(BytePageSize::Size64.capacity() + 1);
660 assert_eq!(buf.page_size(), BytePageSize::Unset);
661 assert_eq!(BytesMut::new().page_size(), BytePageSize::Unset);
662 assert_eq!(BytesMut::with_capacity(64).page_size(), BytePageSize::Unset);
663 assert_eq!(
664 BytesMut::copy_from_slice(b"hello").page_size(),
665 BytePageSize::Unset
666 );
667 drop(buf);
668 for size in crate::size::PAGE_SIZES {
669 let expected = usize::from(size == BytePageSize::Size64);
670 assert_eq!(cached_pages(size), expected);
671 }
672 }
673
674 #[test]
681 fn new_page_falls_back_to_smaller_size() {
682 reset();
683 let metrics = PageCacheMetrics::current();
684 let (s8, s16) = (BytePageSize::Size8, BytePageSize::Size16);
685
686 let st = StorageVec::sized(s8);
687 let addr = st.0;
688 drop(st);
689
690 let st = StorageVec::sized(s16);
691 assert_eq!(st.0, addr);
692 assert_eq!(st.page_size(), s8);
693 assert_eq!(cached_pages(s8), 0);
694 let snapshot = metrics.snapshot();
695 let s = *snapshot.get(s16).unwrap();
696 assert_eq!((s.hits, s.misses, s.fallbacks, s.pages), (0, 1, 1, 0));
697 let s = *snapshot.get(s8).unwrap();
698 assert_eq!(
699 (s.hits, s.misses, s.fallbacks, s.cached, s.pages),
700 (1, 1, 0, 1, 0)
701 );
702
703 let st2 = StorageVec::sized(s16);
704 assert_eq!(st2.page_size(), s16);
705 let snapshot = metrics.snapshot();
706 let s = *snapshot.get(s16).unwrap();
707 assert_eq!((s.hits, s.misses, s.fallbacks), (0, 2, 1));
708 let total = snapshot.total();
709 assert_eq!((total.hits, total.misses, total.fallbacks), (1, 3, 1));
710 drop((st, st2));
711 reset();
712 }
713
714 #[test]
720 fn new_page_falls_back_one_size_only() {
721 reset();
722 let (s4, s8, s16) = (
723 BytePageSize::Size4,
724 BytePageSize::Size8,
725 BytePageSize::Size16,
726 );
727
728 drop(StorageVec::sized_exact(s4));
729 let st = StorageVec::sized(s16);
730 assert_eq!(st.page_size(), s16);
731 assert_eq!(cached_pages(s4), 1);
732
733 let p8 = StorageVec::sized_exact(s8);
734 drop(st);
735 drop(p8);
736 let st = StorageVec::sized(s16);
737 assert_eq!(st.page_size(), s16);
738 assert_eq!((cached_pages(s8), cached_pages(s16)), (1, 0));
739 drop(st);
740 reset();
741 }
742
743 #[test]
751 fn new_buffers_fall_back_to_smaller_page() {
752 reset();
753 let metrics = PageCacheMetrics::current();
754 let (s8, s16) = (BytePageSize::Size8, BytePageSize::Size16);
755
756 drop(StorageVec::sized_exact(s8));
757 let buf = BytesMut::with_page_size(s16);
758 assert_eq!(buf.page_size(), s8);
759 assert_eq!(buf.capacity(), s8.capacity());
760 drop(buf);
761
762 drop(StorageVec::sized_exact(BytePageSize::Size48));
763 let buf = BytesMut::with_page_size(BytePageSize::Unset);
764 assert_eq!(buf.page_size(), BytePageSize::Size48);
765 drop(buf);
766
767 #[cfg(feature = "tracker")]
768 {
769 let tracker = BytesTracker::new();
770 let buf = BytesMut::with_page_size_in(s16, &tracker);
771 assert_eq!(buf.page_size(), s8);
772 assert_eq!(tracker.snapshot().live_bytes, s8.alloc_size());
773 }
774
775 assert_eq!(cached_pages(s8), 1);
776 let mut pages = BytePages::new(s16);
777 pages.put_slice(b"hello");
778 assert_eq!(cached_pages(s8), 0);
779 let expected = if cfg!(feature = "tracker") { 3 } else { 2 };
780 assert_eq!(metrics.snapshot().get(s16).unwrap().fallbacks, expected);
781 drop(pages);
782 reset();
783 }
784}