Skip to content

fix(cpp): preserve full-size tablet allocations and report OOM #1003

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ColinLeeo wants to merge 2 commits into
base: develop
Choose a base branch
from
Open

fix(cpp): preserve full-size tablet allocations and report OOM #1003

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bf4ee84
fix(cpp): prevent oversized tablet allocations from corrupting memory
ColinLeeo Oct 10, 2026
41340a8
fix(cpp): allocate tablet buffers at full size and report OOM
ColinLeeo Oct 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions cpp/src/common/allocator/alloc_base.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,11 +62,17 @@ enum AllocModID {
extern TSFILE_API const char* g_mod_names[__LAST_MOD_ID];

/* very basic alloc/free interface in C style */
void* mem_alloc(uint32_t size, AllocModID mid);
void* mem_alloc(size_t size, AllocModID mid);
void mem_free(void* ptr);
void* mem_realloc(void* ptr, uint32_t size);
void* mem_realloc(void* ptr, size_t size);
#ifdef ENABLE_TEST
TSFILE_API void TEST_fail_next_mem_realloc();
// Fail one allocation of mid after skipping successful_allocations calls on
// the current thread. Other allocation modules are unaffected.
TSFILE_API void TEST_fail_mem_alloc_after(AllocModID mid,
uint32_t successful_allocations);
// Byte count received by the latest injected alloc/realloc failure.
TSFILE_API size_t TEST_get_failed_mem_alloc_size();
#endif

class ModStat {
Expand Down Expand Up @@ -127,7 +133,7 @@ class ModStat {
/* base allocator */
class BaseAllocator {
public:
void* alloc(uint32_t size, AllocModID mid) { return mem_alloc(size, mid); }
void* alloc(size_t size, AllocModID mid) { return mem_alloc(size, mid); }
void free(void* ptr) { mem_free(ptr); }
};

Expand Down
44 changes: 39 additions & 5 deletions cpp/src/common/allocator/mem_alloc.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
#include <atomic>
#include <iomanip>
#include <iostream>
#include <limits>

#include "alloc_base.h"
#include "common/logger/elog.h"
Expand All @@ -37,11 +38,24 @@ namespace common {
#ifdef ENABLE_TEST
namespace {
std::atomic<bool> g_fail_next_mem_realloc(false);
}
thread_local AllocModID g_fail_mem_alloc_mod = __LAST_MOD_ID;
thread_local uint32_t g_mem_alloc_calls_to_skip = 0;
thread_local size_t g_failed_mem_alloc_size = 0;
} // namespace

void TEST_fail_next_mem_realloc() {
g_failed_mem_alloc_size = 0;
g_fail_next_mem_realloc.store(true, std::memory_order_release);
}

void TEST_fail_mem_alloc_after(AllocModID mid,
uint32_t successful_allocations) {
g_fail_mem_alloc_mod = mid;
g_mem_alloc_calls_to_skip = successful_allocations;
g_failed_mem_alloc_size = 0;
}

size_t TEST_get_failed_mem_alloc_size() { return g_failed_mem_alloc_size; }
#endif

const char* g_mod_names[__LAST_MOD_ID] = {
Expand Down Expand Up @@ -89,12 +103,27 @@ constexpr size_t HEADER_PTR_SIZE = 8;
// constexpr size_t ALIGNMENT = alignof(std::max_align_t);
constexpr size_t ALIGNMENT = 8;

void* mem_alloc(uint32_t size, AllocModID mid) {
void* mem_alloc(size_t size, AllocModID mid) {
#ifdef ENABLE_TEST
if (mid == g_fail_mem_alloc_mod) {
if (g_mem_alloc_calls_to_skip == 0) {
g_fail_mem_alloc_mod = __LAST_MOD_ID;
g_failed_mem_alloc_size = size;
return nullptr;
}
--g_mem_alloc_calls_to_skip;
}
#endif
// use 7bit at most
ASSERT(mid <= 127);
static_assert(HEADER_PTR_SIZE <= ALIGNMENT,
"Header must fit within alignment");
constexpr size_t header_size = ALIGNMENT;
// The header stores the byte count in its upper 56 bits.
if (size > (UINT64_MAX >> 8) ||
size > std::numeric_limits<size_t>::max() - header_size) {
return nullptr;
}
const size_t total_size = size + header_size;
auto raw = static_cast<char*>(malloc(total_size));
if (UNLIKELY(raw == nullptr)) {
Expand Down Expand Up @@ -144,24 +173,29 @@ void mem_free(void* ptr) {
static_cast<uint64_t>(*reinterpret_cast<uint32_t*>(raw_ptr + 4)) |
(static_cast<uint64_t>(*reinterpret_cast<uint32_t*>(raw_ptr)) << 32);
auto mid = static_cast<AllocModID>(header & 0x7F);
auto size = static_cast<uint32_t>(header >> 8);
const uint64_t size = header >> 8;
ModStat::get_instance().update_free(mid, size);
::free(raw_ptr);
}

void* mem_realloc(void* ptr, uint32_t size) {
void* mem_realloc(void* ptr, size_t size) {
#ifdef ENABLE_TEST
if (g_fail_next_mem_realloc.exchange(false, std::memory_order_acq_rel)) {
g_failed_mem_alloc_size = size;
return nullptr;
}
#endif
if (size > (UINT64_MAX >> 8) ||
size > std::numeric_limits<size_t>::max() - ALIGNMENT) {
return nullptr;
}
char* p = static_cast<char*>(ptr);
char* raw_ptr = p - ALIGNMENT;
const uint64_t header =
static_cast<uint64_t>(*reinterpret_cast<uint32_t*>(raw_ptr + 4)) |
(static_cast<uint64_t>(*reinterpret_cast<uint32_t*>(raw_ptr)) << 32);
auto mid = static_cast<AllocModID>(header & 0x7F);
auto original_size = static_cast<uint32_t>(header >> 8);
const uint64_t original_size = header >> 8;
p = static_cast<char*>(realloc(raw_ptr, size + ALIGNMENT));
if (UNLIKELY(p == nullptr)) {
return nullptr;
Expand Down
3 changes: 2 additions & 1 deletion cpp/src/common/container/bit_map.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,9 @@ BitMap::~BitMap() {
}

int BitMap::init(uint32_t item_size, bool init_as_zero, AllocModID mod_id) {
uint32_t size = (item_size + 7) / 8;
uint32_t size = item_size / 8 + (item_size % 8 != 0);
bitmap_ = static_cast<char*>(mem_alloc(size, mod_id));
if (bitmap_ == nullptr) return common::E_OOM;
// need set to 0, otherwise there will be wrong data
const char initial_char = init_as_zero ? 0x00 : 0xFF;
memset(bitmap_, initial_char, size);
Expand Down
2 changes: 1 addition & 1 deletion cpp/src/common/container/bit_map.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,7 @@ class BitMap {
if (byte_val) {
return from + bitops::ctz_nonzero(static_cast<uint32_t>(byte_val));
}
const uint32_t byte_end = (total_bits + 7) >> 3;
const uint32_t byte_end = total_bits / 8 + (total_bits % 8 != 0);
for (++byte_idx; byte_idx < byte_end; ++byte_idx) {
if (p[byte_idx]) {
uint32_t pos =
Expand Down
70 changes: 48 additions & 22 deletions cpp/src/common/tablet.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,12 +33,16 @@ namespace storage {

int Tablet::init() {
ASSERT(timestamps_ == nullptr);
timestamps_ = static_cast<int64_t*>(
common::mem_alloc(sizeof(int64_t) * max_row_num_, common::MOD_TABLET));
if (timestamps_ == nullptr) return E_OOM;
cur_row_size_ = 0;
if (max_row_num_ == 0) return E_INVALID_ARG;
if (static_cast<uint64_t>(max_row_num_) > SIZE_MAX / sizeof(int64_t))
return E_OOM;

size_t schema_count = schema_vec_->size();
if (schema_count > SIZE_MAX / sizeof(ValueMatrixEntry) ||
schema_count > SIZE_MAX / sizeof(BitMap)) {
return E_OOM;
}
std::pair<std::map<std::string, int>::iterator, bool> ins_res;
for (size_t c = 0; c < schema_count; c++) {
ins_res = schema_map_.insert(
Expand All @@ -48,9 +52,16 @@ int Tablet::init() {
}
}
ASSERT(schema_map_.size() == schema_count);
// Finish schema-map allocations before acquiring the owned raw buffers.
timestamps_ = static_cast<int64_t*>(
common::mem_alloc(sizeof(int64_t) * max_row_num_, common::MOD_TABLET));
if (timestamps_ == nullptr) return E_OOM;
value_matrix_ = static_cast<ValueMatrixEntry*>(common::mem_alloc(
sizeof(ValueMatrixEntry) * schema_count, common::MOD_TABLET));
if (value_matrix_ == nullptr) return E_OOM;
for (size_t c = 0; c < schema_count; ++c) {
new (&value_matrix_[c]) ValueMatrixEntry();
}
for (size_t c = 0; c < schema_count; ++c) {
const MeasurementSchema& schema = schema_vec_->at(c);

Expand Down Expand Up @@ -106,8 +117,9 @@ int Tablet::init() {
return E_OOM;
}
auto* sc = new (mem) StringColumn();
sc->init(max_row_num_, max_row_num_ * 32);
value_matrix_[c].string_col = sc;
int ret = sc->init(max_row_num_);
if (ret != E_OK) return ret;
break;
}
default:
Expand All @@ -121,7 +133,10 @@ int Tablet::init() {
if (bitmaps_ == nullptr) return E_OOM;
for (size_t c = 0; c < schema_count; c++) {
new (&bitmaps_[c]) BitMap();
bitmaps_[c].init(max_row_num_, false, common::MOD_TABLET);
}
for (size_t c = 0; c < schema_count; c++) {
int ret = bitmaps_[c].init(max_row_num_, false, common::MOD_TABLET);
if (ret != E_OK) return ret;
}

return E_OK;
Expand All @@ -139,26 +154,33 @@ void Tablet::destroy() {
switch (schema.data_type_) {
case DATE:
case INT32:
common::mem_free(value_matrix_[c].int32_data);
if (value_matrix_[c].int32_data)
common::mem_free(value_matrix_[c].int32_data);
break;
case TIMESTAMP:
case INT64:
common::mem_free(value_matrix_[c].int64_data);
if (value_matrix_[c].int64_data)
common::mem_free(value_matrix_[c].int64_data);
break;
case FLOAT:
common::mem_free(value_matrix_[c].float_data);
if (value_matrix_[c].float_data)
common::mem_free(value_matrix_[c].float_data);
break;
case DOUBLE:
common::mem_free(value_matrix_[c].double_data);
if (value_matrix_[c].double_data)
common::mem_free(value_matrix_[c].double_data);
break;
case BOOLEAN:
common::mem_free(value_matrix_[c].bool_data);
if (value_matrix_[c].bool_data)
common::mem_free(value_matrix_[c].bool_data);
break;
case BLOB:
case TEXT:
case STRING:
value_matrix_[c].string_col->destroy();
common::mem_free(value_matrix_[c].string_col);
if (value_matrix_[c].string_col) {
value_matrix_[c].string_col->destroy();
common::mem_free(value_matrix_[c].string_col);
}
break;
default:
break;
Expand Down Expand Up @@ -249,7 +271,7 @@ int Tablet::set_column_values(uint32_t schema_index, const void* data,
// get_bitmap() would leave the flag stale (e.g. cleared by a prior
// clear_all()) and downstream may_have_set_bits() checks would skip
// null-mask handling for the column.
uint32_t bm_bytes = (count + 7) / 8;
uint32_t bm_bytes = count / 8 + (count % 8 != 0);
bitmaps_[schema_index].copy_from(reinterpret_cast<const char*>(bitmap),
bm_bytes);
}
Expand Down Expand Up @@ -305,13 +327,14 @@ int Tablet::set_column_string_values(uint32_t schema_index,
if (total_bytes > 0) {
std::memcpy(sc->buffer, data, total_bytes);
}
std::memcpy(sc->offsets, offsets, (count + 1) * sizeof(int32_t));
std::memcpy(sc->offsets, offsets,
(static_cast<size_t>(count) + 1) * sizeof(int32_t));
sc->buf_used = total_bytes;

if (bitmap == nullptr) {
bitmaps_[schema_index].clear_all();
} else {
uint32_t bm_bytes = (count + 7) / 8;
uint32_t bm_bytes = count / 8 + (count % 8 != 0);
bitmaps_[schema_index].copy_from(reinterpret_cast<const char*>(bitmap),
bm_bytes);
}
Expand All @@ -335,12 +358,11 @@ int Tablet::set_column_string_repeated(uint32_t schema_index, const char* str,
StringColumn* sc = value_matrix_[schema_index].string_col;
if (sc == nullptr) return E_INVALID_ARG;

// str_len * count can overflow uint32_t; do the multiply in uint64_t and
// reject anything that wouldn't fit, otherwise the subsequent loop would
// walk past the truncated buf_capacity allocation.
// Multiply in uint64_t and require the total to fit signed Arrow offsets
// before allocating or writing any data.
uint64_t total_bytes_64 =
static_cast<uint64_t>(str_len) * static_cast<uint64_t>(count);
if (total_bytes_64 > std::numeric_limits<uint32_t>::max()) {
if (total_bytes_64 > std::numeric_limits<int32_t>::max()) {
return E_OVERFLOW;
}
uint32_t total_bytes = static_cast<uint32_t>(total_bytes_64);
Expand All @@ -364,6 +386,7 @@ int Tablet::set_column_string_repeated(uint32_t schema_index, const char* str,
}

void Tablet::reset(uint32_t row_count) {
if (err_code_ != E_OK) return;
ASSERT(row_count <= max_row_num_);
cur_row_size_ = row_count;
reset_string_columns();
Expand All @@ -379,8 +402,10 @@ void Tablet::reset(uint32_t row_count) {
}
}

void* Tablet::get_value(int row_index, uint32_t schema_index,
void* Tablet::get_value(uint32_t row_index, uint32_t schema_index,
common::TSDataType& data_type) const {
if (err_code_ != E_OK) return nullptr;
if (UNLIKELY(row_index >= max_row_num_)) return nullptr;
if (UNLIKELY(schema_index >= schema_vec_->size())) {
return nullptr;
}
Expand Down Expand Up @@ -580,6 +605,7 @@ void Tablet::set_column_categories(
}

void Tablet::reset_string_columns() {
if (err_code_ != E_OK) return;
size_t schema_count = schema_vec_->size();
for (size_t c = 0; c < schema_count; c++) {
const MeasurementSchema& schema = schema_vec_->at(c);
Expand All @@ -595,7 +621,7 @@ std::vector<uint32_t> Tablet::find_all_device_boundaries() const {
if (row_count <= 1) return {};

// Use uint64_t bitmap instead of vector<bool> for faster set/test/scan.
const uint32_t nwords = (row_count + 63) / 64;
const uint32_t nwords = row_count / 64 + (row_count % 64 != 0);
std::vector<uint64_t> boundary(nwords, 0);

// Walk id columns RIGHT to LEFT. In time-series tag systems the rightmost
Expand Down Expand Up @@ -666,7 +692,7 @@ std::vector<uint32_t> Tablet::find_all_device_boundaries() const {
return result;
}

std::shared_ptr<IDeviceID> Tablet::get_device_id(int i) const {
std::shared_ptr<IDeviceID> Tablet::get_device_id(uint32_t i) const {
std::vector<std::string*> id_array;
id_array.push_back(new std::string(insert_target_name_));
for (auto id_column_idx : id_column_indexes_) {
Expand Down
Loading
Loading