Merge pull request #7971 from usamoi/main

refactor(hashtable): adaptive string hash table

src/common/base/Cargo.toml

@@ -37,6 +37,7 @@ pprof = { version = "0.10.1", features = [
     "protobuf-codec",
     "protobuf",
 ] }
+semver = "1.0.10"
 serde = { workspace = true }
 tikv-jemalloc-ctl = { version = "0.4.2", optional = true }
 tikv-jemalloc-sys = "0.4.3"

src/common/base/src/base/runtime_tracker.rs

@@ -18,7 +18,7 @@ use std::sync::atomic::AtomicI64;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
 
-use crate::mem_allocator::ALLOC;
+use crate::mem_allocator::GlobalAllocator;
 
 #[thread_local]
 static mut TRACKER: *mut ThreadTracker = std::ptr::null_mut();
@@ -92,12 +92,15 @@ impl ThreadTracker {
     }
 
     #[inline]
-    pub fn realloc_memory(old_size: i64, new_size: i64) {
-        let addition = new_size - old_size;
-        match addition > 0 {
-            true => Self::alloc_memory(addition),
-            false => Self::dealloc_memory(-addition),
-        }
+    pub fn grow_memory(old_size: i64, new_size: i64) {
+        assert!(old_size <= new_size);
+        Self::alloc_memory(new_size - old_size)
+    }
+
+    #[inline]
+    pub fn shrink_memory(old_size: i64, new_size: i64) {
+        assert!(old_size >= new_size);
+        Self::dealloc_memory(old_size - new_size)
     }
 }
 
@@ -171,7 +174,7 @@ impl RuntimeTracker {
             let tracker = std::mem::replace(&mut TRACKER, std::ptr::null_mut());
 
             std::ptr::drop_in_place(tracker as usize as *mut ThreadTracker);
-            ALLOC.dealloc(tracker as *mut u8, Layout::new::<ThreadTracker>())
+            GlobalAllocator.dealloc(tracker as *mut u8, Layout::new::<ThreadTracker>())
         }
     }
 

src/common/base/src/lib.rs

@@ -12,7 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#![feature(allocator_api)]
 #![feature(thread_local)]
+#![feature(ptr_metadata)]
 #![feature(result_flattening)]
 #![feature(try_trait_v2)]
 #![allow(incomplete_features)]

src/common/base/src/mem_allocator/global_allocator.rs

@@ -0,0 +1,128 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::alloc::AllocError;
+use std::alloc::Allocator;
+use std::alloc::GlobalAlloc;
+use std::alloc::Layout;
+use std::ptr::null_mut;
+use std::ptr::NonNull;
+
+use super::je_allocator::JEAllocator;
+use super::system_allocator::SystemAllocator;
+
+#[global_allocator]
+pub static GLOBAL_ALLOCATOR: GlobalAllocator = GlobalAllocator;
+
+#[derive(Debug, Clone, Copy, Default)]
+pub struct GlobalAllocator;
+
+type Fallback = JEAllocator<SystemAllocator>;
+
+unsafe impl Allocator for GlobalAllocator {
+    #[inline(always)]
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        Fallback::new(Default::default()).allocate(layout)
+    }
+
+    #[inline(always)]
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        Fallback::new(Default::default()).deallocate(ptr, layout)
+    }
+
+    #[inline(always)]
+    fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        Fallback::new(Default::default()).allocate_zeroed(layout)
+    }
+
+    #[inline(always)]
+    unsafe fn grow(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        Fallback::new(Default::default()).grow(ptr, old_layout, new_layout)
+    }
+
+    #[inline(always)]
+    unsafe fn grow_zeroed(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        Fallback::new(Default::default()).grow_zeroed(ptr, old_layout, new_layout)
+    }
+
+    #[inline(always)]
+    unsafe fn shrink(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        Fallback::new(Default::default()).shrink(ptr, old_layout, new_layout)
+    }
+}
+
+unsafe impl GlobalAlloc for GlobalAllocator {
+    #[inline]
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        if let Ok(ptr) = GlobalAllocator.allocate(layout) {
+            ptr.as_ptr() as *mut u8
+        } else {
+            null_mut()
+        }
+    }
+
+    #[inline]
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        let ptr = NonNull::new(ptr).unwrap_unchecked();
+        GlobalAllocator.deallocate(ptr, layout);
+    }
+
+    #[inline]
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        if let Ok(ptr) = GlobalAllocator.allocate_zeroed(layout) {
+            ptr.as_ptr() as *mut u8
+        } else {
+            null_mut()
+        }
+    }
+
+    #[inline]
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        use std::cmp::Ordering::*;
+        let ptr = NonNull::new(ptr).unwrap_unchecked();
+        let new_layout = Layout::from_size_align(new_size, layout.align()).unwrap();
+        match layout.size().cmp(&new_size) {
+            Less => {
+                if let Ok(ptr) = GlobalAllocator.grow(ptr, layout, new_layout) {
+                    ptr.as_ptr() as *mut u8
+                } else {
+                    null_mut()
+                }
+            }
+            Greater => {
+                if let Ok(ptr) = GlobalAllocator.shrink(ptr, layout, new_layout) {
+                    ptr.as_ptr() as *mut u8
+                } else {
+                    null_mut()
+                }
+            }
+            Equal => ptr.as_ptr() as *mut u8,
+        }
+    }
+}