Skip to content

Commit 1ad2acb

Browse files
zhang2014zhang2014
authored
feat(base): enhanced query-level memory management (#17358)
* chore(base): test memory alloc bucket * chore(base): embed tracker into large memory * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory. * chore(base): embed tracker into large memory.
1 parent 1dd6ea9 commit 1ad2acb

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

51 files changed

+2284
-962
lines changed

‎src/binaries/query/ee_main.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
mod entry;
1919

20-
use databend_common_base::mem_allocator::GlobalAllocator;
20+
use databend_common_base::mem_allocator::TrackingGlobalAllocator;
2121
use databend_common_base::runtime::Runtime;
2222
use databend_common_base::runtime::ThreadTracker;
2323
use databend_common_config::InnerConfig;
@@ -34,7 +34,7 @@ use crate::entry::run_cmd;
3434
use crate::entry::start_services;
3535

3636
#[global_allocator]
37-
pub static GLOBAL_ALLOCATOR: GlobalAllocator = GlobalAllocator;
37+
pub static GLOBAL_ALLOCATOR: TrackingGlobalAllocator = TrackingGlobalAllocator::create();
3838

3939
fn main() {
4040
let binary_version = (*databend_common_config::DATABEND_COMMIT_VERSION).clone();

‎src/binaries/query/entry.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
use std::env;
1616
use std::time::Duration;
1717

18-
use databend_common_base::mem_allocator::GlobalAllocator;
18+
use databend_common_base::mem_allocator::TrackingGlobalAllocator;
1919
use databend_common_base::runtime::set_alloc_error_hook;
2020
use databend_common_base::runtime::GLOBAL_MEM_STAT;
2121
use databend_common_config::Commands;
@@ -305,8 +305,8 @@ pub async fn start_services(conf: &InnerConfig) -> Result<(), MainError> {
305305
"unlimited".to_string()
306306
}
307307
});
308-
println!(" allocator: {}", GlobalAllocator::name());
309-
println!(" config: {}", GlobalAllocator::conf());
308+
println!(" allocator: {}", TrackingGlobalAllocator::name());
309+
println!(" config: {}", TrackingGlobalAllocator::conf());
310310

311311
println!();
312312
println!("Cluster: {}", {

‎src/binaries/query/oss_main.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
mod entry;
1919

20-
use databend_common_base::mem_allocator::GlobalAllocator;
20+
use databend_common_base::mem_allocator::TrackingGlobalAllocator;
2121
use databend_common_base::runtime::Runtime;
2222
use databend_common_base::runtime::ThreadTracker;
2323
use databend_common_config::InnerConfig;
@@ -35,7 +35,7 @@ use crate::entry::run_cmd;
3535
use crate::entry::start_services;
3636

3737
#[global_allocator]
38-
pub static GLOBAL_ALLOCATOR: GlobalAllocator = GlobalAllocator;
38+
pub static GLOBAL_ALLOCATOR: TrackingGlobalAllocator = TrackingGlobalAllocator::create();
3939

4040
fn main() {
4141
let binary_version = (*databend_common_config::DATABEND_COMMIT_VERSION).clone();

‎src/common/base/src/lib.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@
2626
#![feature(variant_count)]
2727
#![feature(ptr_alignment_type)]
2828
#![feature(vec_into_raw_parts)]
29+
#![feature(slice_ptr_get)]
30+
#![feature(alloc_layout_extra)]
31+
#![feature(let_chains)]
2932

3033
pub mod base;
3134
pub mod containers;

‎src/common/base/src/mem_allocator/global.rs‎

Lines changed: 53 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,42 @@ use std::alloc::Layout;
1919
use std::ptr::null_mut;
2020
use std::ptr::NonNull;
2121

22+
use crate::mem_allocator::tracker::MetaTrackerAllocator;
2223
use crate::mem_allocator::DefaultAllocator;
2324

25+
pub type DefaultGlobalAllocator = GlobalAllocator<DefaultAllocator>;
26+
pub type TrackingGlobalAllocator = GlobalAllocator<MetaTrackerAllocator<DefaultAllocator>>;
27+
2428
/// Global allocator, default is JeAllocator.
2529
2630
#[derive(Debug, Clone, Copy, Default)]
27-
pub struct GlobalAllocator;
31+
pub struct GlobalAllocator<T> {
32+
inner: T,
33+
}
34+
35+
impl GlobalAllocator<MetaTrackerAllocator<DefaultAllocator>> {
36+
pub const fn create() -> GlobalAllocator<MetaTrackerAllocator<DefaultAllocator>> {
37+
GlobalAllocator {
38+
inner: MetaTrackerAllocator::create(DefaultAllocator::create()),
39+
}
40+
}
41+
42+
pub fn name() -> String {
43+
DefaultAllocator::name()
44+
}
45+
46+
pub fn conf() -> String {
47+
DefaultAllocator::conf()
48+
}
49+
}
50+
51+
impl GlobalAllocator<DefaultAllocator> {
52+
pub const fn create() -> GlobalAllocator<DefaultAllocator> {
53+
GlobalAllocator {
54+
inner: DefaultAllocator::create(),
55+
}
56+
}
2857

29-
impl GlobalAllocator {
3058
pub fn name() -> String {
3159
DefaultAllocator::name()
3260
}
@@ -36,20 +64,20 @@ impl GlobalAllocator {
3664
}
3765
}
3866

39-
unsafe impl Allocator for GlobalAllocator {
67+
unsafe impl<T: Allocator> Allocator for GlobalAllocator<T> {
4068
#[inline(always)]
4169
fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
42-
DefaultAllocator::default().allocate(layout)
70+
self.inner.allocate(layout)
4371
}
4472

4573
#[inline(always)]
4674
fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
47-
DefaultAllocator::default().allocate_zeroed(layout)
75+
self.inner.allocate_zeroed(layout)
4876
}
4977

5078
#[inline(always)]
5179
unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
52-
DefaultAllocator::default().deallocate(ptr, layout)
80+
self.inner.deallocate(ptr, layout)
5381
}
5482

5583
#[inline(always)]
@@ -59,7 +87,7 @@ unsafe impl Allocator for GlobalAllocator {
5987
old_layout: Layout,
6088
new_layout: Layout,
6189
) -> Result<NonNull<[u8]>, AllocError> {
62-
DefaultAllocator::default().grow(ptr, old_layout, new_layout)
90+
self.inner.grow(ptr, old_layout, new_layout)
6391
}
6492

6593
#[inline(always)]
@@ -69,7 +97,7 @@ unsafe impl Allocator for GlobalAllocator {
6997
old_layout: Layout,
7098
new_layout: Layout,
7199
) -> Result<NonNull<[u8]>, AllocError> {
72-
DefaultAllocator::default().grow_zeroed(ptr, old_layout, new_layout)
100+
self.inner.grow_zeroed(ptr, old_layout, new_layout)
73101
}
74102

75103
#[inline(always)]
@@ -79,32 +107,30 @@ unsafe impl Allocator for GlobalAllocator {
79107
old_layout: Layout,
80108
new_layout: Layout,
81109
) -> Result<NonNull<[u8]>, AllocError> {
82-
DefaultAllocator::default().shrink(ptr, old_layout, new_layout)
110+
self.inner.shrink(ptr, old_layout, new_layout)
83111
}
84112
}
85113

86-
unsafe impl GlobalAlloc for GlobalAllocator {
114+
unsafe impl<T: Allocator> GlobalAlloc for GlobalAllocator<T> {
87115
#[inline]
88116
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
89-
if let Ok(ptr) = GlobalAllocator.allocate(layout) {
90-
ptr.as_ptr() as *mut u8
91-
} else {
92-
null_mut()
117+
match self.allocate(layout) {
118+
Ok(ptr) => ptr.as_ptr() as *mut u8,
119+
Err(_) => null_mut(),
93120
}
94121
}
95122

96123
#[inline]
97124
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
98125
let ptr = NonNull::new(ptr).unwrap_unchecked();
99-
GlobalAllocator.deallocate(ptr, layout);
126+
self.deallocate(ptr, layout);
100127
}
101128

102129
#[inline]
103130
unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
104-
if let Ok(ptr) = GlobalAllocator.allocate_zeroed(layout) {
105-
ptr.as_ptr() as *mut u8
106-
} else {
107-
null_mut()
131+
match self.allocate_zeroed(layout) {
132+
Ok(ptr) => ptr.as_ptr() as *mut u8,
133+
Err(_) => null_mut(),
108134
}
109135
}
110136

@@ -115,21 +141,15 @@ unsafe impl GlobalAlloc for GlobalAllocator {
115141
let ptr = NonNull::new(ptr).unwrap_unchecked();
116142
let new_layout = Layout::from_size_align(new_size, layout.align()).unwrap();
117143
match layout.size().cmp(&new_size) {
118-
Less => {
119-
if let Ok(ptr) = GlobalAllocator.grow(ptr, layout, new_layout) {
120-
ptr.as_ptr() as *mut u8
121-
} else {
122-
null_mut()
123-
}
124-
}
125-
Greater => {
126-
if let Ok(ptr) = GlobalAllocator.shrink(ptr, layout, new_layout) {
127-
ptr.as_ptr() as *mut u8
128-
} else {
129-
null_mut()
130-
}
131-
}
132144
Equal => ptr.as_ptr(),
145+
Less => match self.grow(ptr, layout, new_layout) {
146+
Ok(ptr) => ptr.as_ptr() as *mut u8,
147+
Err(_) => null_mut(),
148+
},
149+
Greater => match self.shrink(ptr, layout, new_layout) {
150+
Ok(ptr) => ptr.as_ptr() as *mut u8,
151+
Err(_) => null_mut(),
152+
},
133153
}
134154
}
135155
}

‎src/common/base/src/mem_allocator/jemalloc.rs‎

Lines changed: 4 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,10 @@
2121
pub struct JEAllocator;
2222

2323
impl JEAllocator {
24+
pub const fn create() -> JEAllocator {
25+
JEAllocator
26+
}
27+
2428
pub fn name() -> String {
2529
"jemalloc".to_string()
2630
}
@@ -44,7 +48,6 @@ pub mod linux {
4448
use tikv_jemalloc_sys as ffi;
4549

4650
use super::JEAllocator;
47-
use crate::runtime::ThreadTracker;
4851

4952
#[cfg(any(target_arch = "arm", target_arch = "mips", target_arch = "powerpc"))]
5053
const ALIGNOF_MAX_ALIGN_T: usize = 8;
@@ -77,8 +80,6 @@ pub mod linux {
7780
unsafe impl Allocator for JEAllocator {
7881
#[inline(always)]
7982
fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
80-
ThreadTracker::alloc(layout.size() as i64)?;
81-
8283
let data_address = if layout.size() == 0 {
8384
unsafe { NonNull::new(layout.align() as *mut ()).unwrap_unchecked() }
8485
} else {
@@ -92,8 +93,6 @@ pub mod linux {
9293

9394
#[inline(always)]
9495
fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
95-
ThreadTracker::alloc(layout.size() as i64)?;
96-
9796
let data_address = if layout.size() == 0 {
9897
unsafe { NonNull::new(layout.align() as *mut ()).unwrap_unchecked() }
9998
} else {
@@ -108,8 +107,6 @@ pub mod linux {
108107

109108
#[inline(always)]
110109
unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
111-
ThreadTracker::dealloc(layout.size() as i64);
112-
113110
if layout.size() == 0 {
114111
debug_assert_eq!(ptr.as_ptr() as usize, layout.align());
115112
} else {
@@ -127,9 +124,6 @@ pub mod linux {
127124
debug_assert_eq!(old_layout.align(), new_layout.align());
128125
debug_assert!(old_layout.size() <= new_layout.size());
129126

130-
ThreadTracker::dealloc(old_layout.size() as i64);
131-
ThreadTracker::alloc(new_layout.size() as i64)?;
132-
133127
let data_address = if new_layout.size() == 0 {
134128
NonNull::new(new_layout.align() as *mut ()).unwrap_unchecked()
135129
} else if old_layout.size() == 0 {
@@ -156,9 +150,6 @@ pub mod linux {
156150
debug_assert_eq!(old_layout.align(), new_layout.align());
157151
debug_assert!(old_layout.size() <= new_layout.size());
158152

159-
ThreadTracker::dealloc(old_layout.size() as i64);
160-
ThreadTracker::alloc(new_layout.size() as i64)?;
161-
162153
let data_address = if new_layout.size() == 0 {
163154
NonNull::new(new_layout.align() as *mut ()).unwrap_unchecked()
164155
} else if old_layout.size() == 0 {
@@ -195,9 +186,6 @@ pub mod linux {
195186
debug_assert_eq!(old_layout.align(), new_layout.align());
196187
debug_assert!(old_layout.size() >= new_layout.size());
197188

198-
ThreadTracker::dealloc(old_layout.size() as i64);
199-
ThreadTracker::alloc(new_layout.size() as i64)?;
200-
201189
if old_layout.size() == 0 {
202190
debug_assert_eq!(ptr.as_ptr() as usize, old_layout.align());
203191
let slice = std::slice::from_raw_parts_mut(ptr.as_ptr(), 0);

0 commit comments

Comments
 (0)