Add a public API to expose allocation reports (#226)

nical · web-flow · commit 4b5c0d4f17dd · 2024-06-17T16:29:38.000+02:00
diff --git a/src/allocator/dedicated_block_allocator/mod.rs b/src/allocator/dedicated_block_allocator/mod.rs
@@ -116,6 +116,7 @@ impl SubAllocator for DedicatedBlockAllocator {
                 .name
                 .clone()
                 .unwrap_or_else(|| "<Unnamed Dedicated allocation>".to_owned()),
+            offset: 0,
             size: self.size,
             #[cfg(feature = "visualizer")]
             backtrace: self.backtrace.clone(),
diff --git a/src/allocator/free_list_allocator/mod.rs b/src/allocator/free_list_allocator/mod.rs
@@ -398,6 +398,7 @@ impl SubAllocator for FreeListAllocator {
                     .name
                     .clone()
                     .unwrap_or_else(|| "<Unnamed FreeList allocation>".to_owned()),
+                offset: chunk.offset,
                 size: chunk.size,
                 #[cfg(feature = "visualizer")]
                 backtrace: chunk.backtrace.clone(),
diff --git a/src/allocator/mod.rs b/src/allocator/mod.rs
@@ -1,4 +1,4 @@
-use std::{backtrace::Backtrace, sync::Arc};
+use std::{backtrace::Backtrace, fmt, ops::Range, sync::Arc};
 
 use log::*;
 
@@ -29,20 +29,83 @@ impl AllocationType {
     }
 }
 
+/// Describes an allocation in the [`AllocatorReport`].
 #[derive(Clone)]
-pub(crate) struct AllocationReport {
-    pub(crate) name: String,
-    pub(crate) size: u64,
+pub struct AllocationReport {
+    /// The name provided to the `allocate()` function.
+    pub name: String,
+    /// The offset in bytes of the allocation in its memory block.
+    pub offset: u64,
+    /// The size in bytes of the allocation.
+    pub size: u64,
     #[cfg(feature = "visualizer")]
     pub(crate) backtrace: Arc<Backtrace>,
 }
 
+/// Describes a memory block in the [`AllocatorReport`].
+#[derive(Clone)]
+pub struct MemoryBlockReport {
+    /// The size in bytes of this memory block.
+    pub size: u64,
+    /// The range of allocations in [`AllocatorReport::allocations`] that are associated
+    /// to this memory block.
+    pub allocations: Range<usize>,
+}
+
+/// A report that can be generated for informational purposes using `Allocator::generate_report()`.
+#[derive(Clone)]
+pub struct AllocatorReport {
+    /// All live allocations, sub-allocated from memory blocks.
+    pub allocations: Vec<AllocationReport>,
+    /// All memory blocks.
+    pub blocks: Vec<MemoryBlockReport>,
+    /// Sum of the memory used by all allocations, in bytes.
+    pub total_allocated_bytes: u64,
+    /// Sum of the memory reserved by all memory blocks including unallocated regions, in bytes.
+    pub total_reserved_bytes: u64,
+}
+
+impl fmt::Debug for AllocationReport {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let name = if !self.name.is_empty() {
+            self.name.as_str()
+        } else {
+            "--"
+        };
+        write!(f, "{name:?}: {}", fmt_bytes(self.size))
+    }
+}
+
+impl fmt::Debug for AllocatorReport {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut allocations = self.allocations.clone();
+        allocations.sort_by_key(|alloc| std::cmp::Reverse(alloc.size));
+
+        let max_num_allocations_to_print = f.precision().unwrap_or(usize::MAX);
+        allocations.truncate(max_num_allocations_to_print);
+
+        f.debug_struct("AllocatorReport")
+            .field(
+                "summary",
+                &std::format_args!(
+                    "{} / {}",
+                    fmt_bytes(self.total_allocated_bytes),
+                    fmt_bytes(self.total_reserved_bytes)
+                ),
+            )
+            .field("blocks", &self.blocks.len())
+            .field("allocations", &self.allocations.len())
+            .field("largest", &allocations.as_slice())
+            .finish()
+    }
+}
+
 #[cfg(feature = "visualizer")]
 pub(crate) trait SubAllocatorBase: crate::visualizer::SubAllocatorVisualizer {}
 #[cfg(not(feature = "visualizer"))]
 pub(crate) trait SubAllocatorBase {}
 
-pub(crate) trait SubAllocator: SubAllocatorBase + std::fmt::Debug + Sync + Send {
+pub(crate) trait SubAllocator: SubAllocatorBase + fmt::Debug + Sync + Send {
     fn allocate(
         &mut self,
         size: u64,
@@ -82,8 +145,6 @@ pub(crate) trait SubAllocator: SubAllocatorBase + std::fmt::Debug + Sync + Send
     }
 }
 
-pub(crate) const VISUALIZER_TABLE_MAX_ENTRY_NAME_LEN: usize = 40;
-
 pub(crate) fn fmt_bytes(mut amount: u64) -> String {
     const SUFFIX: [&str; 5] = ["B", "KB", "MB", "GB", "TB"];
 
diff --git a/src/d3d12/mod.rs b/src/d3d12/mod.rs
@@ -88,8 +88,8 @@ use super::allocator;
 use super::allocator::AllocationType;
 
 use crate::{
-    allocator::fmt_bytes, AllocationError, AllocationSizes, AllocatorDebugSettings, MemoryLocation,
-    Result,
+    allocator::{AllocatorReport, MemoryBlockReport},
+    AllocationError, AllocationSizes, AllocatorDebugSettings, MemoryLocation, Result,
 };
 
 /// [`ResourceCategory`] is used for supporting [`D3D12_RESOURCE_HEAP_TIER_1`].
@@ -1102,50 +1102,38 @@ impl Allocator {
             Ok(())
         }
     }
-}
 
-impl fmt::Debug for Allocator {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let mut allocation_report = vec![];
-        let mut total_reserved_size_in_bytes = 0;
+    pub fn generate_report(&self) -> AllocatorReport {
+        let mut allocations = vec![];
+        let mut blocks = vec![];
+        let mut total_reserved_bytes = 0;
 
         for memory_type in &self.memory_types {
             for block in memory_type.memory_blocks.iter().flatten() {
-                total_reserved_size_in_bytes += block.size;
-                allocation_report.extend(block.sub_allocator.report_allocations())
+                total_reserved_bytes += block.size;
+                let first_allocation = allocations.len();
+                allocations.extend(block.sub_allocator.report_allocations());
+                blocks.push(MemoryBlockReport {
+                    size: block.size,
+                    allocations: first_allocation..allocations.len(),
+                });
             }
         }
 
-        let total_used_size_in_bytes = allocation_report.iter().map(|report| report.size).sum();
-
-        allocation_report.sort_by_key(|alloc| std::cmp::Reverse(alloc.size));
-
-        writeln!(
-            f,
-            "================================================================",
-        )?;
-        writeln!(
-            f,
-            "ALLOCATION BREAKDOWN ({} / {})",
-            fmt_bytes(total_used_size_in_bytes),
-            fmt_bytes(total_reserved_size_in_bytes),
-        )?;
+        let total_allocated_bytes = allocations.iter().map(|report| report.size).sum();
 
-        let max_num_allocations_to_print = f.precision().map_or(usize::MAX, |n| n);
-        for (idx, alloc) in allocation_report.iter().enumerate() {
-            if idx >= max_num_allocations_to_print {
-                break;
-            }
-            writeln!(
-                f,
-                "{:max_len$.max_len$}\t- {}",
-                alloc.name,
-                fmt_bytes(alloc.size),
-                max_len = allocator::VISUALIZER_TABLE_MAX_ENTRY_NAME_LEN,
-            )?;
+        AllocatorReport {
+            allocations,
+            blocks,
+            total_allocated_bytes,
+            total_reserved_bytes,
         }
+    }
+}
 
-        Ok(())
+impl fmt::Debug for Allocator {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.generate_report().fmt(f)
     }
 }
 
diff --git a/src/lib.rs b/src/lib.rs
@@ -212,6 +212,8 @@ pub use result::*;
 
 pub(crate) mod allocator;
 
+pub use allocator::{AllocationReport, AllocatorReport, MemoryBlockReport};
+
 #[cfg(feature = "visualizer")]
 pub mod visualizer;
 
diff --git a/src/metal/mod.rs b/src/metal/mod.rs
@@ -2,7 +2,8 @@
 use std::{backtrace::Backtrace, sync::Arc};
 
 use crate::{
-    allocator, AllocationError, AllocationSizes, AllocatorDebugSettings, MemoryLocation, Result,
+    allocator::{self, AllocatorReport, MemoryBlockReport},
+    AllocationError, AllocationSizes, AllocatorDebugSettings, MemoryLocation, Result,
 };
 use log::{debug, Level};
 
@@ -146,7 +147,7 @@ pub struct CommittedAllocationStatistics {
 }
 struct MemoryBlock {
     heap: Arc<metal::Heap>,
-    _size: u64,
+    size: u64,
     sub_allocator: Box<dyn allocator::SubAllocator>,
 }
 
@@ -169,7 +170,7 @@ impl MemoryBlock {
 
         Ok(Self {
             heap,
-            _size: size,
+            size,
             sub_allocator,
         })
     }
@@ -485,4 +486,31 @@ impl Allocator {
     pub fn report_memory_leaks(&self, _log_level: Level) {
         todo!()
     }
+
+    pub fn generate_report(&self) -> AllocatorReport {
+        let mut allocations = vec![];
+        let mut blocks = vec![];
+        let mut total_reserved_bytes = 0;
+
+        for memory_type in &self.memory_types {
+            for block in memory_type.memory_blocks.iter().flatten() {
+                total_reserved_bytes += block.size;
+                let first_allocation = allocations.len();
+                allocations.extend(block.sub_allocator.report_allocations());
+                blocks.push(MemoryBlockReport {
+                    size: block.size,
+                    allocations: first_allocation..allocations.len(),
+                });
+            }
+        }
+
+        let total_allocated_bytes = allocations.iter().map(|report| report.size).sum();
+
+        AllocatorReport {
+            allocations,
+            blocks,
+            total_allocated_bytes,
+            total_reserved_bytes,
+        }
+    }
 }
diff --git a/src/visualizer/allocation_reports.rs b/src/visualizer/allocation_reports.rs
@@ -115,6 +115,7 @@ pub(crate) fn render_allocation_reports_ui(
                     name,
                     size,
                     backtrace,
+                    ..
                 } = alloc;
 
                 row.col(|ui| {
diff --git a/src/vulkan/mod.rs b/src/vulkan/mod.rs
@@ -12,8 +12,8 @@ use log::{debug, Level};
 
 use super::allocator;
 use crate::{
-    allocator::fmt_bytes, AllocationError, AllocationSizes, AllocatorDebugSettings, MemoryLocation,
-    Result,
+    allocator::{AllocatorReport, MemoryBlockReport},
+    AllocationError, AllocationSizes, AllocatorDebugSettings, MemoryLocation, Result,
 };
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
@@ -691,47 +691,7 @@ pub struct Allocator {
 
 impl fmt::Debug for Allocator {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let mut allocation_report = vec![];
-        let mut total_reserved_size_in_bytes = 0;
-
-        for memory_type in &self.memory_types {
-            for block in memory_type.memory_blocks.iter().flatten() {
-                total_reserved_size_in_bytes += block.size;
-                allocation_report.extend(block.sub_allocator.report_allocations())
-            }
-        }
-
-        let total_used_size_in_bytes = allocation_report.iter().map(|report| report.size).sum();
-
-        allocation_report.sort_by_key(|alloc| std::cmp::Reverse(alloc.size));
-
-        writeln!(
-            f,
-            "================================================================",
-        )?;
-        writeln!(
-            f,
-            "ALLOCATION BREAKDOWN ({} / {})",
-            fmt_bytes(total_used_size_in_bytes),
-            fmt_bytes(total_reserved_size_in_bytes),
-        )?;
-
-        let max_num_allocations_to_print = f.precision().map_or(usize::MAX, |n| n);
-        for (idx, alloc) in allocation_report.iter().enumerate() {
-            if idx >= max_num_allocations_to_print {
-                break;
-            }
-
-            writeln!(
-                f,
-                "{:max_len$.max_len$}\t- {}",
-                alloc.name,
-                fmt_bytes(alloc.size),
-                max_len = allocator::VISUALIZER_TABLE_MAX_ENTRY_NAME_LEN,
-            )?;
-        }
-
-        Ok(())
+        self.generate_report().fmt(f)
     }
 }
 
@@ -972,6 +932,33 @@ impl Allocator {
             })
             .map(|memory_type| memory_type.memory_type_index as _)
     }
+
+    pub fn generate_report(&self) -> AllocatorReport {
+        let mut allocations = vec![];
+        let mut blocks = vec![];
+        let mut total_reserved_bytes = 0;
+
+        for memory_type in &self.memory_types {
+            for block in memory_type.memory_blocks.iter().flatten() {
+                total_reserved_bytes += block.size;
+                let first_allocation = allocations.len();
+                allocations.extend(block.sub_allocator.report_allocations());
+                blocks.push(MemoryBlockReport {
+                    size: block.size,
+                    allocations: first_allocation..allocations.len(),
+                });
+            }
+        }
+
+        let total_allocated_bytes = allocations.iter().map(|report| report.size).sum();
+
+        AllocatorReport {
+            allocations,
+            blocks,
+            total_allocated_bytes,
+            total_reserved_bytes,
+        }
+    }
 }
 
 impl Drop for Allocator {
