diff --git a/cpp/src/common/allocator/alloc_base.h b/cpp/src/common/allocator/alloc_base.h index 2b97c8f83..a29ad42ea 100644 --- a/cpp/src/common/allocator/alloc_base.h +++ b/cpp/src/common/allocator/alloc_base.h @@ -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 { @@ -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); } }; diff --git a/cpp/src/common/allocator/mem_alloc.cc b/cpp/src/common/allocator/mem_alloc.cc index deda0b8d8..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" @@ -37,11 +38,24 @@ 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; +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] = { @@ -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::max() - header_size) { + return nullptr; + } const size_t total_size = size + header_size; auto raw = static_cast(malloc(total_size)); if (UNLIKELY(raw == nullptr)) { @@ -144,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 3b1af6ab2..a6bb8c3e5 100644 --- a/cpp/src/common/container/bit_map.cc +++ b/cpp/src/common/container/bit_map.cc @@ -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(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/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 3bfc54bb9..1fbf3e6e5 100644 --- a/cpp/src/common/tablet.cc +++ b/cpp/src/common/tablet.cc @@ -33,12 +33,16 @@ 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; + if (static_cast(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::iterator, bool> ins_res; for (size_t c = 0; c < schema_count; c++) { ins_res = schema_map_.insert( @@ -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( + 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 +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: @@ -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; @@ -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; @@ -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(bitmap), bm_bytes); } @@ -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(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); } @@ -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(); @@ -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; } @@ -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); @@ -595,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 @@ -666,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 095157ac7..8b5b87ad4 100644 --- a/cpp/src/common/tablet.h +++ b/cpp/src/common/tablet.h @@ -49,21 +49,30 @@ 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) {} - 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); + 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( - 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 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; + 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); @@ -141,7 +152,7 @@ class Tablet { 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 +160,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 +167,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 +175,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 +195,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), @@ -333,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 @@ -352,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, @@ -407,7 +408,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..425baf378 100644 --- a/cpp/src/cwrapper/tsfile_cwrapper.cc +++ b/cpp/src/cwrapper/tsfile_cwrapper.cc @@ -368,7 +368,7 @@ 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; } try { @@ -2329,12 +2329,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 +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. 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. + // 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 4a87db90c..ebbd6930e 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 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. @@ -1290,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^30 + * @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 ce75ca2ee..f33471e7a 100644 --- a/cpp/test/common/tablet_test.cc +++ b/cpp/test/common/tablet_test.cc @@ -24,6 +24,159 @@ namespace storage { +// 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, + 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_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_OOM); + // The allocation failure is one-shot. + Tablet next({"value"}, {type}, 1u); + EXPECT_EQ(next.err_code_, common::E_OK); + } + } +} + +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, INT32_MAX}) { + if (capacity > 0 && sizeof(size_t) < sizeof(uint64_t)) continue; + const int expected = + 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); + } + Tablet zero(names, types, 0u); + 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. + 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..269985572 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,61 @@ TEST_F(CWrapperPublicWriterTest, TableWriterRejectsInvalidSchema) { remove(filename.c_str()); } +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, + 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; + 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_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); + } +} + +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); @@ -185,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_INVALID_ARG); - 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 14be89c64..368863393 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 << 32, (1 << 32) + 1} { + 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..e8e673fbe 100644 --- a/go/tsfile/cgo_bridge.go +++ b/go/tsfile/cgo_bridge.go @@ -363,12 +363,8 @@ 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 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 24cf32488..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<<30). 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<<30 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{cTabletMaxRowsMax, 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 37c7c5493..38bc8055c 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 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)