From d10c207371f13a041390808c1ff96f72e1e462da Mon Sep 17 00:00:00 2001
From: "Lei, HUANG" <mrsatangel@gmail.com>
Date: Thu, 6 Feb 2025 09:17:37 +0000
Subject: [PATCH] poc-write-path:  Add allocation tracking to `BulkMemtable`
 methods

 - Updated `write_bulk` in `bulk.rs` to track memory allocation using `alloc_tracker.on_allocation`.
 - Modified `freeze` in `bulk.rs` to signal completion of allocation with `alloc_tracker.done_allocating`.
---
 src/mito2/src/memtable/bulk.rs | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs
index 482bc9e503..971bf44511 100644
--- a/src/mito2/src/memtable/bulk.rs
+++ b/src/mito2/src/memtable/bulk.rs
@@ -133,6 +133,7 @@ impl Memtable for BulkMemtable {
     }
 
     fn write_bulk(&self, fragment: BulkPart) -> Result<()> {
+        self.alloc_tracker.on_allocation(fragment.data.len());
         let mut parts = self.parts.write().unwrap();
         parts.push(fragment);
         Ok(())
@@ -163,6 +164,7 @@ impl Memtable for BulkMemtable {
     }
 
     fn freeze(&self) -> Result<()> {
+        self.alloc_tracker.done_allocating();
         Ok(())
     }