1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

mod console_markdown_report;
mod metrics;

pub use self::{
    console_markdown_report::ConsoleMarkdownReport, metrics::Metrics,
};

use super::{Allocation, DeviceAllocator};

use anyhow::Result;
use ash::vk;
use std::collections::HashMap;

/// Types which implement this trait can be used by the Metrics Allocator to
/// render a report on memory allocations.
pub trait MetricsReport {
    /// Render the metrics report.
    ///
    /// The output is implementation-defined (console, file, format, etc..).
    fn render(
        &self,
        name: &str,
        total: &Metrics,
        metrics_by_type: &HashMap<u32, Metrics>,
    );
}

/// A device allocator decorator which records the number of allocations and
/// other metrics. A summary of results is printed when the allocator is
/// destroyed.
pub struct MetricsAllocator<Alloc: DeviceAllocator> {
    allocator: Alloc,
    name: String,
    by_type: HashMap<u32, Metrics>,
    total: Metrics,
    report: Box<dyn MetricsReport>,
}

impl<Alloc: DeviceAllocator> MetricsAllocator<Alloc> {
    /// Decorate an existing allocator with support for metrics.
    pub fn new<Report: 'static + MetricsReport>(
        name: impl Into<String>,
        report: Report,
        allocator: Alloc,
    ) -> Self {
        Self {
            name: name.into(),
            allocator,
            by_type: HashMap::new(),
            total: Metrics::default(),
            report: Box::new(report),
        }
    }

    fn record_allocation(&mut self, allocation: &Allocation) {
        self.total.measure_alloctaion(&allocation);
        self.by_type
            .entry(allocation.memory_type_index)
            .or_default()
            .measure_alloctaion(allocation);
    }

    fn record_free(&mut self, allocation: &Allocation) {
        if allocation.is_null() {
            return;
        }
        self.total.measure_free();
        self.by_type
            .entry(allocation.memory_type_index)
            .and_modify(|metrics| metrics.measure_free());
    }
}

impl<Alloc: DeviceAllocator> DeviceAllocator for MetricsAllocator<Alloc> {
    unsafe fn allocate(
        &mut self,
        allocate_info: vk::MemoryAllocateInfo,
    ) -> Result<Allocation> {
        let allocation = self.allocator.allocate(allocate_info)?;
        self.record_allocation(&allocation);
        Ok(allocation)
    }

    unsafe fn free(&mut self, allocation: &Allocation) -> Result<()> {
        self.record_free(allocation);
        self.allocator.free(allocation)
    }
}

impl<T: DeviceAllocator> Drop for MetricsAllocator<T> {
    fn drop(&mut self) {
        self.report.render(&self.name, &self.total, &self.by_type);
    }
}