[Enhancement] defensive strategy on on large mem allocation (#57041)

Signed-off-by: Kevin Cai <[email protected]>
(cherry picked from commit b9075b7)

Author: kevincai

be/src/common/config.h

@@ -83,6 +83,10 @@ CONF_String(mem_limit, "90%");
 // Enable the jemalloc tracker, which is responsible for reserving memory
 CONF_Bool(enable_jemalloc_memory_tracker, "true");
 
+// Whether abort the process if a large memory allocation is detected which the requested
+// size is larger than the available physical memory without wrapping with TRY_CATCH_BAD_ALLOC
+CONF_mBool(abort_on_large_memory_allocation, "false");
+
 // The port heartbeat service used.
 CONF_Int32(heartbeat_service_port, "9050");
 // The count of heart beat service.

be/src/common/daemon.cpp

@@ -52,6 +52,7 @@
 #include "runtime/time_types.h"
 #include "runtime/user_function_cache.h"
 #include "service/backend_options.h"
+#include "service/mem_hook.h"
 #include "storage/options.h"
 #include "storage/storage_engine.h"
 #include "util/cpu_info.h"
@@ -345,6 +346,12 @@ void Daemon::init(bool as_cn, const std::vector<StorePath>& paths) {
 
     init_signals();
     init_minidump();
+
+    // Don't bother set the limit if the process is running with very limited memory capacity
+    if (MemInfo::physical_mem() > 1024 * 1024 * 1024) {
+        // set mem hook to reject the memory allocation if large than available physical memory detected.
+        set_large_memory_alloc_failure_threshold(MemInfo::physical_mem());
+    }
 }
 
 void Daemon::stop() {

be/src/service/mem_hook.cpp

@@ -12,9 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "mem_hook.h"
+
 #include <iostream>
 
 #include "common/compiler_util.h"
+#include "common/config.h"
 #include "glog/logging.h"
 #include "jemalloc/jemalloc.h"
 #include "runtime/current_thread.h"
@@ -92,6 +95,17 @@ std::atomic<int64_t> g_mem_usage(0);
 #define IS_BAD_ALLOC_CATCHED() false
 #endif
 
+static int64_t g_large_memory_alloc_failure_threshold = 0;
+
+namespace starrocks {
+// thread-safety is not a concern here since this is a really rare op
+int64_t set_large_memory_alloc_failure_threshold(int64_t val) {
+    int64_t old_val = g_large_memory_alloc_failure_threshold;
+    g_large_memory_alloc_failure_threshold = val;
+    return old_val;
+}
+} // namespace starrocks
+
 const size_t large_memory_alloc_report_threshold = 1073741824;
 inline thread_local bool skip_report = false;
 inline void report_large_memory_alloc(size_t size) {
@@ -101,7 +115,8 @@ inline void report_large_memory_alloc(size_t size) {
             auto qid = starrocks::CurrentThread::current().query_id();
             auto fid = starrocks::CurrentThread::current().fragment_instance_id();
             LOG(WARNING) << "large memory alloc, query_id:" << print_id(qid) << " instance: " << print_id(fid)
-                         << " acquire:" << size << " bytes, stack:\n"
+                         << " acquire:" << size << " bytes, is_bad_alloc_caught: " << IS_BAD_ALLOC_CATCHED()
+                         << ", stack:\n"
                          << starrocks::get_stack_trace();
         } catch (...) {
             // do nothing
@@ -111,6 +126,18 @@ inline void report_large_memory_alloc(size_t size) {
 }
 #define STARROCKS_REPORT_LARGE_MEM_ALLOC(size) report_large_memory_alloc(size)
 
+inline bool block_large_memory_alloc(size_t size) {
+    if (UNLIKELY(g_large_memory_alloc_failure_threshold > 0 && size > g_large_memory_alloc_failure_threshold)) {
+        // DON'T try to allocate the memory at all
+        if (starrocks::config::abort_on_large_memory_allocation) {
+            std::abort();
+        } else {
+            return true;
+        }
+    }
+    return false;
+}
+
 DEFINE_SCOPED_FAIL_POINT(mem_alloc_error);
 
 #ifdef FIU_ENABLE
@@ -142,6 +169,10 @@ void* my_malloc(size_t size) __THROW {
         }
         return ptr;
     } else {
+        if (UNLIKELY(block_large_memory_alloc(size))) {
+            return nullptr;
+        }
+
         void* ptr = STARROCKS_MALLOC(size);
         // NOTE: do NOT call `tc_malloc_size` here, it may call the new operator, which in turn will
         // call the `my_malloc`, and result in a deadloop.
@@ -228,6 +259,9 @@ void* my_calloc(size_t n, size_t size) __THROW {
         }
         return ptr;
     } else {
+        if (UNLIKELY(block_large_memory_alloc(n * size))) {
+            return nullptr;
+        }
         void* ptr = STARROCKS_CALLOC(n, size);
         int64_t alloc_size = STARROCKS_MALLOC_SIZE(ptr);
         MEMORY_CONSUME_SIZE(alloc_size);

be/src/service/mem_hook.h

@@ -0,0 +1,23 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// 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
+//
+//     https://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.
+
+#pragma once
+
+#include <cstdint>
+
+namespace starrocks {
+
+int64_t set_large_memory_alloc_failure_threshold(int64_t);
+
+} // namespace starrocks

be/test/CMakeLists.txt

@@ -486,6 +486,7 @@ set(EXEC_FILES
         ./gutil/cpu_test.cc
         ./gutil/sysinfo-test.cc
         ./service/lake_service_test.cpp
+        ./service/mem_hook_test.cpp
         ./service/service_be/internal_service_test.cpp
         ./storage/metadata_util_test.cpp
         ./storage/delta_writer_test.cpp

be/test/service/mem_hook_test.cpp

@@ -0,0 +1,78 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// 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
+//
+//     https://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.
+
+// FIXME: The case doesn't work on ASAN case since the malloc will be hooked
+// when doing the ASAN init. Can't make it work so just disable it for now.
+#ifndef ADDRESS_SANITIZER
+
+#include "service/mem_hook.h"
+
+#include <gtest/gtest.h>
+
+#include <vector>
+
+namespace starrocks {
+
+static void try_malloc_memory(size_t size) {
+    std::vector<int8_t> arr;
+    arr.reserve(size);
+    EXPECT_EQ(size, arr.capacity());
+}
+
+static void try_calloc_memory(size_t n, size_t size) {
+    void* ptr = calloc(n, size);
+    if (ptr == nullptr) {
+        throw std::bad_alloc();
+    }
+    free(ptr);
+}
+
+TEST(MemhookTest, test_malloc_mem_hook_block_without_try_catch) {
+    set_large_memory_alloc_failure_threshold(0);
+    int64_t blocking_size = 1024 * 1024;
+
+    set_large_memory_alloc_failure_threshold(blocking_size);
+    // successful because blocking_size <= g_large_memory_alloc_failure_threshold
+    EXPECT_NO_THROW(try_malloc_memory(blocking_size));
+
+    set_large_memory_alloc_failure_threshold(blocking_size - 1);
+    // fail with std::bad_alloc for the same size
+    EXPECT_THROW(try_malloc_memory(blocking_size), std::bad_alloc);
+
+    // reset to no limit
+    set_large_memory_alloc_failure_threshold(0);
+}
+
+TEST(MemhookTest, test_malloc_mem_hook_block_with_try_catch) {
+    // IS_BAD_ALLOC_CATCHED is always `false` when builds with -DBE_TEST.
+    // There is no easy way to simulate the mem_tracker limit check here.
+    // Leave the test body empty as intended for now.
+}
+
+TEST(MemhookTest, test_calloc_mem_hook_block_without_try_catch) {
+    int64_t blocking_size = 1024 * 1024;
+
+    set_large_memory_alloc_failure_threshold(blocking_size);
+    // successful because blocking_size <= g_large_memory_alloc_failure_threshold
+    EXPECT_NO_THROW(try_calloc_memory(blocking_size / sizeof(int), sizeof(int)));
+
+    set_large_memory_alloc_failure_threshold(blocking_size - 1);
+    // fail with std::bad_alloc for the same size
+    EXPECT_THROW(try_calloc_memory(blocking_size / sizeof(int), sizeof(int)), std::bad_alloc);
+
+    // reset to no limit
+    set_large_memory_alloc_failure_threshold(0);
+}
+} // namespace starrocks
+#endif