From bf4ee8439e377380b2ac574733579b390baf37f0 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Sat, 10 Oct 2026 11:52:00 +0800 Subject: [PATCH 1/2] fix(cpp): prevent oversized tablet allocations from corrupting memory --- cpp/src/common/allocator/alloc_base.h | 4 + cpp/src/common/allocator/mem_alloc.cc | 19 ++- cpp/src/common/container/bit_map.cc | 1 + cpp/src/common/tablet.cc | 57 +++++--- cpp/src/common/tablet.h | 50 +++---- cpp/src/cwrapper/tsfile_cwrapper.cc | 29 ++-- cpp/src/cwrapper/tsfile_cwrapper.h | 5 +- cpp/test/common/tablet_test.cc | 125 ++++++++++++++++++ .../cwrapper/cwrapper_public_writer_test.cc | 52 +++++++- go/tsfile/api_test.go | 19 +++ go/tsfile/cgo_bridge.go | 10 +- go/tsfile/cgo_bridge_test.go | 6 +- go/tsfile/tablet.go | 2 + 13 files changed, 318 insertions(+), 61 deletions(-) diff --git a/cpp/src/common/allocator/alloc_base.h b/cpp/src/common/allocator/alloc_base.h index 2b97c8f83..82bc734ea 100644 --- a/cpp/src/common/allocator/alloc_base.h +++ b/cpp/src/common/allocator/alloc_base.h @@ -67,6 +67,10 @@ void mem_free(void* ptr); void* mem_realloc(void* ptr, uint32_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); #endif class ModStat { diff --git a/cpp/src/common/allocator/mem_alloc.cc b/cpp/src/common/allocator/mem_alloc.cc index deda0b8d8..9d9a5222c 100644 --- a/cpp/src/common/allocator/mem_alloc.cc +++ b/cpp/src/common/allocator/mem_alloc.cc @@ -37,11 +37,19 @@ namespace common { #ifdef ENABLE_TEST namespace { std::atomic 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; +} // namespace void TEST_fail_next_mem_realloc() { 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; +} #endif const char* g_mod_names[__LAST_MOD_ID] = { @@ -90,6 +98,15 @@ constexpr size_t HEADER_PTR_SIZE = 8; constexpr size_t ALIGNMENT = 8; void* mem_alloc(uint32_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; + return nullptr; + } + --g_mem_alloc_calls_to_skip; + } +#endif // use 7bit at most ASSERT(mid <= 127); static_assert(HEADER_PTR_SIZE <= ALIGNMENT, diff --git a/cpp/src/common/container/bit_map.cc b/cpp/src/common/container/bit_map.cc index 3b1af6ab2..f786226f0 100644 --- a/cpp/src/common/container/bit_map.cc +++ b/cpp/src/common/container/bit_map.cc @@ -34,6 +34,7 @@ BitMap::~BitMap() { int BitMap::init(uint32_t item_size, bool init_as_zero, AllocModID mod_id) { uint32_t size = (item_size + 7) / 8; bitmap_ = static_cast(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); diff --git a/cpp/src/common/tablet.cc b/cpp/src/common/tablet.cc index 3bfc54bb9..74627ee50 100644 --- a/cpp/src/common/tablet.cc +++ b/cpp/src/common/tablet.cc @@ -33,12 +33,17 @@ namespace storage { int Tablet::init() { ASSERT(timestamps_ == nullptr); - timestamps_ = static_cast( - 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; + // Check before any allocation: otherwise mem_alloc truncates a size_t + // request to uint32_t while memset / memcpy still use the full size. + if (max_row_num_ > MAX_ROWS) return E_OVERFLOW; size_t schema_count = schema_vec_->size(); + if (schema_count > UINT32_MAX / sizeof(ValueMatrixEntry) || + schema_count > UINT32_MAX / sizeof(BitMap)) { + return E_OVERFLOW; + } std::pair::iterator, bool> ins_res; for (size_t c = 0; c < schema_count; c++) { ins_res = schema_map_.insert( @@ -48,9 +53,16 @@ int Tablet::init() { } } ASSERT(schema_map_.size() == schema_count); + // Finish schema-map allocations before acquiring the owned raw buffers. + timestamps_ = static_cast( + common::mem_alloc(sizeof(int64_t) * max_row_num_, common::MOD_TABLET)); + if (timestamps_ == nullptr) return E_OOM; value_matrix_ = static_cast(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); @@ -106,8 +118,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: @@ -121,7 +134,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; @@ -139,26 +155,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; @@ -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(str_len) * static_cast(count); - if (total_bytes_64 > std::numeric_limits::max()) { + if (total_bytes_64 > std::numeric_limits::max()) { return E_OVERFLOW; } uint32_t total_bytes = static_cast(total_bytes_64); @@ -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(); @@ -381,6 +404,7 @@ void Tablet::reset(uint32_t row_count) { void* Tablet::get_value(int row_index, uint32_t schema_index, common::TSDataType& data_type) const { + if (err_code_ != E_OK) return nullptr; if (UNLIKELY(schema_index >= schema_vec_->size())) { return nullptr; } @@ -580,6 +604,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); diff --git a/cpp/src/common/tablet.h b/cpp/src/common/tablet.h index 095157ac7..1bbf1349d 100644 --- a/cpp/src/common/tablet.h +++ b/cpp/src/common/tablet.h @@ -57,13 +57,22 @@ class Tablet { StringColumn() : offsets(nullptr), buffer(nullptr), buf_capacity(0), buf_used(0) {} - void init(uint32_t max_rows, uint32_t init_buf_capacity) { + int init(uint32_t max_rows) { + const uint64_t offsets_bytes = + sizeof(int32_t) * (static_cast(max_rows) + 1); + if (offsets_bytes > UINT32_MAX) return common::E_OVERFLOW; offsets = (int32_t*)common::mem_alloc( - sizeof(int32_t) * (max_rows + 1), common::MOD_TABLET); + static_cast(offsets_bytes), common::MOD_TABLET); + if (offsets == nullptr) return common::E_OOM; offsets[0] = 0; - buf_capacity = init_buf_capacity; + // Estimate 32 bytes per row, but cap eager allocation at 1 MiB. + // Large tablets grow the data buffer as values are appended. + buf_capacity = static_cast(std::min( + static_cast(max_rows) * 32, 1u << 20)); buffer = (char*)common::mem_alloc(buf_capacity, common::MOD_TABLET); + if (buffer == nullptr) return common::E_OOM; buf_used = 0; + return common::E_OK; } void destroy() { @@ -81,13 +90,15 @@ class Tablet { } int append(uint32_t row, const char* data, uint32_t len) { + const uint64_t required = static_cast(buf_used) + len; + // Arrow offsets are signed int32, even though the capacities are + // unsigned. Reject before either the addition or offsets wrap. + if (required > INT32_MAX) return common::E_OVERFLOW; // Grow buffer if needed - if (buf_used + len > buf_capacity) { - uint64_t new_capacity_64 = - static_cast(buf_capacity) * 2 + len; - if (UNLIKELY(new_capacity_64 > UINT32_MAX)) { - return common::E_OVERFLOW; - } + if (required > buf_capacity) { + uint64_t new_capacity_64 = std::min( + std::max(required, static_cast(buf_capacity) * 2), + INT32_MAX); uint32_t new_capacity = static_cast(new_capacity_64); char* new_buffer = (char*)common::mem_realloc(buffer, new_capacity); @@ -135,13 +146,16 @@ class Tablet { public: static const uint32_t DEFAULT_MAX_ROWS = 1024; + // Every tablet has an int64 timestamp buffer allocated through mem_alloc, + // whose byte-size argument is uint32_t. + static const uint32_t MAX_ROWS = UINT32_MAX / sizeof(int64_t); int err_code_ = common::E_OK; public: Tablet(const std::string& device_id, std::shared_ptr> schema_vec, int max_rows = DEFAULT_MAX_ROWS) - : max_row_num_(max_rows), + : max_row_num_(max_rows > 0 ? max_rows : 0), insert_target_name_(device_id), schema_vec_(schema_vec), timestamps_(nullptr), @@ -149,11 +163,6 @@ class Tablet { bitmaps_(nullptr) { ASSERT(device_id.size() >= 1); ASSERT(schema_vec != NULL); - ASSERT(max_rows > 0 && max_rows < (1 << 30)); - if (max_rows < 0) { - ASSERT(false); - max_row_num_ = DEFAULT_MAX_ROWS; - } err_code_ = init(); } @@ -161,7 +170,7 @@ class Tablet { const std::vector* measurement_list, const std::vector* data_type_list, int max_row_num = DEFAULT_MAX_ROWS) - : max_row_num_(max_row_num), + : max_row_num_(max_row_num > 0 ? max_row_num : 0), insert_target_name_(device_id), timestamps_(nullptr), value_matrix_(nullptr), @@ -169,11 +178,6 @@ class Tablet { ASSERT(!device_id.empty()); ASSERT(measurement_list != nullptr); ASSERT(data_type_list != nullptr); - ASSERT(max_row_num > 0 && max_row_num < (1 << 30)); - if (max_row_num < 0) { - ASSERT(false); - max_row_num_ = DEFAULT_MAX_ROWS; - } ASSERT(measurement_list->size() == data_type_list->size()); std::vector measurement_vec; @@ -194,7 +198,7 @@ class Tablet { const std::vector& data_types, const std::vector& column_categories, int max_rows = DEFAULT_MAX_ROWS) - : max_row_num_(max_rows), + : max_row_num_(max_rows > 0 ? max_rows : 0), cur_row_size_(0), insert_target_name_(insert_target_name), timestamps_(nullptr), @@ -407,7 +411,7 @@ class Tablet { template int process_val(uint32_t row_index, uint32_t schema_index, T val); uint32_t max_row_num_; - uint32_t cur_row_size_; + uint32_t cur_row_size_ = 0; std::string insert_target_name_; std::shared_ptr> schema_vec_; std::map schema_map_; diff --git a/cpp/src/cwrapper/tsfile_cwrapper.cc b/cpp/src/cwrapper/tsfile_cwrapper.cc index ff70990ca..c572ec2ae 100644 --- a/cpp/src/cwrapper/tsfile_cwrapper.cc +++ b/cpp/src/cwrapper/tsfile_cwrapper.cc @@ -368,7 +368,11 @@ Tablet tablet_new(char** column_name_list, TSDataType* data_types, } *err_code = common::E_INVALID_ARG; if (column_num == 0 || max_rows == 0 || column_name_list == nullptr || - data_types == nullptr || max_rows >= (1u << 30)) { + data_types == nullptr) { + return nullptr; + } + if (max_rows > storage::Tablet::MAX_ROWS) { + *err_code = common::E_OVERFLOW; return nullptr; } try { @@ -2329,12 +2333,19 @@ Tablet _tablet_new_with_target_name(const char* device_id, data_type_list.push_back( static_cast(*(data_types + i))); } + storage::Tablet* tablet; if (device_id != nullptr) { - return new storage::Tablet(device_id, &measurement_list, - &data_type_list, max_rows); + tablet = new storage::Tablet(device_id, &measurement_list, + &data_type_list, max_rows); } else { - return new storage::Tablet(measurement_list, data_type_list, max_rows); + tablet = + new storage::Tablet(measurement_list, data_type_list, max_rows); + } + if (tablet->err_code_ != common::E_OK) { + delete tablet; + return nullptr; } + return tablet; } ERRNO _tsfile_writer_register_table(TsFileWriter writer, TableSchema* schema) { @@ -2495,11 +2506,11 @@ Tablet tablet_new_with_target_name(const char* target_name, int max_rows) { // Public C entry point: reject arguments that would make the private // delegate read out-of-bounds memory (negative column_num, null lists or - // null column names) before touching caller buffers. The native Tablet - // constructor asserts 0 < max_rows < (1 << 30); reject values outside - // that domain here instead of letting an assertion abort the host - // process. Returns NULL on invalid arguments or allocation failure. - if (column_num < 0 || max_rows <= 0 || max_rows >= (1 << 30)) { + // null column names) before touching caller buffers. Timestamp allocation + // must also fit the native allocator's uint32_t byte-size argument. + // Returns NULL on invalid arguments, overflow or allocation failure. + if (column_num < 0 || max_rows <= 0 || + static_cast(max_rows) > storage::Tablet::MAX_ROWS) { return nullptr; } if (column_num > 0) { diff --git a/cpp/src/cwrapper/tsfile_cwrapper.h b/cpp/src/cwrapper/tsfile_cwrapper.h index 4a87db90c..1065e53ee 100644 --- a/cpp/src/cwrapper/tsfile_cwrapper.h +++ b/cpp/src/cwrapper/tsfile_cwrapper.h @@ -582,7 +582,8 @@ void tsfile_free_tsfile_properties(TsFileProperty* properties, uint32_t length); * @param column_name_list [in] Column names array. Size=column_num. * @param data_types [in] Data types array. Size=column_num. * @param column_num [in] Number of columns. Must be ≥1. - * @param max_rows [in] Pre-allocated row capacity. Must be ≥1. + * @param max_rows [in] Pre-allocated row capacity. Must satisfy + * 1 <= max_rows < 2^29; larger capacities return RET_OVERFLOW. * @param err_code [out] Receives RET_OK on success or the failure code. Must * not be NULL. * @return Tablet Valid handle, or NULL on failure. @@ -1291,7 +1292,7 @@ TsFileGenericWriter tsfile_generic_writer_new(const char* pathname, * when column_num > 0 * @param column_num number of columns; must be >= 0 * @param max_rows maximum number of rows the tablet can hold; must satisfy - * 0 < max_rows < 2^30 + * 0 < max_rows < 2^29 * @return a caller-owned tablet, or NULL if the arguments are invalid or * memory allocation fails */ diff --git a/cpp/test/common/tablet_test.cc b/cpp/test/common/tablet_test.cc index ce75ca2ee..bd14a66f3 100644 --- a/cpp/test/common/tablet_test.cc +++ b/cpp/test/common/tablet_test.cc @@ -24,6 +24,131 @@ namespace storage { +// Keep the old implementation safe to test: if it attempts to allocate before +// rejecting the capacity, fail that allocation instead of reserving GiBs. +TEST(TabletTest, OversizedCapacityRejectedBeforeAllocation) { + const std::vector types = { + common::BOOLEAN, common::INT32, common::DATE, common::FLOAT, + common::INT64, common::TIMESTAMP, common::DOUBLE, common::TEXT, + common::BLOB, common::STRING}; + const uint32_t capacities[] = {1u << 29, (1u << 30) - 2, (1u << 30) - 1, + UINT32_MAX}; + for (auto type : types) { + for (auto capacity : capacities) { + SCOPED_TRACE(::testing::Message() + << "type=" << static_cast(type) + << " rows=" << capacity); + common::TEST_fail_mem_alloc_after(common::MOD_TABLET, 0); + Tablet tablet({"value"}, {type}, capacity); + ASSERT_EQ(tablet.err_code_, common::E_OVERFLOW); + EXPECT_EQ(tablet.get_cur_row_size(), 0u); + EXPECT_EQ(tablet.add_timestamp(0, 1), common::E_OVERFLOW); + // The overflow check must leave the failure injection unconsumed. + Tablet next({"value"}, {type}, 1u); + EXPECT_EQ(next.err_code_, common::E_OOM); + } + } +} + +TEST(TabletTest, CapacityCheckedByEveryConstructor) { + std::vector names = {"value"}; + std::vector types = {common::DOUBLE}; + auto schema = std::make_shared>(); + schema->emplace_back("value", common::DOUBLE, common::PLAIN, + common::UNCOMPRESSED); + for (int capacity : {0, -1, 1 << 29, (1 << 30) - 1}) { + const int expected = + capacity <= 0 ? common::E_INVALID_ARG : common::E_OVERFLOW; + Tablet from_schema("dev", schema, capacity); + EXPECT_EQ(from_schema.err_code_, expected); + Tablet from_lists("dev", &names, &types, capacity); + EXPECT_EQ(from_lists.err_code_, expected); + Tablet with_categories("table", names, types, + {common::ColumnCategory::FIELD}, capacity); + EXPECT_EQ(with_categories.err_code_, expected); + } + Tablet zero(names, types, 0u); + EXPECT_EQ(zero.err_code_, common::E_INVALID_ARG); +} + +TEST(TabletTest, AllocationFailuresReturnOomAndReleasePartialBuffers) { + // timestamps, matrix, INT32, STRING object/offsets/data, DOUBLE, + // BLOB object/offsets/data, bitmap array, and four bitmap buffers. + for (uint32_t fail_after = 0; fail_after < 15; ++fail_after) { + SCOPED_TRACE(fail_after); + const int64_t before = + common::ModStat::get_instance().get_stat(common::MOD_TABLET); + common::TEST_fail_mem_alloc_after(common::MOD_TABLET, fail_after); + { + Tablet tablet( + {"i", "s", "d", "b"}, + {common::INT32, common::STRING, common::DOUBLE, common::BLOB}, + 8u); + EXPECT_EQ(tablet.err_code_, common::E_OOM); + EXPECT_EQ(tablet.get_cur_row_size(), 0u); + EXPECT_EQ(tablet.add_timestamp(0, 1), common::E_OOM); + common::TSDataType type; + EXPECT_EQ(tablet.get_value(0, 0u, type), nullptr); + tablet.reset(); + EXPECT_EQ(tablet.get_cur_row_size(), 0u); + } + EXPECT_EQ(common::ModStat::get_instance().get_stat(common::MOD_TABLET), + before); + // Each injection is one-shot; a subsequent tablet must work normally. + Tablet next({"value"}, {common::INT64}, 1u); + EXPECT_EQ(next.err_code_, common::E_OK); + } +} + +TEST(TabletTest, MillionRowsSupportsNumericAndStringColumns) { + const uint32_t rows = 1u << 20; + Tablet tablet({"d", "text", "blob", "string"}, + {common::DOUBLE, common::TEXT, common::BLOB, common::STRING}, + rows); + ASSERT_EQ(tablet.err_code_, common::E_OK); + const std::string text = "value-中文"; + const common::String value(text); + for (uint32_t row = 0; row < rows; ++row) { + ASSERT_EQ(tablet.add_timestamp(row, row), common::E_OK); + ASSERT_EQ(tablet.add_value(row, 0u, static_cast(row)), + common::E_OK); + for (uint32_t col = 1; col < 4; ++col) { + ASSERT_EQ(tablet.add_value(row, col, value), common::E_OK); + } + } + ASSERT_EQ(tablet.get_cur_row_size(), rows); + for (uint32_t row = 0; row < rows; ++row) { + common::TSDataType type; + auto* number = static_cast(tablet.get_value(row, 0u, type)); + ASSERT_NE(number, nullptr); + ASSERT_EQ(*number, static_cast(row)); + for (uint32_t col = 1; col < 4; ++col) { + auto* stored = + static_cast(tablet.get_value(row, col, type)); + ASSERT_NE(stored, nullptr); + ASSERT_EQ(stored->len_, text.size()); + ASSERT_EQ(memcmp(stored->buf_, text.data(), text.size()), 0); + } + } +} + +TEST(TabletTest, StringBytesMustFitSignedOffsets) { + Tablet tablet({"value"}, {common::STRING}, 2u); + ASSERT_EQ(tablet.err_code_, common::E_OK); + ASSERT_EQ(tablet.add_value(0u, 0u, common::String("x", 1)), common::E_OK); + // This length used to wrap buf_used + len and skip the growth check. + EXPECT_EQ(tablet.add_value(1u, 0u, common::String("x", UINT32_MAX)), + common::E_OVERFLOW); + EXPECT_EQ(tablet.set_column_string_repeated(0u, "x", 1u << 30, 2u), + common::E_OVERFLOW); + common::TSDataType type; + auto* stored = static_cast(tablet.get_value(0u, 0u, type)); + ASSERT_NE(stored, nullptr); + ASSERT_EQ(stored->len_, 1u); + EXPECT_EQ(stored->buf_[0], 'x'); + EXPECT_EQ(tablet.add_value(1u, 0u, common::String("ok", 2)), common::E_OK); +} + TEST(TabletTest, BasicFunctionality) { std::string device_name = "test_device"; std::vector schema_vec; diff --git a/cpp/test/cwrapper/cwrapper_public_writer_test.cc b/cpp/test/cwrapper/cwrapper_public_writer_test.cc index 3df668711..5b08d0e8c 100644 --- a/cpp/test/cwrapper/cwrapper_public_writer_test.cc +++ b/cpp/test/cwrapper/cwrapper_public_writer_test.cc @@ -25,6 +25,8 @@ #include #include +#include "common/allocator/alloc_base.h" + #ifdef _WIN32 #include #else @@ -165,6 +167,54 @@ TEST_F(CWrapperPublicWriterTest, TableWriterRejectsInvalidSchema) { remove(filename.c_str()); } +TEST_F(CWrapperPublicWriterTest, TabletOversizedCapacityReturnsOverflow) { + char name[] = "value"; + char* names[] = {name}; + const TSDataType types[] = {TS_DATATYPE_BOOLEAN, TS_DATATYPE_INT32, + TS_DATATYPE_DATE, TS_DATATYPE_FLOAT, + TS_DATATYPE_INT64, TS_DATATYPE_TIMESTAMP, + TS_DATATYPE_DOUBLE, TS_DATATYPE_TEXT, + TS_DATATYPE_BLOB, TS_DATATYPE_STRING}; + for (TSDataType type : types) { + for (uint32_t rows : + {1u << 29, (1u << 30) - 2, (1u << 30) - 1, 1u << 30, UINT32_MAX}) { + SCOPED_TRACE(::testing::Message() + << "type=" << type << " rows=" << rows); + ERRNO err = RET_OK; + Tablet tablet = tablet_new(names, &type, 1, rows, &err); + EXPECT_EQ(tablet, nullptr); + EXPECT_EQ(err, RET_OVERFLOW); + if (tablet) free_tablet(&tablet); + } + EXPECT_EQ(tablet_new_with_target_name("dev", names, &type, 1, 1 << 29), + nullptr); + } +} + +TEST_F(CWrapperPublicWriterTest, TabletAllocationFailureReturnsNull) { + char i[] = "i", s[] = "s", d[] = "d", b[] = "b"; + char* names[] = {i, s, d, b}; + TSDataType types[] = {TS_DATATYPE_INT32, TS_DATATYPE_STRING, + TS_DATATYPE_DOUBLE, TS_DATATYPE_BLOB}; + for (uint32_t fail_after = 0; fail_after < 15; ++fail_after) { + SCOPED_TRACE(fail_after); + const int64_t before = + common::ModStat::get_instance().get_stat(common::MOD_TABLET); + ERRNO err = RET_OK; + common::TEST_fail_mem_alloc_after(common::MOD_TABLET, fail_after); + Tablet tablet = tablet_new(names, types, 4, 8, &err); + EXPECT_EQ(tablet, nullptr); + EXPECT_EQ(err, RET_OOM); + if (tablet) free_tablet(&tablet); + common::TEST_fail_mem_alloc_after(common::MOD_TABLET, fail_after); + tablet = tablet_new_with_target_name("dev", names, types, 4, 8); + EXPECT_EQ(tablet, nullptr); + if (tablet) free_tablet(&tablet); + EXPECT_EQ(common::ModStat::get_instance().get_stat(common::MOD_TABLET), + before); + } +} + TEST_F(CWrapperPublicWriterTest, PublicTableAndArrowNullArgs) { ERRNO err = RET_OK; EXPECT_EQ(tablet_new(nullptr, nullptr, 0, 1, &err), nullptr); @@ -186,7 +236,7 @@ TEST_F(CWrapperPublicWriterTest, PublicTableAndArrowNullArgs) { char* one_name[] = {first}; TSDataType one_type[] = {TS_DATATYPE_INT64}; EXPECT_EQ(tablet_new(one_name, one_type, 1, 1u << 30, &err), nullptr); - EXPECT_EQ(err, RET_INVALID_ARG); + EXPECT_EQ(err, RET_OVERFLOW); TSDataType unsupported_type[] = {TS_DATATYPE_VECTOR}; EXPECT_EQ(tablet_new(one_name, unsupported_type, 1, 1, &err), nullptr); diff --git a/go/tsfile/api_test.go b/go/tsfile/api_test.go index 14be89c64..4e208e7ad 100644 --- a/go/tsfile/api_test.go +++ b/go/tsfile/api_test.go @@ -22,6 +22,25 @@ import ( "testing" ) +func TestTabletRejectsOversizedCapacityForEveryType(t *testing.T) { + for _, dataType := range []DataType{ + DataTypeBoolean, DataTypeInt32, DataTypeDate, DataTypeFloat, + DataTypeInt64, DataTypeTimestamp, DataTypeDouble, + DataTypeText, DataTypeBlob, DataTypeString, + } { + for _, rows := range []int{1 << 29, (1 << 30) - 2, (1 << 30) - 1, 1 << 30} { + tablet, err := NewTablet([]TabletColumn{{Name: "value", DataType: dataType}}, rows) + if tablet != nil { + tablet.Close() + t.Fatalf("type %d rows %d: expected nil tablet", dataType, rows) + } + if !errors.Is(err, ErrOverflow) { + t.Fatalf("type %d rows %d: got %v, want ErrOverflow", dataType, rows, err) + } + } + } +} + func TestTabletValidationAndClose(t *testing.T) { if _, err := NewTablet(nil, 1); !errors.Is(err, ErrInvalidArgument) { t.Fatalf("empty columns: %v", err) diff --git a/go/tsfile/cgo_bridge.go b/go/tsfile/cgo_bridge.go index 5208c7219..d03b99a8d 100644 --- a/go/tsfile/cgo_bridge.go +++ b/go/tsfile/cgo_bridge.go @@ -363,12 +363,10 @@ func releaseTablet(ptr unsafe.Pointer) cerrno { return C.RET_OK } -// cTabletMaxRowsMax is the native row-capacity limit: storage::Tablet -// asserts max_rows > 0 && max_rows < (1 << 30), so the bridge only admits -// 1 <= maxRows < cTabletMaxRowsMax. Beyond the native assertion, larger -// capacities would make the native row-buffer index math (int32 offsets of -// length max_rows + 1) unsafe. -const cTabletMaxRowsMax = 1 << 30 +// cTabletMaxRowsMax is the exclusive native row-capacity limit. Every Tablet +// allocates int64 timestamps through a uint32-sized allocator, so capacities +// at or above 1<<29 would truncate the allocation's byte size. +const cTabletMaxRowsMax = 1 << 29 // validateTabletMaxRows is a pure integer validator for the tablet row // capacity: it reports whether maxRows fits the native domain diff --git a/go/tsfile/cgo_bridge_test.go b/go/tsfile/cgo_bridge_test.go index 24cf32488..4bda750e9 100644 --- a/go/tsfile/cgo_bridge_test.go +++ b/go/tsfile/cgo_bridge_test.go @@ -380,9 +380,9 @@ func TestBridgeValidateCBytesLen(t *testing.T) { } // TestBridgeValidateTabletMaxRows pins the pure tablet row-capacity -// validator against the native domain [1, 1<<30). Every case is a bare +// validator against the native domain [1, 1<<29). Every case is a bare // integer: no tablet and no C allocation is involved, so the boundary at -// 1<<30 is tested without reserving row buffers. +// 1<<29 is tested without reserving row buffers. func TestBridgeValidateTabletMaxRows(t *testing.T) { for _, n := range []int{1, 8, 1024, cTabletMaxRowsMax - 1} { if err := validateTabletMaxRows("op", n); err != nil { @@ -423,7 +423,7 @@ func TestBridgeValidateCInt32(t *testing.T) { // host) must be unreachable. No handle is returned and no native tablet is // created, so no row buffers are allocated. func TestBridgeTabletRejectsOversizedMaxRows(t *testing.T) { - for _, maxRows := range []int{cTabletMaxRowsMax, 1 << 31, 1 << 32, math.MaxInt32, math.MaxInt64} { + for _, maxRows := range []int{1 << 29, (1 << 30) - 2, (1 << 30) - 1, 1 << 30, 1 << 31, 1 << 32, math.MaxInt32, math.MaxInt64} { h, err := newTabletHandle([]string{"s1"}, []DataType{DataTypeInt64}, maxRows) if h != nil { t.Fatalf("maxRows %d: expected no handle", maxRows) diff --git a/go/tsfile/tablet.go b/go/tsfile/tablet.go index 37c7c5493..6a52c80cc 100644 --- a/go/tsfile/tablet.go +++ b/go/tsfile/tablet.go @@ -32,6 +32,8 @@ type Tablet struct { } // NewTablet allocates a tablet with fixed columns and row capacity. +// maxRows must satisfy 1 <= maxRows < 1<<29. Larger capacities return ErrOverflow; +// capacities within this range may still fail if memory is unavailable. func NewTablet(columns []TabletColumn, maxRows int) (*Tablet, error) { if len(columns) == 0 { return nil, fmt.Errorf("%w: at least one column is required", ErrInvalidArgument) From 41340a82a31e9091083586d119246723c04050e1 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Sat, 10 Oct 2026 12:12:31 +0800 Subject: [PATCH 2/2] fix(cpp): allocate tablet buffers at full size and report OOM --- cpp/src/common/allocator/alloc_base.h | 8 ++-- cpp/src/common/allocator/mem_alloc.cc | 25 ++++++++-- cpp/src/common/container/bit_map.cc | 2 +- cpp/src/common/container/bit_map.h | 2 +- cpp/src/common/tablet.cc | 25 +++++----- cpp/src/common/tablet.h | 27 +++++------ cpp/src/cwrapper/tsfile_cwrapper.cc | 12 ++--- cpp/src/cwrapper/tsfile_cwrapper.h | 7 ++- cpp/src/writer/tsfile_writer.cc | 16 +++---- cpp/src/writer/tsfile_writer.h | 2 +- cpp/test/common/allocator/alloc_base_test.cc | 41 ++++++++++++++++- cpp/test/common/container/bit_map_test.cc | 8 ++++ cpp/test/common/tablet_test.cc | 46 +++++++++++++++---- .../cwrapper/cwrapper_public_writer_test.cc | 14 ++++-- go/tsfile/api_test.go | 2 +- go/tsfile/cgo_bridge.go | 6 +-- go/tsfile/cgo_bridge_test.go | 12 ++--- go/tsfile/tablet.go | 4 +- 18 files changed, 173 insertions(+), 86 deletions(-) diff --git a/cpp/src/common/allocator/alloc_base.h b/cpp/src/common/allocator/alloc_base.h index 82bc734ea..a29ad42ea 100644 --- a/cpp/src/common/allocator/alloc_base.h +++ b/cpp/src/common/allocator/alloc_base.h @@ -62,15 +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 { @@ -131,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); } }; diff --git a/cpp/src/common/allocator/mem_alloc.cc b/cpp/src/common/allocator/mem_alloc.cc index 9d9a5222c..34ad5409c 100644 --- a/cpp/src/common/allocator/mem_alloc.cc +++ b/cpp/src/common/allocator/mem_alloc.cc @@ -25,6 +25,7 @@ #include #include #include +#include #include "alloc_base.h" #include "common/logger/elog.h" @@ -39,9 +40,11 @@ namespace { std::atomic 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); } @@ -49,7 +52,10 @@ 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] = { @@ -97,11 +103,12 @@ 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; @@ -112,6 +119,11 @@ void* mem_alloc(uint32_t size, AllocModID mid) { 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::max() - header_size) { + return nullptr; + } const size_t total_size = size + header_size; auto raw = static_cast(malloc(total_size)); if (UNLIKELY(raw == nullptr)) { @@ -161,24 +173,29 @@ void mem_free(void* ptr) { static_cast(*reinterpret_cast(raw_ptr + 4)) | (static_cast(*reinterpret_cast(raw_ptr)) << 32); auto mid = static_cast(header & 0x7F); - auto size = static_cast(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::max() - ALIGNMENT) { + return nullptr; + } char* p = static_cast(ptr); char* raw_ptr = p - ALIGNMENT; const uint64_t header = static_cast(*reinterpret_cast(raw_ptr + 4)) | (static_cast(*reinterpret_cast(raw_ptr)) << 32); auto mid = static_cast(header & 0x7F); - auto original_size = static_cast(header >> 8); + const uint64_t original_size = header >> 8; p = static_cast(realloc(raw_ptr, size + ALIGNMENT)); if (UNLIKELY(p == nullptr)) { return nullptr; diff --git a/cpp/src/common/container/bit_map.cc b/cpp/src/common/container/bit_map.cc index f786226f0..a6bb8c3e5 100644 --- a/cpp/src/common/container/bit_map.cc +++ b/cpp/src/common/container/bit_map.cc @@ -32,7 +32,7 @@ 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(mem_alloc(size, mod_id)); if (bitmap_ == nullptr) return common::E_OOM; // need set to 0, otherwise there will be wrong data diff --git a/cpp/src/common/container/bit_map.h b/cpp/src/common/container/bit_map.h index 90ed0e0b6..678147f42 100644 --- a/cpp/src/common/container/bit_map.h +++ b/cpp/src/common/container/bit_map.h @@ -166,7 +166,7 @@ class BitMap { if (byte_val) { return from + bitops::ctz_nonzero(static_cast(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 = diff --git a/cpp/src/common/tablet.cc b/cpp/src/common/tablet.cc index 74627ee50..1fbf3e6e5 100644 --- a/cpp/src/common/tablet.cc +++ b/cpp/src/common/tablet.cc @@ -35,14 +35,13 @@ int Tablet::init() { ASSERT(timestamps_ == nullptr); cur_row_size_ = 0; if (max_row_num_ == 0) return E_INVALID_ARG; - // Check before any allocation: otherwise mem_alloc truncates a size_t - // request to uint32_t while memset / memcpy still use the full size. - if (max_row_num_ > MAX_ROWS) return E_OVERFLOW; + if (static_cast(max_row_num_) > SIZE_MAX / sizeof(int64_t)) + return E_OOM; size_t schema_count = schema_vec_->size(); - if (schema_count > UINT32_MAX / sizeof(ValueMatrixEntry) || - schema_count > UINT32_MAX / sizeof(BitMap)) { - return E_OVERFLOW; + if (schema_count > SIZE_MAX / sizeof(ValueMatrixEntry) || + schema_count > SIZE_MAX / sizeof(BitMap)) { + return E_OOM; } std::pair::iterator, bool> ins_res; for (size_t c = 0; c < schema_count; c++) { @@ -272,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(bitmap), bm_bytes); } @@ -328,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(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(bitmap), bm_bytes); } @@ -402,9 +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; } @@ -620,7 +621,7 @@ std::vector Tablet::find_all_device_boundaries() const { if (row_count <= 1) return {}; // Use uint64_t bitmap instead of vector 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 boundary(nwords, 0); // Walk id columns RIGHT to LEFT. In time-series tag systems the rightmost @@ -691,7 +692,7 @@ std::vector Tablet::find_all_device_boundaries() const { return result; } -std::shared_ptr Tablet::get_device_id(int i) const { +std::shared_ptr Tablet::get_device_id(uint32_t i) const { std::vector id_array; id_array.push_back(new std::string(insert_target_name_)); for (auto id_column_idx : id_column_indexes_) { diff --git a/cpp/src/common/tablet.h b/cpp/src/common/tablet.h index 1bbf1349d..8b5b87ad4 100644 --- a/cpp/src/common/tablet.h +++ b/cpp/src/common/tablet.h @@ -49,10 +49,10 @@ class Tablet { // Arrow-style string column: offsets + contiguous buffer. // string[i] = buffer + offsets[i], len = offsets[i+1] - offsets[i] struct StringColumn { - int32_t* offsets; // length: max_rows + 1 (Arrow-compatible) - char* buffer; // contiguous string data - uint32_t buf_capacity; // allocated buffer size - uint32_t buf_used; // bytes written so far + int32_t* offsets; // length: max_rows + 1 (Arrow-compatible) + char* buffer; // contiguous string data + size_t buf_capacity; // allocated buffer size + uint32_t buf_used; // bytes written so far StringColumn() : offsets(nullptr), buffer(nullptr), buf_capacity(0), buf_used(0) {} @@ -60,15 +60,15 @@ class Tablet { int init(uint32_t max_rows) { const uint64_t offsets_bytes = sizeof(int32_t) * (static_cast(max_rows) + 1); - if (offsets_bytes > UINT32_MAX) return common::E_OVERFLOW; + const uint64_t data_bytes = static_cast(max_rows) * 32; + if (offsets_bytes > SIZE_MAX || data_bytes > SIZE_MAX) + return common::E_OOM; offsets = (int32_t*)common::mem_alloc( - static_cast(offsets_bytes), common::MOD_TABLET); + static_cast(offsets_bytes), common::MOD_TABLET); if (offsets == nullptr) return common::E_OOM; offsets[0] = 0; - // Estimate 32 bytes per row, but cap eager allocation at 1 MiB. - // Large tablets grow the data buffer as values are appended. - buf_capacity = static_cast(std::min( - static_cast(max_rows) * 32, 1u << 20)); + // Estimate 32 bytes per row and preserve the full preallocation. + buf_capacity = static_cast(data_bytes); buffer = (char*)common::mem_alloc(buf_capacity, common::MOD_TABLET); if (buffer == nullptr) return common::E_OOM; buf_used = 0; @@ -146,9 +146,6 @@ class Tablet { public: static const uint32_t DEFAULT_MAX_ROWS = 1024; - // Every tablet has an int64 timestamp buffer allocated through mem_alloc, - // whose byte-size argument is uint32_t. - static const uint32_t MAX_ROWS = UINT32_MAX / sizeof(int64_t); int err_code_ = common::E_OK; public: @@ -337,7 +334,7 @@ class Tablet { // its backing buffers. row_count is typically 0 before refilling. void reset(uint32_t row_count = 0); - void* get_value(int row_index, uint32_t schema_index, + void* get_value(uint32_t row_index, uint32_t schema_index, common::TSDataType& data_type) const; /** * @brief Template function to add a value of type T to the specified row @@ -356,7 +353,7 @@ class Tablet { void set_column_categories( const std::vector& column_categories); - std::shared_ptr get_device_id(int i) const; + std::shared_ptr get_device_id(uint32_t i) const; std::vector find_all_device_boundaries() const; // When the caller guarantees that all rows belong to a single device, diff --git a/cpp/src/cwrapper/tsfile_cwrapper.cc b/cpp/src/cwrapper/tsfile_cwrapper.cc index c572ec2ae..425baf378 100644 --- a/cpp/src/cwrapper/tsfile_cwrapper.cc +++ b/cpp/src/cwrapper/tsfile_cwrapper.cc @@ -371,10 +371,6 @@ Tablet tablet_new(char** column_name_list, TSDataType* data_types, data_types == nullptr) { return nullptr; } - if (max_rows > storage::Tablet::MAX_ROWS) { - *err_code = common::E_OVERFLOW; - return nullptr; - } try { std::vector measurement_list; std::vector data_type_list; @@ -2506,11 +2502,9 @@ Tablet tablet_new_with_target_name(const char* target_name, int max_rows) { // Public C entry point: reject arguments that would make the private // delegate read out-of-bounds memory (negative column_num, null lists or - // null column names) before touching caller buffers. Timestamp allocation - // must also fit the native allocator's uint32_t byte-size argument. - // Returns NULL on invalid arguments, overflow or allocation failure. - if (column_num < 0 || max_rows <= 0 || - static_cast(max_rows) > storage::Tablet::MAX_ROWS) { + // null column names) before touching caller buffers. + // Returns NULL on invalid arguments or allocation failure. + if (column_num < 0 || max_rows <= 0) { return nullptr; } if (column_num > 0) { diff --git a/cpp/src/cwrapper/tsfile_cwrapper.h b/cpp/src/cwrapper/tsfile_cwrapper.h index 1065e53ee..ebbd6930e 100644 --- a/cpp/src/cwrapper/tsfile_cwrapper.h +++ b/cpp/src/cwrapper/tsfile_cwrapper.h @@ -582,8 +582,8 @@ void tsfile_free_tsfile_properties(TsFileProperty* properties, uint32_t length); * @param column_name_list [in] Column names array. Size=column_num. * @param data_types [in] Data types array. Size=column_num. * @param column_num [in] Number of columns. Must be ≥1. - * @param max_rows [in] Pre-allocated row capacity. Must satisfy - * 1 <= max_rows < 2^29; larger capacities return RET_OVERFLOW. + * @param max_rows [in] Pre-allocated row capacity. Must be >=1. + * Allocation failures return RET_OOM. * @param err_code [out] Receives RET_OK on success or the failure code. Must * not be NULL. * @return Tablet Valid handle, or NULL on failure. @@ -1291,8 +1291,7 @@ TsFileGenericWriter tsfile_generic_writer_new(const char* pathname, * @param data_types data types of the columns (borrowed); must not be NULL * when column_num > 0 * @param column_num number of columns; must be >= 0 - * @param max_rows maximum number of rows the tablet can hold; must satisfy - * 0 < max_rows < 2^29 + * @param max_rows maximum number of rows the tablet can hold; must be > 0 * @return a caller-owned tablet, or NULL if the arguments are invalid or * memory allocation fails */ diff --git a/cpp/src/writer/tsfile_writer.cc b/cpp/src/writer/tsfile_writer.cc index ed3e09be7..e7539f7fb 100644 --- a/cpp/src/writer/tsfile_writer.cc +++ b/cpp/src/writer/tsfile_writer.cc @@ -1325,10 +1325,10 @@ int TsFileWriter::write_table(Tablet& tablet) { std::vector device_ctxs; std::map, size_t, IDeviceIDComparator> device_ctx_index; - int start_idx = 0; + uint32_t start_idx = 0; for (auto& pair : device_id_end_index_pairs) { auto device_id = pair.first; - int end_idx = pair.second; + uint32_t end_idx = pair.second; if (end_idx == 0) continue; const uint32_t si = static_cast(start_idx); @@ -1558,10 +1558,10 @@ int TsFileWriter::write_table(Tablet& tablet) { } } } else { - int start_idx = 0; + uint32_t start_idx = 0; for (auto& device_id_end_index_pair : device_id_end_index_pairs) { auto device_id = device_id_end_index_pair.first; - int end_idx = device_id_end_index_pair.second; + uint32_t end_idx = device_id_end_index_pair.second; if (end_idx == 0) continue; const uint32_t si = static_cast(start_idx); @@ -1636,9 +1636,9 @@ int TsFileWriter::write_table(Tablet& tablet) { // per-device last_time_ floor to the highest timestamp this tablet // contributed for each device. if (enforce_recovered_last_time_order_ && tablet.timestamps_) { - int update_start = 0; + uint32_t update_start = 0; for (auto& pair : device_id_end_index_pairs) { - int end_idx = pair.second; + uint32_t end_idx = pair.second; if (end_idx == 0) continue; if (end_idx > update_start) { auto schema_it = schemas_.find(pair.first); @@ -1660,9 +1660,9 @@ int TsFileWriter::write_table(Tablet& tablet) { return ret; } -std::vector, int>> +std::vector, uint32_t>> TsFileWriter::split_tablet_by_device(const Tablet& tablet) { - std::vector, int>> result; + std::vector, uint32_t>> result; if (tablet.id_column_indexes_.empty() || tablet.single_device_) { // No tag columns or caller guarantees single device — skip boundary diff --git a/cpp/src/writer/tsfile_writer.h b/cpp/src/writer/tsfile_writer.h index 0606a7a33..ee4c3baf6 100644 --- a/cpp/src/writer/tsfile_writer.h +++ b/cpp/src/writer/tsfile_writer.h @@ -194,7 +194,7 @@ class TsFileWriter { uint32_t end_idx = UINT32_MAX); int register_timeseries(const std::string& device_path, MeasurementSchema* measurement_schema); - std::vector, int>> + std::vector, uint32_t>> split_tablet_by_device(const Tablet& tablet); private: diff --git a/cpp/test/common/allocator/alloc_base_test.cc b/cpp/test/common/allocator/alloc_base_test.cc index 6946aa8c2..2f968b279 100644 --- a/cpp/test/common/allocator/alloc_base_test.cc +++ b/cpp/test/common/allocator/alloc_base_test.cc @@ -20,6 +20,8 @@ #include +#include + namespace common { const uint32_t HEADER_SIZE_4B = 4; @@ -61,6 +63,43 @@ TEST(AllocBaseTest, AllocLargeSize) { mem_free(ptr); } +TEST(AllocBaseTest, SizeAboveUint32ReachesAllocatorWithoutTruncation) { + if (sizeof(size_t) < sizeof(uint64_t)) GTEST_SKIP(); + const size_t size = static_cast(UINT32_MAX) + 1; + TEST_fail_mem_alloc_after(MOD_DEFAULT, 0); + BaseAllocator allocator; + EXPECT_EQ(allocator.alloc(size, MOD_DEFAULT), nullptr); + EXPECT_EQ(TEST_get_failed_mem_alloc_size(), size); + + char* ptr = static_cast(allocator.alloc(8, MOD_DEFAULT)); + ASSERT_NE(ptr, nullptr); + memset(ptr, 'x', 8); + TEST_fail_next_mem_realloc(); + EXPECT_EQ(mem_realloc(ptr, size), nullptr); + EXPECT_EQ(TEST_get_failed_mem_alloc_size(), size); + for (int i = 0; i < 8; ++i) EXPECT_EQ(ptr[i], 'x'); + allocator.free(ptr); +} + +TEST(AllocBaseTest, UnrepresentableSizeReturnsNullAndPreservesReallocInput) { + const size_t size = std::numeric_limits::max(); + void* oversized = mem_alloc(size, MOD_DEFAULT); + EXPECT_EQ(oversized, nullptr); + if (oversized) mem_free(oversized); + + char* ptr = static_cast(mem_alloc(8, MOD_DEFAULT)); + ASSERT_NE(ptr, nullptr); + memset(ptr, 'x', 8); + void* resized = mem_realloc(ptr, size); + EXPECT_EQ(resized, nullptr); + if (resized) { + mem_free(resized); + } else { + for (int i = 0; i < 8; ++i) EXPECT_EQ(ptr[i], 'x'); + mem_free(ptr); + } +} + TEST(AllocBaseTest, ReallocateToLargerSize) { const uint32_t initial_size = 1000; const uint32_t new_size = 2000; @@ -93,4 +132,4 @@ TEST(AllocBaseTest, ReallocateToSmallerSize) { mem_free(ptr); } -} // namespace common \ No newline at end of file +} // namespace common diff --git a/cpp/test/common/container/bit_map_test.cc b/cpp/test/common/container/bit_map_test.cc index 9c70e661c..20c02ab3b 100644 --- a/cpp/test/common/container/bit_map_test.cc +++ b/cpp/test/common/container/bit_map_test.cc @@ -33,6 +33,14 @@ TEST_F(BitMapTest, Initialization) { EXPECT_TRUE(bitmap.get_bitmap() != nullptr); } +TEST_F(BitMapTest, Uint32MaxCapacityDoesNotWrapByteCount) { + common::TEST_fail_mem_alloc_after(common::MOD_TSBLOCK, 0); + common::BitMap bitmap; + EXPECT_EQ(bitmap.init(UINT32_MAX), common::E_OOM); + EXPECT_EQ(common::TEST_get_failed_mem_alloc_size(), size_t(1) << 29); + EXPECT_EQ(bitmap.get_size(), 0u); +} + TEST_F(BitMapTest, Reset) { common::BitMap bitmap; bitmap.init(100); diff --git a/cpp/test/common/tablet_test.cc b/cpp/test/common/tablet_test.cc index bd14a66f3..f33471e7a 100644 --- a/cpp/test/common/tablet_test.cc +++ b/cpp/test/common/tablet_test.cc @@ -24,9 +24,10 @@ namespace storage { -// Keep the old implementation safe to test: if it attempts to allocate before -// rejecting the capacity, fail that allocation instead of reserving GiBs. -TEST(TabletTest, OversizedCapacityRejectedBeforeAllocation) { +// Fail the first allocation so large capacities can be checked without +// reserving GiBs. The allocator must receive the full timestamp byte count. +TEST(TabletTest, LargeCapacityAttemptsFullAllocationAndReturnsOom) { + if (sizeof(size_t) < sizeof(uint64_t)) GTEST_SKIP(); const std::vector types = { common::BOOLEAN, common::INT32, common::DATE, common::FLOAT, common::INT64, common::TIMESTAMP, common::DOUBLE, common::TEXT, @@ -40,12 +41,14 @@ TEST(TabletTest, OversizedCapacityRejectedBeforeAllocation) { << " rows=" << capacity); common::TEST_fail_mem_alloc_after(common::MOD_TABLET, 0); Tablet tablet({"value"}, {type}, capacity); - ASSERT_EQ(tablet.err_code_, common::E_OVERFLOW); + ASSERT_EQ(tablet.err_code_, common::E_OOM); + EXPECT_EQ(common::TEST_get_failed_mem_alloc_size(), + static_cast(capacity) * sizeof(int64_t)); EXPECT_EQ(tablet.get_cur_row_size(), 0u); - EXPECT_EQ(tablet.add_timestamp(0, 1), common::E_OVERFLOW); - // The overflow check must leave the failure injection unconsumed. + EXPECT_EQ(tablet.add_timestamp(0, 1), common::E_OOM); + // The allocation failure is one-shot. Tablet next({"value"}, {type}, 1u); - EXPECT_EQ(next.err_code_, common::E_OOM); + EXPECT_EQ(next.err_code_, common::E_OK); } } } @@ -56,13 +59,20 @@ TEST(TabletTest, CapacityCheckedByEveryConstructor) { auto schema = std::make_shared>(); schema->emplace_back("value", common::DOUBLE, common::PLAIN, common::UNCOMPRESSED); - for (int capacity : {0, -1, 1 << 29, (1 << 30) - 1}) { + for (int capacity : {0, -1, 1 << 29, (1 << 30) - 1, INT32_MAX}) { + if (capacity > 0 && sizeof(size_t) < sizeof(uint64_t)) continue; const int expected = - capacity <= 0 ? common::E_INVALID_ARG : common::E_OVERFLOW; + capacity <= 0 ? common::E_INVALID_ARG : common::E_OOM; + if (capacity > 0) + common::TEST_fail_mem_alloc_after(common::MOD_TABLET, 0); Tablet from_schema("dev", schema, capacity); EXPECT_EQ(from_schema.err_code_, expected); + if (capacity > 0) + common::TEST_fail_mem_alloc_after(common::MOD_TABLET, 0); Tablet from_lists("dev", &names, &types, capacity); EXPECT_EQ(from_lists.err_code_, expected); + if (capacity > 0) + common::TEST_fail_mem_alloc_after(common::MOD_TABLET, 0); Tablet with_categories("table", names, types, {common::ColumnCategory::FIELD}, capacity); EXPECT_EQ(with_categories.err_code_, expected); @@ -71,6 +81,24 @@ TEST(TabletTest, CapacityCheckedByEveryConstructor) { EXPECT_EQ(zero.err_code_, common::E_INVALID_ARG); } +TEST(TabletTest, StringColumnsPreallocate32BytesPerRow) { + const uint32_t rows = 1u << 16; + for (auto type : {common::TEXT, common::BLOB, common::STRING}) { + const int64_t before = + common::ModStat::get_instance().get_stat(common::MOD_TABLET); + // timestamps, matrix, StringColumn, offsets, then the data buffer. + common::TEST_fail_mem_alloc_after(common::MOD_TABLET, 4); + { + Tablet tablet({"value"}, {type}, rows); + EXPECT_EQ(tablet.err_code_, common::E_OOM); + EXPECT_EQ(common::TEST_get_failed_mem_alloc_size(), + static_cast(rows) * 32); + } + EXPECT_EQ(common::ModStat::get_instance().get_stat(common::MOD_TABLET), + before); + } +} + TEST(TabletTest, AllocationFailuresReturnOomAndReleasePartialBuffers) { // timestamps, matrix, INT32, STRING object/offsets/data, DOUBLE, // BLOB object/offsets/data, bitmap array, and four bitmap buffers. diff --git a/cpp/test/cwrapper/cwrapper_public_writer_test.cc b/cpp/test/cwrapper/cwrapper_public_writer_test.cc index 5b08d0e8c..269985572 100644 --- a/cpp/test/cwrapper/cwrapper_public_writer_test.cc +++ b/cpp/test/cwrapper/cwrapper_public_writer_test.cc @@ -167,7 +167,8 @@ TEST_F(CWrapperPublicWriterTest, TableWriterRejectsInvalidSchema) { remove(filename.c_str()); } -TEST_F(CWrapperPublicWriterTest, TabletOversizedCapacityReturnsOverflow) { +TEST_F(CWrapperPublicWriterTest, TabletLargeCapacityAttemptsAllocation) { + if (sizeof(size_t) < sizeof(uint64_t)) GTEST_SKIP(); char name[] = "value"; char* names[] = {name}; const TSDataType types[] = {TS_DATATYPE_BOOLEAN, TS_DATATYPE_INT32, @@ -181,11 +182,17 @@ TEST_F(CWrapperPublicWriterTest, TabletOversizedCapacityReturnsOverflow) { SCOPED_TRACE(::testing::Message() << "type=" << type << " rows=" << rows); ERRNO err = RET_OK; + common::TEST_fail_mem_alloc_after(common::MOD_TABLET, 0); Tablet tablet = tablet_new(names, &type, 1, rows, &err); EXPECT_EQ(tablet, nullptr); - EXPECT_EQ(err, RET_OVERFLOW); + EXPECT_EQ(err, RET_OOM); + if (sizeof(size_t) >= sizeof(uint64_t)) { + EXPECT_EQ(common::TEST_get_failed_mem_alloc_size(), + static_cast(rows) * sizeof(int64_t)); + } if (tablet) free_tablet(&tablet); } + common::TEST_fail_mem_alloc_after(common::MOD_TABLET, 0); EXPECT_EQ(tablet_new_with_target_name("dev", names, &type, 1, 1 << 29), nullptr); } @@ -235,9 +242,6 @@ TEST_F(CWrapperPublicWriterTest, PublicTableAndArrowNullArgs) { char* one_name[] = {first}; TSDataType one_type[] = {TS_DATATYPE_INT64}; - EXPECT_EQ(tablet_new(one_name, one_type, 1, 1u << 30, &err), nullptr); - EXPECT_EQ(err, RET_OVERFLOW); - TSDataType unsupported_type[] = {TS_DATATYPE_VECTOR}; EXPECT_EQ(tablet_new(one_name, unsupported_type, 1, 1, &err), nullptr); EXPECT_EQ(err, RET_TYPE_NOT_SUPPORTED); diff --git a/go/tsfile/api_test.go b/go/tsfile/api_test.go index 4e208e7ad..368863393 100644 --- a/go/tsfile/api_test.go +++ b/go/tsfile/api_test.go @@ -28,7 +28,7 @@ func TestTabletRejectsOversizedCapacityForEveryType(t *testing.T) { DataTypeInt64, DataTypeTimestamp, DataTypeDouble, DataTypeText, DataTypeBlob, DataTypeString, } { - for _, rows := range []int{1 << 29, (1 << 30) - 2, (1 << 30) - 1, 1 << 30} { + for _, rows := range []int{1 << 32, (1 << 32) + 1} { tablet, err := NewTablet([]TabletColumn{{Name: "value", DataType: dataType}}, rows) if tablet != nil { tablet.Close() diff --git a/go/tsfile/cgo_bridge.go b/go/tsfile/cgo_bridge.go index d03b99a8d..e8e673fbe 100644 --- a/go/tsfile/cgo_bridge.go +++ b/go/tsfile/cgo_bridge.go @@ -363,10 +363,8 @@ func releaseTablet(ptr unsafe.Pointer) cerrno { return C.RET_OK } -// cTabletMaxRowsMax is the exclusive native row-capacity limit. Every Tablet -// allocates int64 timestamps through a uint32-sized allocator, so capacities -// at or above 1<<29 would truncate the allocation's byte size. -const cTabletMaxRowsMax = 1 << 29 +// cTabletMaxRowsMax is the exclusive limit of the C ABI's uint32_t row count. +const cTabletMaxRowsMax = 1 << 32 // validateTabletMaxRows is a pure integer validator for the tablet row // capacity: it reports whether maxRows fits the native domain diff --git a/go/tsfile/cgo_bridge_test.go b/go/tsfile/cgo_bridge_test.go index 4bda750e9..f469daf22 100644 --- a/go/tsfile/cgo_bridge_test.go +++ b/go/tsfile/cgo_bridge_test.go @@ -380,11 +380,11 @@ func TestBridgeValidateCBytesLen(t *testing.T) { } // TestBridgeValidateTabletMaxRows pins the pure tablet row-capacity -// validator against the native domain [1, 1<<29). Every case is a bare +// validator against the native domain [1, 1<<32). Every case is a bare // integer: no tablet and no C allocation is involved, so the boundary at -// 1<<29 is tested without reserving row buffers. +// 1<<32 is tested without reserving row buffers. func TestBridgeValidateTabletMaxRows(t *testing.T) { - for _, n := range []int{1, 8, 1024, cTabletMaxRowsMax - 1} { + for _, n := range []int{1, 8, 1024, 1 << 29, (1 << 30) - 1, math.MaxInt32, cTabletMaxRowsMax - 1} { if err := validateTabletMaxRows("op", n); err != nil { t.Fatalf("validateTabletMaxRows(%d) = %v, want nil", n, err) } @@ -394,7 +394,7 @@ func TestBridgeValidateTabletMaxRows(t *testing.T) { t.Fatalf("validateTabletMaxRows(%d) = %v, want ErrInvalidArgument", n, err) } } - for _, n := range []int{cTabletMaxRowsMax, cTabletMaxRowsMax + 1, math.MaxInt32, math.MaxInt64} { + for _, n := range []int{cTabletMaxRowsMax, cTabletMaxRowsMax + 1, math.MaxInt64} { if err := validateTabletMaxRows("op", n); !errors.Is(err, ErrOverflow) { t.Fatalf("validateTabletMaxRows(%d) = %v, want ErrOverflow", n, err) } @@ -419,11 +419,11 @@ func TestBridgeValidateCInt32(t *testing.T) { // TestBridgeTabletRejectsOversizedMaxRows proves newTabletHandle rejects // row capacities outside the native domain before any C call: the truncated -// C.int path (values >= 1<<32 wrap to small or negative ints on a 64-bit +// C.uint32_t path (values >= 1<<32 wrap to small integers on a 64-bit // host) must be unreachable. No handle is returned and no native tablet is // created, so no row buffers are allocated. func TestBridgeTabletRejectsOversizedMaxRows(t *testing.T) { - for _, maxRows := range []int{1 << 29, (1 << 30) - 2, (1 << 30) - 1, 1 << 30, 1 << 31, 1 << 32, math.MaxInt32, math.MaxInt64} { + for _, maxRows := range []int{1 << 32, (1 << 32) + 1, math.MaxInt64} { h, err := newTabletHandle([]string{"s1"}, []DataType{DataTypeInt64}, maxRows) if h != nil { t.Fatalf("maxRows %d: expected no handle", maxRows) diff --git a/go/tsfile/tablet.go b/go/tsfile/tablet.go index 6a52c80cc..38bc8055c 100644 --- a/go/tsfile/tablet.go +++ b/go/tsfile/tablet.go @@ -32,8 +32,8 @@ type Tablet struct { } // NewTablet allocates a tablet with fixed columns and row capacity. -// maxRows must satisfy 1 <= maxRows < 1<<29. Larger capacities return ErrOverflow; -// capacities within this range may still fail if memory is unavailable. +// maxRows must fit the C ABI's positive uint32 row count. Allocation failures +// return an out-of-memory error. func NewTablet(columns []TabletColumn, maxRows int) (*Tablet, error) { if len(columns) == 0 { return nil, fmt.Errorf("%w: at least one column is required", ErrInvalidArgument)