From 8ad7c72cc8732172ad427b06da6eaa0e7568fda4 Mon Sep 17 00:00:00 2001 From: zhongzheyun Date: Thu, 27 Aug 2026 17:21:26 +0800 Subject: [PATCH 1/3] feat(build): make TBB dependency configurable --- CMakeLists.txt | 10 +- cmake_modules/DefineOptions.cmake | 2 + cmake_modules/ThirdpartyToolchain.cmake | 4 +- docs/source/building.rst | 17 ++ src/paimon/CMakeLists.txt | 5 +- .../common/utils/concurrent_backend_factory.h | 76 ++++++++ .../common/utils/concurrent_bounded_queue.h | 166 ++++++++++++++++++ .../utils/concurrent_bounded_queue_test.cpp | 130 ++++++++++++++ src/paimon/common/utils/concurrent_hash_map.h | 131 +++++++++++++- .../common/utils/concurrent_hash_map_test.cpp | 70 ++++++++ .../async_key_value_producer_and_consumer.cpp | 28 +-- .../async_key_value_producer_and_consumer.h | 14 +- src/paimon/format/avro/CMakeLists.txt | 2 +- src/paimon/format/orc/CMakeLists.txt | 2 +- 14 files changed, 623 insertions(+), 34 deletions(-) create mode 100644 src/paimon/common/utils/concurrent_backend_factory.h create mode 100644 src/paimon/common/utils/concurrent_bounded_queue.h create mode 100644 src/paimon/common/utils/concurrent_bounded_queue_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 084e6bf03..38c656ca1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -58,6 +58,7 @@ option(PAIMON_BUILD_BENCHMARKS "Build benchmarks" OFF) option(PAIMON_USE_ASAN "Use Address Sanitizer" OFF) option(PAIMON_USE_UBSAN "Use Undefined Behavior Sanitizer" OFF) option(PAIMON_USE_CXX11_ABI "Use C++11 ABI" ON) +option(PAIMON_USE_TBB "Use oneTBB concurrent containers" ON) option(PAIMON_ENABLE_AVRO "Whether to enable avro file format" ON) option(PAIMON_ENABLE_ORC "Whether to enable orc file format" ON) option(PAIMON_ENABLE_JINDO "Whether to enable jindo file system" OFF) @@ -104,6 +105,9 @@ else() endif() add_definitions(-D_GLIBCXX_USE_CXX11_ABI=0) endif() +if(PAIMON_USE_TBB) + add_definitions(-DPAIMON_USE_TBB) +endif() if(PAIMON_ENABLE_LUMINA) add_definitions(-DPAIMON_ENABLE_LUMINA) endif() @@ -360,7 +364,11 @@ include_directories("${CMAKE_SOURCE_DIR}/third_party/roaring_bitmap") include_directories("${CMAKE_SOURCE_DIR}/third_party/xxhash") include_directories(SYSTEM ${ARROW_INCLUDE_DIR}) -include_directories(SYSTEM ${TBB_INCLUDE_DIR}) +set(PAIMON_TBB_LIBS) +if(PAIMON_USE_TBB) + include_directories(SYSTEM ${TBB_INCLUDE_DIR}) + list(APPEND PAIMON_TBB_LIBS tbb) +endif() include_directories(SYSTEM ${GLOG_INCLUDE_DIR}) add_compile_definitions("GLOG_USE_GLOG_EXPORT") diff --git a/cmake_modules/DefineOptions.cmake b/cmake_modules/DefineOptions.cmake index 21f653155..d865fd9b2 100644 --- a/cmake_modules/DefineOptions.cmake +++ b/cmake_modules/DefineOptions.cmake @@ -173,6 +173,8 @@ if("${CMAKE_SOURCE_DIR}" STREQUAL "${CMAKE_CURRENT_SOURCE_DIR}") define_option(PAIMON_DEPENDENCY_USE_SHARED "Prefer shared libraries for system third-party packages" OFF) + define_option(PAIMON_USE_TBB "Use oneTBB concurrent containers" ON) + define_option_string(Arrow_SOURCE "Dependency source for Apache Arrow; SYSTEM is unsupported" "" diff --git a/cmake_modules/ThirdpartyToolchain.cmake b/cmake_modules/ThirdpartyToolchain.cmake index 776519104..78482f34d 100644 --- a/cmake_modules/ThirdpartyToolchain.cmake +++ b/cmake_modules/ThirdpartyToolchain.cmake @@ -1950,7 +1950,9 @@ resolve_dependency(ZLIB) resolve_dependency(LZ4) resolve_dependency(Arrow) paimon_warn_if_mixed_arrow_dependencies() -resolve_dependency(TBB) +if(PAIMON_USE_TBB) + resolve_dependency(TBB) +endif() resolve_dependency(glog) if(PAIMON_ENABLE_AVRO) diff --git a/docs/source/building.rst b/docs/source/building.rst index 32c7e67cd..d901a370b 100644 --- a/docs/source/building.rst +++ b/docs/source/building.rst @@ -184,6 +184,10 @@ boolean flags to ``cmake``. * ``-DPAIMON_ENABLE_TANTIVY=ON``: Enable the experimental Tantivy full-text index Rust FFI. * ``-DPAIMON_ENABLE_REST=ON``: Support for the REST catalog (``metastore=rest``), requires the libcurl and OpenSSL development packages. +* ``-DPAIMON_USE_TBB=ON``: Use oneTBB for the internal concurrent hash map and + bounded queue implementations. This is enabled by default. Set it to ``OFF`` + to use the C++17 standard-library implementations without resolving, building, + or linking TBB. Third-party dependency source ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -217,6 +221,19 @@ require project-specific patches, so their supported source values are -Dfmt_ROOT=/opt/fmt \ -Dzstd_SOURCE=BUNDLED +``TBB_SOURCE`` is only considered when ``PAIMON_USE_TBB=ON``. To build without +any TBB dependency, configure with ``-DPAIMON_USE_TBB=OFF``. + +Custom concurrent backends can be compiled and statically registered through +``ConcurrentBackendFactory`` regardless of this option. With +``PAIMON_USE_TBB=ON``, registrations still run but the TBB-backed containers do +not query the factory. With it set to ``OFF``, a registered backend is selected, +falling back to the built-in C++17 implementation when none is registered. The +registration object and its container specialization must be part of the same +final executable or shared library. If registration code is stored only in a +static archive, make sure its object file is retained by the linker (for example, +by referencing an exported symbol or linking that archive whole). + Use ``PAIMON_PACKAGE_PREFIX`` to provide one common prefix for dependencies whose own ``_ROOT`` variable is not set. Because the patched Arrow and ORC dependencies cannot be resolved from the system, a global ``SYSTEM`` build diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 3de2b667e..e0b252800 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -467,7 +467,7 @@ add_paimon_lib(paimon ${PAIMON_CORE_SRCS} DEPENDENCIES arrow - tbb + ${PAIMON_TBB_LIBS} glog fmt roaring_bitmap @@ -479,7 +479,7 @@ add_paimon_lib(paimon DataSketches STATIC_LINK_LIBS arrow - tbb + ${PAIMON_TBB_LIBS} glog fmt roaring_bitmap @@ -632,6 +632,7 @@ if(PAIMON_BUILD_TESTS) common/utils/arrow/mem_utils_test.cpp common/utils/arrow/status_utils_test.cpp common/utils/concurrent_hash_map_test.cpp + common/utils/concurrent_bounded_queue_test.cpp common/utils/projected_row_test.cpp common/utils/projected_array_test.cpp common/utils/bit_set_test.cpp diff --git a/src/paimon/common/utils/concurrent_backend_factory.h b/src/paimon/common/utils/concurrent_backend_factory.h new file mode 100644 index 000000000..fac2fc0dd --- /dev/null +++ b/src/paimon/common/utils/concurrent_backend_factory.h @@ -0,0 +1,76 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include + +namespace paimon { + +/// Stores one statically linked creator for a concrete concurrent backend interface. +/// +/// A plugin registers a creator for each backend specialization it implements. Registration is +/// effective only when Paimon is built without TBB; TBB-backed containers do not query this +/// factory. The first registration for a specialization wins. The registration and its container +/// specialization must be linked into the same executable or shared library. When registration is +/// packaged in a static archive, its object file must be retained by the final link (for example, +/// by referencing a symbol from it or linking the archive whole). +template +class ConcurrentBackendFactory { + public: + using Creator = std::function()>; + + ConcurrentBackendFactory() = delete; + ~ConcurrentBackendFactory() = delete; + + static bool Register(Creator creator) { + std::lock_guard lock(GetMutex()); + Creator& registered_creator = GetCreator(); + if (registered_creator) { + return false; + } + registered_creator = std::move(creator); + return true; + } + + static std::unique_ptr Create() { + Creator creator; + { + std::lock_guard lock(GetMutex()); + creator = GetCreator(); + } + return creator ? creator() : nullptr; + } + + private: + static Creator& GetCreator() { + static Creator creator; + return creator; + } + + static std::mutex& GetMutex() { + static std::mutex mutex; + return mutex; + } +}; + +} // namespace paimon diff --git a/src/paimon/common/utils/concurrent_bounded_queue.h b/src/paimon/common/utils/concurrent_bounded_queue.h new file mode 100644 index 000000000..87e353e27 --- /dev/null +++ b/src/paimon/common/utils/concurrent_bounded_queue.h @@ -0,0 +1,166 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include + +#include "paimon/common/utils/concurrent_backend_factory.h" +#ifdef PAIMON_USE_TBB +#include "tbb/concurrent_queue.h" +#else +#include +#include +#include +#include +#endif + +namespace paimon { + +template +class ConcurrentBoundedQueueBackend { + public: + virtual ~ConcurrentBoundedQueueBackend() = default; + + virtual void SetCapacity(size_t capacity) = 0; + virtual void Push(T&& value) = 0; + virtual bool TryPop(T& value) = 0; + virtual bool Empty() const = 0; +}; + +#ifndef PAIMON_USE_TBB +namespace detail { + +template +class StdConcurrentBoundedQueueBackend : public ConcurrentBoundedQueueBackend { + public: + void SetCapacity(size_t capacity) override { + { + std::unique_lock lock(mutex_); + capacity_ = capacity; + } + capacity_available_.notify_all(); + } + + void Push(T&& value) override { + std::unique_lock lock(mutex_); + capacity_available_.wait(lock, [this]() { return queue_.size() < capacity_; }); + queue_.push(std::move(value)); + } + + bool TryPop(T& value) override { + { + std::unique_lock lock(mutex_); + if (queue_.empty()) { + return false; + } + value = std::move(queue_.front()); + queue_.pop(); + } + capacity_available_.notify_one(); + return true; + } + + bool Empty() const override { + std::unique_lock lock(mutex_); + return queue_.empty(); + } + + private: + std::queue queue_; + size_t capacity_ = std::numeric_limits::max(); + mutable std::mutex mutex_; + std::condition_variable capacity_available_; +}; + +} // namespace detail +#endif + +template +class ConcurrentBoundedQueue { + public: +#ifdef PAIMON_USE_TBB + ConcurrentBoundedQueue() = default; +#else + ConcurrentBoundedQueue() + : backend_(ConcurrentBackendFactory >::Create()) { + if (backend_ == nullptr) { + backend_ = std::make_unique >(); + } + } +#endif + ~ConcurrentBoundedQueue() = default; + + ConcurrentBoundedQueue(const ConcurrentBoundedQueue&) = delete; + ConcurrentBoundedQueue& operator=(const ConcurrentBoundedQueue&) = delete; + ConcurrentBoundedQueue(ConcurrentBoundedQueue&&) = delete; + ConcurrentBoundedQueue& operator=(ConcurrentBoundedQueue&&) = delete; + + void SetCapacity(size_t capacity) { +#ifdef PAIMON_USE_TBB + queue_.set_capacity(static_cast(capacity)); +#else + backend_->SetCapacity(capacity); +#endif + } + + void Push(const T& value) { +#ifdef PAIMON_USE_TBB + queue_.push(value); +#else + T copied_value = value; + backend_->Push(std::move(copied_value)); +#endif + } + + void Push(T&& value) { +#ifdef PAIMON_USE_TBB + queue_.push(std::move(value)); +#else + backend_->Push(std::move(value)); +#endif + } + + bool TryPop(T& value) { +#ifdef PAIMON_USE_TBB + return queue_.try_pop(value); +#else + return backend_->TryPop(value); +#endif + } + + bool Empty() const { +#ifdef PAIMON_USE_TBB + return queue_.empty(); +#else + return backend_->Empty(); +#endif + } + + private: +#ifdef PAIMON_USE_TBB + tbb::concurrent_bounded_queue queue_; +#else + std::unique_ptr > backend_; +#endif +}; + +} // namespace paimon diff --git a/src/paimon/common/utils/concurrent_bounded_queue_test.cpp b/src/paimon/common/utils/concurrent_bounded_queue_test.cpp new file mode 100644 index 000000000..e56f79199 --- /dev/null +++ b/src/paimon/common/utils/concurrent_bounded_queue_test.cpp @@ -0,0 +1,130 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/common/utils/concurrent_bounded_queue.h" + +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +namespace paimon::test { +namespace { + +struct PluginQueueValue { + int32_t value = 0; +}; + +class PluginQueueBackend : public ConcurrentBoundedQueueBackend { + public: + void SetCapacity(size_t capacity) override { + capacity_ = capacity; + } + + void Push(PluginQueueValue&& value) override { + queue_.push(std::move(value)); + } + + bool TryPop(PluginQueueValue& value) override { + if (queue_.empty()) { + return false; + } + value = std::move(queue_.front()); + queue_.pop(); + return true; + } + + bool Empty() const override { + return queue_.empty(); + } + + size_t Capacity() const { + return capacity_; + } + + private: + std::queue queue_; + size_t capacity_ = 0; +}; + +int32_t plugin_queue_backend_create_count = 0; +const bool plugin_queue_backend_registered = + ConcurrentBackendFactory>::Register([]() { + ++plugin_queue_backend_create_count; + return std::make_unique(); + }); + +} // namespace + +TEST(ConcurrentBoundedQueueTest, TestPushAndTryPop) { + ConcurrentBoundedQueue queue; + queue.SetCapacity(2); + ASSERT_TRUE(queue.Empty()); + + queue.Push(1); + queue.Push(2); + ASSERT_FALSE(queue.Empty()); + + int32_t value = 0; + ASSERT_TRUE(queue.TryPop(value)); + ASSERT_EQ(value, 1); + ASSERT_TRUE(queue.TryPop(value)); + ASSERT_EQ(value, 2); + ASSERT_FALSE(queue.TryPop(value)); + ASSERT_TRUE(queue.Empty()); +} + +TEST(ConcurrentBoundedQueueTest, TestPushWaitsForCapacity) { + ConcurrentBoundedQueue queue; + queue.SetCapacity(1); + queue.Push(1); + + std::future push_future = std::async(std::launch::async, [&queue]() { queue.Push(2); }); + std::future_status initial_status = push_future.wait_for(std::chrono::milliseconds(50)); + + int32_t value = 0; + ASSERT_TRUE(queue.TryPop(value)); + ASSERT_EQ(value, 1); + ASSERT_EQ(initial_status, std::future_status::timeout); + ASSERT_EQ(push_future.wait_for(std::chrono::seconds(1)), std::future_status::ready); + ASSERT_TRUE(queue.TryPop(value)); + ASSERT_EQ(value, 2); +} + +TEST(ConcurrentBoundedQueueTest, TestStaticallyRegisteredBackendSelection) { + ASSERT_TRUE(plugin_queue_backend_registered); + ConcurrentBoundedQueue queue; + queue.SetCapacity(1); + queue.Push(PluginQueueValue{42}); + + PluginQueueValue value; + ASSERT_TRUE(queue.TryPop(value)); + ASSERT_EQ(value.value, 42); +#ifdef PAIMON_USE_TBB + ASSERT_EQ(plugin_queue_backend_create_count, 0); +#else + ASSERT_EQ(plugin_queue_backend_create_count, 1); +#endif +} + +} // namespace paimon::test diff --git a/src/paimon/common/utils/concurrent_hash_map.h b/src/paimon/common/utils/concurrent_hash_map.h index 79a21f098..e98d0d750 100644 --- a/src/paimon/common/utils/concurrent_hash_map.h +++ b/src/paimon/common/utils/concurrent_hash_map.h @@ -22,56 +22,173 @@ #include #include #include +#include +#include +#include #include +#include #include +#include #include +#include "paimon/common/utils/concurrent_backend_factory.h" #include "paimon/common/utils/murmurhash_utils.h" +#ifdef PAIMON_USE_TBB #include "tbb/concurrent_hash_map.h" +#endif namespace paimon { -template > + +template +class DefaultHashCompare { + public: + size_t hash(const Key& key) const { + return std::hash{}(key); + } + + bool equal(const Key& lhs, const Key& rhs) const { + return lhs == rhs; + } +}; + +template +class HashCompareHasher { + public: + size_t operator()(const Key& key) const { + return HashCompare{}.hash(key); + } +}; + +template +class HashCompareEqual { + public: + bool operator()(const Key& lhs, const Key& rhs) const { + return HashCompare{}.equal(lhs, rhs); + } +}; + +template > +class ConcurrentHashMapBackend { + public: + virtual ~ConcurrentHashMapBackend() = default; + + virtual std::optional Find(const Key& key) const = 0; + virtual void Insert(const Key& key, const T& value) = 0; + virtual void Erase(const Key& key) = 0; + virtual size_t Size() const = 0; +}; + +#ifndef PAIMON_USE_TBB +namespace detail { + +template +class StdConcurrentHashMapBackend : public ConcurrentHashMapBackend { + private: + using HashMap = std::unordered_map, + HashCompareEqual>; + + public: + std::optional Find(const Key& key) const override { + std::shared_lock lock(mutex_); + typename HashMap::const_iterator iter = hash_map_.find(key); + if (iter != hash_map_.end()) { + return iter->second; + } + return std::nullopt; + } + + void Insert(const Key& key, const T& value) override { + std::unique_lock lock(mutex_); + hash_map_.insert_or_assign(key, value); + } + + void Erase(const Key& key) override { + std::unique_lock lock(mutex_); + hash_map_.erase(key); + } + + size_t Size() const override { + std::shared_lock lock(mutex_); + return hash_map_.size(); + } + + private: + HashMap hash_map_; + mutable std::shared_mutex mutex_; +}; + +} // namespace detail +#endif + +template > class ConcurrentHashMap { private: - using HashMap = tbb::concurrent_hash_map; + using Backend = ConcurrentHashMapBackend; public: +#ifdef PAIMON_USE_TBB ConcurrentHashMap() = default; +#else + ConcurrentHashMap() : backend_(ConcurrentBackendFactory::Create()) { + if (backend_ == nullptr) { + backend_ = std::make_unique>(); + } + } +#endif ~ConcurrentHashMap() = default; - // No copying allowed ConcurrentHashMap(const ConcurrentHashMap&) = delete; void operator=(const ConcurrentHashMap&) = delete; ConcurrentHashMap(ConcurrentHashMap&&) = delete; ConcurrentHashMap& operator=(ConcurrentHashMap&&) = delete; std::optional Find(const Key& key) const { - typename HashMap::const_accessor accessor; +#ifdef PAIMON_USE_TBB + typename tbb::concurrent_hash_map::const_accessor accessor; if (hash_map_.find(accessor, key)) { return accessor->second; } return std::nullopt; +#else + return backend_->Find(key); +#endif } void Insert(const Key& key, const T& value) { - typename HashMap::accessor accessor; +#ifdef PAIMON_USE_TBB + typename tbb::concurrent_hash_map::accessor accessor; hash_map_.insert(accessor, key); accessor->second = value; +#else + backend_->Insert(key, value); +#endif } void Erase(const Key& key) { - typename HashMap::accessor accessor; +#ifdef PAIMON_USE_TBB + typename tbb::concurrent_hash_map::accessor accessor; if (hash_map_.find(accessor, key)) { hash_map_.erase(accessor); } +#else + backend_->Erase(key); +#endif } size_t Size() const { +#ifdef PAIMON_USE_TBB return hash_map_.size(); +#else + return backend_->Size(); +#endif } private: - HashMap hash_map_; +#ifdef PAIMON_USE_TBB + tbb::concurrent_hash_map hash_map_; +#else + std::unique_ptr backend_; +#endif }; class VectorStringHashCompare { diff --git a/src/paimon/common/utils/concurrent_hash_map_test.cpp b/src/paimon/common/utils/concurrent_hash_map_test.cpp index d3462afb5..ca49da7ef 100644 --- a/src/paimon/common/utils/concurrent_hash_map_test.cpp +++ b/src/paimon/common/utils/concurrent_hash_map_test.cpp @@ -23,12 +23,69 @@ #include #include +#include +#include #include #include "gtest/gtest.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +struct PluginMapKey { + int32_t value = 0; +}; + +class PluginMapKeyHashCompare { + public: + size_t hash(const PluginMapKey& key) const { + return std::hash{}(key.value); + } + + bool equal(const PluginMapKey& lhs, const PluginMapKey& rhs) const { + return lhs.value == rhs.value; + } +}; + +using PluginMapBackend = + ConcurrentHashMapBackend; + +class PluginMapBackendImpl : public PluginMapBackend { + public: + std::optional Find(const PluginMapKey& key) const override { + if (value_ && value_->first.value == key.value) { + return value_->second; + } + return std::nullopt; + } + + void Insert(const PluginMapKey& key, const std::string& value) override { + value_ = std::make_pair(key, value); + } + + void Erase(const PluginMapKey& key) override { + if (value_ && value_->first.value == key.value) { + value_ = std::nullopt; + } + } + + size_t Size() const override { + return value_ ? 1 : 0; + } + + private: + std::optional> value_; +}; + +int32_t plugin_map_backend_create_count = 0; +const bool plugin_map_backend_registered = + ConcurrentBackendFactory::Register([]() { + ++plugin_map_backend_create_count; + return std::make_unique(); + }); + +} // namespace TEST(ConcurrentHashMapTest, TestSimple) { ConcurrentHashMap hash_map; @@ -154,4 +211,17 @@ TEST(ConcurrentHashMapTest, TestMultiThreadInsertAndFindAndDelete) { } } +TEST(ConcurrentHashMapTest, TestStaticallyRegisteredBackendSelection) { + ASSERT_TRUE(plugin_map_backend_registered); + ConcurrentHashMap hash_map; + hash_map.Insert(PluginMapKey{1}, "plugin"); + ASSERT_EQ(hash_map.Find(PluginMapKey{1}), "plugin"); + +#ifdef PAIMON_USE_TBB + ASSERT_EQ(plugin_map_backend_create_count, 0); +#else + ASSERT_EQ(plugin_map_backend_create_count, 1); +#endif +} + } // namespace paimon::test diff --git a/src/paimon/core/io/async_key_value_producer_and_consumer.cpp b/src/paimon/core/io/async_key_value_producer_and_consumer.cpp index 1792b43cd..682f91592 100644 --- a/src/paimon/core/io/async_key_value_producer_and_consumer.cpp +++ b/src/paimon/core/io/async_key_value_producer_and_consumer.cpp @@ -40,8 +40,8 @@ AsyncKeyValueProducerAndConsumer::AsyncKeyValueProducerAndConsumer( pool_(pool), sort_merge_reader_(std::move(sort_merge_reader)), create_consumer_(std::move(create_consumer)) { - kv_queue_.set_capacity(consumer_thread_num * 2); - result_queue_.set_capacity(RESULT_BATCH_COUNT); + kv_queue_.SetCapacity(consumer_thread_num * 2); + result_queue_.SetCapacity(RESULT_BATCH_COUNT); } template @@ -94,9 +94,9 @@ Result AsyncKeyValueProducerAndConsumer::NextBatch() { } R result; - while (!result_queue_.try_pop(result)) { + while (!result_queue_.TryPop(result)) { PAIMON_RETURN_NOT_OK(CheckStatusAndCleanUp()); - if (consumer_finished_count_ == consumer_thread_num_ && result_queue_.empty()) { + if (consumer_finished_count_ == consumer_thread_num_ && result_queue_.Empty()) { // all consume thread finished next_batch_finished_ = true; return R(); @@ -124,7 +124,7 @@ Status AsyncKeyValueProducerAndConsumer::ProduceLoop() { } batch.push_back(std::move(iterator->Next())); if (static_cast(batch.size()) >= batch_size_) { - kv_queue_.push(std::move(batch)); + kv_queue_.Push(std::move(batch)); batch = std::vector(); batch.reserve(batch_size_); } @@ -132,10 +132,10 @@ Status AsyncKeyValueProducerAndConsumer::ProduceLoop() { } // Push remaining rows if (!batch.empty()) { - kv_queue_.push(std::move(batch)); + kv_queue_.Push(std::move(batch)); } // Push empty batch as EOF signal - kv_queue_.push(std::vector()); + kv_queue_.Push(std::vector()); return Status::OK(); } @@ -156,7 +156,7 @@ void AsyncKeyValueProducerAndConsumer::CleanUp() { template void AsyncKeyValueProducerAndConsumer::CleanUpQueue() { R read_batch; - while (result_queue_.try_pop(read_batch)) { + while (result_queue_.TryPop(read_batch)) { if constexpr (std::is_same_v) { if (!BatchReader::IsEofBatch(read_batch)) { ReaderUtils::ReleaseReadBatch(std::move(read_batch)); @@ -169,7 +169,7 @@ void AsyncKeyValueProducerAndConsumer::CleanUpQueue() { } std::vector kv_batch; - while (kv_queue_.try_pop(kv_batch)) { + while (kv_queue_.TryPop(kv_batch)) { } } @@ -180,8 +180,8 @@ template AsyncKeyValueConsumer::AsyncKeyValueConsumer( std::unique_ptr>&& key_value_consumer, std::atomic& consume_finished, std::atomic& consumer_finished_count, - tbb::concurrent_bounded_queue>& kv_queue, - tbb::concurrent_bounded_queue& result_queue) + ConcurrentBoundedQueue>& kv_queue, + ConcurrentBoundedQueue& result_queue) : key_value_consumer_(std::move(key_value_consumer)), consume_finished_(consume_finished), consumer_finished_count_(consumer_finished_count), @@ -206,17 +206,17 @@ template Status AsyncKeyValueConsumer::ConsumeLoop() { while (!consume_finished_) { std::vector key_value_vec; - if (!kv_queue_.try_pop(key_value_vec)) { + if (!kv_queue_.TryPop(key_value_vec)) { usleep(100); continue; } if (key_value_vec.empty()) { // Empty batch is EOF signal; re-push for other consumers - kv_queue_.push(std::move(key_value_vec)); + kv_queue_.Push(std::move(key_value_vec)); break; } PAIMON_ASSIGN_OR_RAISE(R result, key_value_consumer_->NextBatch(key_value_vec)); - result_queue_.push(std::move(result)); + result_queue_.Push(std::move(result)); } consumer_finished_count_++; return Status::OK(); diff --git a/src/paimon/core/io/async_key_value_producer_and_consumer.h b/src/paimon/core/io/async_key_value_producer_and_consumer.h index af8bbed68..99ae53d5b 100644 --- a/src/paimon/core/io/async_key_value_producer_and_consumer.h +++ b/src/paimon/core/io/async_key_value_producer_and_consumer.h @@ -27,12 +27,12 @@ #include #include "arrow/api.h" +#include "paimon/common/utils/concurrent_bounded_queue.h" #include "paimon/core/io/row_to_arrow_array_converter.h" #include "paimon/core/key_value.h" #include "paimon/core/mergetree/compact/sort_merge_reader.h" #include "paimon/result.h" #include "paimon/status.h" -#include "tbb/concurrent_queue.h" namespace paimon { template @@ -94,8 +94,8 @@ class AsyncKeyValueProducerAndConsumer { std::shared_future producer_future_; std::vector>> consumers_; std::atomic consumer_finished_count_ = 0; - tbb::concurrent_bounded_queue> kv_queue_; - tbb::concurrent_bounded_queue result_queue_; + ConcurrentBoundedQueue> kv_queue_; + ConcurrentBoundedQueue result_queue_; }; template @@ -104,8 +104,8 @@ class AsyncKeyValueConsumer { AsyncKeyValueConsumer(std::unique_ptr>&& key_value_consumer, std::atomic& consume_finished, std::atomic& consumer_finished_count, - tbb::concurrent_bounded_queue>& kv_queue, - tbb::concurrent_bounded_queue& result_queue); + ConcurrentBoundedQueue>& kv_queue, + ConcurrentBoundedQueue& result_queue); ~AsyncKeyValueConsumer() { CleanUp(); @@ -122,8 +122,8 @@ class AsyncKeyValueConsumer { std::shared_future consumer_future_; std::atomic& consume_finished_; std::atomic& consumer_finished_count_; - tbb::concurrent_bounded_queue>& kv_queue_; - tbb::concurrent_bounded_queue& result_queue_; + ConcurrentBoundedQueue>& kv_queue_; + ConcurrentBoundedQueue& result_queue_; }; } // namespace paimon diff --git a/src/paimon/format/avro/CMakeLists.txt b/src/paimon/format/avro/CMakeLists.txt index 1c4256741..9b61c0ecc 100644 --- a/src/paimon/format/avro/CMakeLists.txt +++ b/src/paimon/format/avro/CMakeLists.txt @@ -39,7 +39,7 @@ if(PAIMON_ENABLE_AVRO) glog fmt avro - tbb + ${PAIMON_TBB_LIBS} Threads::Threads SHARED_LINK_LIBS paimon_shared diff --git a/src/paimon/format/orc/CMakeLists.txt b/src/paimon/format/orc/CMakeLists.txt index d86750ea7..7641686ef 100644 --- a/src/paimon/format/orc/CMakeLists.txt +++ b/src/paimon/format/orc/CMakeLists.txt @@ -39,7 +39,7 @@ if(PAIMON_ENABLE_ORC) glog fmt orc::orc - tbb + ${PAIMON_TBB_LIBS} Threads::Threads SHARED_LINK_LIBS paimon_shared From d90fe94d1a43604bb4f201bade15fdc07bd24f6e Mon Sep 17 00:00:00 2001 From: zhongzheyun Date: Thu, 27 Aug 2026 18:03:32 +0800 Subject: [PATCH 2/3] refactor(utils): separate non-TBB concurrent containers --- docs/source/building.rst | 10 -- .../common/utils/concurrent_backend_factory.h | 76 ----------- .../common/utils/concurrent_bounded_queue.h | 127 +++-------------- .../utils/concurrent_bounded_queue_test.cpp | 93 ++----------- src/paimon/common/utils/concurrent_hash_map.h | 116 ++-------------- .../common/utils/concurrent_hash_map_test.cpp | 69 ---------- .../utils/std_concurrent_bounded_queue.h | 92 +++++++++++++ .../common/utils/std_concurrent_hash_map.h | 129 ++++++++++++++++++ .../async_key_value_producer_and_consumer.cpp | 24 ++-- 9 files changed, 275 insertions(+), 461 deletions(-) delete mode 100644 src/paimon/common/utils/concurrent_backend_factory.h create mode 100644 src/paimon/common/utils/std_concurrent_bounded_queue.h create mode 100644 src/paimon/common/utils/std_concurrent_hash_map.h diff --git a/docs/source/building.rst b/docs/source/building.rst index d901a370b..88dee2d75 100644 --- a/docs/source/building.rst +++ b/docs/source/building.rst @@ -224,16 +224,6 @@ require project-specific patches, so their supported source values are ``TBB_SOURCE`` is only considered when ``PAIMON_USE_TBB=ON``. To build without any TBB dependency, configure with ``-DPAIMON_USE_TBB=OFF``. -Custom concurrent backends can be compiled and statically registered through -``ConcurrentBackendFactory`` regardless of this option. With -``PAIMON_USE_TBB=ON``, registrations still run but the TBB-backed containers do -not query the factory. With it set to ``OFF``, a registered backend is selected, -falling back to the built-in C++17 implementation when none is registered. The -registration object and its container specialization must be part of the same -final executable or shared library. If registration code is stored only in a -static archive, make sure its object file is retained by the linker (for example, -by referencing an exported symbol or linking that archive whole). - Use ``PAIMON_PACKAGE_PREFIX`` to provide one common prefix for dependencies whose own ``_ROOT`` variable is not set. Because the patched Arrow and ORC dependencies cannot be resolved from the system, a global ``SYSTEM`` build diff --git a/src/paimon/common/utils/concurrent_backend_factory.h b/src/paimon/common/utils/concurrent_backend_factory.h deleted file mode 100644 index fac2fc0dd..000000000 --- a/src/paimon/common/utils/concurrent_backend_factory.h +++ /dev/null @@ -1,76 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -#pragma once - -#include -#include -#include -#include - -namespace paimon { - -/// Stores one statically linked creator for a concrete concurrent backend interface. -/// -/// A plugin registers a creator for each backend specialization it implements. Registration is -/// effective only when Paimon is built without TBB; TBB-backed containers do not query this -/// factory. The first registration for a specialization wins. The registration and its container -/// specialization must be linked into the same executable or shared library. When registration is -/// packaged in a static archive, its object file must be retained by the final link (for example, -/// by referencing a symbol from it or linking the archive whole). -template -class ConcurrentBackendFactory { - public: - using Creator = std::function()>; - - ConcurrentBackendFactory() = delete; - ~ConcurrentBackendFactory() = delete; - - static bool Register(Creator creator) { - std::lock_guard lock(GetMutex()); - Creator& registered_creator = GetCreator(); - if (registered_creator) { - return false; - } - registered_creator = std::move(creator); - return true; - } - - static std::unique_ptr Create() { - Creator creator; - { - std::lock_guard lock(GetMutex()); - creator = GetCreator(); - } - return creator ? creator() : nullptr; - } - - private: - static Creator& GetCreator() { - static Creator creator; - return creator; - } - - static std::mutex& GetMutex() { - static std::mutex mutex; - return mutex; - } -}; - -} // namespace paimon diff --git a/src/paimon/common/utils/concurrent_bounded_queue.h b/src/paimon/common/utils/concurrent_bounded_queue.h index 87e353e27..642d5b227 100644 --- a/src/paimon/common/utils/concurrent_bounded_queue.h +++ b/src/paimon/common/utils/concurrent_bounded_queue.h @@ -19,94 +19,24 @@ #pragma once +#ifdef PAIMON_USE_TBB + #include -#include #include -#include "paimon/common/utils/concurrent_backend_factory.h" -#ifdef PAIMON_USE_TBB #include "tbb/concurrent_queue.h" -#else -#include -#include -#include -#include -#endif namespace paimon { -template -class ConcurrentBoundedQueueBackend { - public: - virtual ~ConcurrentBoundedQueueBackend() = default; - - virtual void SetCapacity(size_t capacity) = 0; - virtual void Push(T&& value) = 0; - virtual bool TryPop(T& value) = 0; - virtual bool Empty() const = 0; -}; - -#ifndef PAIMON_USE_TBB -namespace detail { - -template -class StdConcurrentBoundedQueueBackend : public ConcurrentBoundedQueueBackend { - public: - void SetCapacity(size_t capacity) override { - { - std::unique_lock lock(mutex_); - capacity_ = capacity; - } - capacity_available_.notify_all(); - } - - void Push(T&& value) override { - std::unique_lock lock(mutex_); - capacity_available_.wait(lock, [this]() { return queue_.size() < capacity_; }); - queue_.push(std::move(value)); - } - - bool TryPop(T& value) override { - { - std::unique_lock lock(mutex_); - if (queue_.empty()) { - return false; - } - value = std::move(queue_.front()); - queue_.pop(); - } - capacity_available_.notify_one(); - return true; - } - - bool Empty() const override { - std::unique_lock lock(mutex_); - return queue_.empty(); - } - - private: - std::queue queue_; - size_t capacity_ = std::numeric_limits::max(); - mutable std::mutex mutex_; - std::condition_variable capacity_available_; -}; - -} // namespace detail -#endif - template class ConcurrentBoundedQueue { public: -#ifdef PAIMON_USE_TBB + using size_type = std::ptrdiff_t; + using value_type = T; + using reference = T&; + using const_reference = const T&; + ConcurrentBoundedQueue() = default; -#else - ConcurrentBoundedQueue() - : backend_(ConcurrentBackendFactory >::Create()) { - if (backend_ == nullptr) { - backend_ = std::make_unique >(); - } - } -#endif ~ConcurrentBoundedQueue() = default; ConcurrentBoundedQueue(const ConcurrentBoundedQueue&) = delete; @@ -114,53 +44,34 @@ class ConcurrentBoundedQueue { ConcurrentBoundedQueue(ConcurrentBoundedQueue&&) = delete; ConcurrentBoundedQueue& operator=(ConcurrentBoundedQueue&&) = delete; - void SetCapacity(size_t capacity) { -#ifdef PAIMON_USE_TBB - queue_.set_capacity(static_cast(capacity)); -#else - backend_->SetCapacity(capacity); -#endif + void set_capacity(size_type capacity) { + queue_.set_capacity(capacity); } - void Push(const T& value) { -#ifdef PAIMON_USE_TBB + void push(const T& value) { queue_.push(value); -#else - T copied_value = value; - backend_->Push(std::move(copied_value)); -#endif } - void Push(T&& value) { -#ifdef PAIMON_USE_TBB + void push(T&& value) { queue_.push(std::move(value)); -#else - backend_->Push(std::move(value)); -#endif } - bool TryPop(T& value) { -#ifdef PAIMON_USE_TBB + bool try_pop(T& value) { return queue_.try_pop(value); -#else - return backend_->TryPop(value); -#endif } - bool Empty() const { -#ifdef PAIMON_USE_TBB + bool empty() const { return queue_.empty(); -#else - return backend_->Empty(); -#endif } private: -#ifdef PAIMON_USE_TBB tbb::concurrent_bounded_queue queue_; -#else - std::unique_ptr > backend_; -#endif }; } // namespace paimon + +#else + +#include "paimon/common/utils/std_concurrent_bounded_queue.h" + +#endif diff --git a/src/paimon/common/utils/concurrent_bounded_queue_test.cpp b/src/paimon/common/utils/concurrent_bounded_queue_test.cpp index e56f79199..ed2479aa6 100644 --- a/src/paimon/common/utils/concurrent_bounded_queue_test.cpp +++ b/src/paimon/common/utils/concurrent_bounded_queue_test.cpp @@ -22,109 +22,44 @@ #include #include #include -#include -#include -#include #include "gtest/gtest.h" namespace paimon::test { -namespace { - -struct PluginQueueValue { - int32_t value = 0; -}; - -class PluginQueueBackend : public ConcurrentBoundedQueueBackend { - public: - void SetCapacity(size_t capacity) override { - capacity_ = capacity; - } - - void Push(PluginQueueValue&& value) override { - queue_.push(std::move(value)); - } - - bool TryPop(PluginQueueValue& value) override { - if (queue_.empty()) { - return false; - } - value = std::move(queue_.front()); - queue_.pop(); - return true; - } - - bool Empty() const override { - return queue_.empty(); - } - - size_t Capacity() const { - return capacity_; - } - - private: - std::queue queue_; - size_t capacity_ = 0; -}; - -int32_t plugin_queue_backend_create_count = 0; -const bool plugin_queue_backend_registered = - ConcurrentBackendFactory>::Register([]() { - ++plugin_queue_backend_create_count; - return std::make_unique(); - }); - -} // namespace TEST(ConcurrentBoundedQueueTest, TestPushAndTryPop) { ConcurrentBoundedQueue queue; - queue.SetCapacity(2); - ASSERT_TRUE(queue.Empty()); + queue.set_capacity(2); + ASSERT_TRUE(queue.empty()); - queue.Push(1); - queue.Push(2); - ASSERT_FALSE(queue.Empty()); + queue.push(1); + queue.push(2); + ASSERT_FALSE(queue.empty()); int32_t value = 0; - ASSERT_TRUE(queue.TryPop(value)); + ASSERT_TRUE(queue.try_pop(value)); ASSERT_EQ(value, 1); - ASSERT_TRUE(queue.TryPop(value)); + ASSERT_TRUE(queue.try_pop(value)); ASSERT_EQ(value, 2); - ASSERT_FALSE(queue.TryPop(value)); - ASSERT_TRUE(queue.Empty()); + ASSERT_FALSE(queue.try_pop(value)); + ASSERT_TRUE(queue.empty()); } TEST(ConcurrentBoundedQueueTest, TestPushWaitsForCapacity) { ConcurrentBoundedQueue queue; - queue.SetCapacity(1); - queue.Push(1); + queue.set_capacity(1); + queue.push(1); - std::future push_future = std::async(std::launch::async, [&queue]() { queue.Push(2); }); + std::future push_future = std::async(std::launch::async, [&queue]() { queue.push(2); }); std::future_status initial_status = push_future.wait_for(std::chrono::milliseconds(50)); int32_t value = 0; - ASSERT_TRUE(queue.TryPop(value)); + ASSERT_TRUE(queue.try_pop(value)); ASSERT_EQ(value, 1); ASSERT_EQ(initial_status, std::future_status::timeout); ASSERT_EQ(push_future.wait_for(std::chrono::seconds(1)), std::future_status::ready); - ASSERT_TRUE(queue.TryPop(value)); + ASSERT_TRUE(queue.try_pop(value)); ASSERT_EQ(value, 2); } -TEST(ConcurrentBoundedQueueTest, TestStaticallyRegisteredBackendSelection) { - ASSERT_TRUE(plugin_queue_backend_registered); - ConcurrentBoundedQueue queue; - queue.SetCapacity(1); - queue.Push(PluginQueueValue{42}); - - PluginQueueValue value; - ASSERT_TRUE(queue.TryPop(value)); - ASSERT_EQ(value.value, 42); -#ifdef PAIMON_USE_TBB - ASSERT_EQ(plugin_queue_backend_create_count, 0); -#else - ASSERT_EQ(plugin_queue_backend_create_count, 1); -#endif -} - } // namespace paimon::test diff --git a/src/paimon/common/utils/concurrent_hash_map.h b/src/paimon/common/utils/concurrent_hash_map.h index e98d0d750..cedb49111 100644 --- a/src/paimon/common/utils/concurrent_hash_map.h +++ b/src/paimon/common/utils/concurrent_hash_map.h @@ -19,23 +19,18 @@ #pragma once +#ifdef PAIMON_USE_TBB + #include #include #include #include -#include -#include #include -#include #include -#include #include -#include "paimon/common/utils/concurrent_backend_factory.h" #include "paimon/common/utils/murmurhash_utils.h" -#ifdef PAIMON_USE_TBB #include "tbb/concurrent_hash_map.h" -#endif namespace paimon { @@ -51,90 +46,10 @@ class DefaultHashCompare { } }; -template -class HashCompareHasher { - public: - size_t operator()(const Key& key) const { - return HashCompare{}.hash(key); - } -}; - -template -class HashCompareEqual { - public: - bool operator()(const Key& lhs, const Key& rhs) const { - return HashCompare{}.equal(lhs, rhs); - } -}; - -template > -class ConcurrentHashMapBackend { - public: - virtual ~ConcurrentHashMapBackend() = default; - - virtual std::optional Find(const Key& key) const = 0; - virtual void Insert(const Key& key, const T& value) = 0; - virtual void Erase(const Key& key) = 0; - virtual size_t Size() const = 0; -}; - -#ifndef PAIMON_USE_TBB -namespace detail { - -template -class StdConcurrentHashMapBackend : public ConcurrentHashMapBackend { - private: - using HashMap = std::unordered_map, - HashCompareEqual>; - - public: - std::optional Find(const Key& key) const override { - std::shared_lock lock(mutex_); - typename HashMap::const_iterator iter = hash_map_.find(key); - if (iter != hash_map_.end()) { - return iter->second; - } - return std::nullopt; - } - - void Insert(const Key& key, const T& value) override { - std::unique_lock lock(mutex_); - hash_map_.insert_or_assign(key, value); - } - - void Erase(const Key& key) override { - std::unique_lock lock(mutex_); - hash_map_.erase(key); - } - - size_t Size() const override { - std::shared_lock lock(mutex_); - return hash_map_.size(); - } - - private: - HashMap hash_map_; - mutable std::shared_mutex mutex_; -}; - -} // namespace detail -#endif - template > class ConcurrentHashMap { - private: - using Backend = ConcurrentHashMapBackend; - public: -#ifdef PAIMON_USE_TBB ConcurrentHashMap() = default; -#else - ConcurrentHashMap() : backend_(ConcurrentBackendFactory::Create()) { - if (backend_ == nullptr) { - backend_ = std::make_unique>(); - } - } -#endif ~ConcurrentHashMap() = default; ConcurrentHashMap(const ConcurrentHashMap&) = delete; @@ -143,52 +58,32 @@ class ConcurrentHashMap { ConcurrentHashMap& operator=(ConcurrentHashMap&&) = delete; std::optional Find(const Key& key) const { -#ifdef PAIMON_USE_TBB typename tbb::concurrent_hash_map::const_accessor accessor; if (hash_map_.find(accessor, key)) { return accessor->second; } return std::nullopt; -#else - return backend_->Find(key); -#endif } void Insert(const Key& key, const T& value) { -#ifdef PAIMON_USE_TBB typename tbb::concurrent_hash_map::accessor accessor; hash_map_.insert(accessor, key); accessor->second = value; -#else - backend_->Insert(key, value); -#endif } void Erase(const Key& key) { -#ifdef PAIMON_USE_TBB typename tbb::concurrent_hash_map::accessor accessor; if (hash_map_.find(accessor, key)) { hash_map_.erase(accessor); } -#else - backend_->Erase(key); -#endif } size_t Size() const { -#ifdef PAIMON_USE_TBB return hash_map_.size(); -#else - return backend_->Size(); -#endif } private: -#ifdef PAIMON_USE_TBB tbb::concurrent_hash_map hash_map_; -#else - std::unique_ptr backend_; -#endif }; class VectorStringHashCompare { @@ -206,4 +101,11 @@ class VectorStringHashCompare { return a == b; } }; + } // namespace paimon + +#else + +#include "paimon/common/utils/std_concurrent_hash_map.h" + +#endif diff --git a/src/paimon/common/utils/concurrent_hash_map_test.cpp b/src/paimon/common/utils/concurrent_hash_map_test.cpp index ca49da7ef..95a2f7fa9 100644 --- a/src/paimon/common/utils/concurrent_hash_map_test.cpp +++ b/src/paimon/common/utils/concurrent_hash_map_test.cpp @@ -23,7 +23,6 @@ #include #include -#include #include #include @@ -31,61 +30,6 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { -namespace { - -struct PluginMapKey { - int32_t value = 0; -}; - -class PluginMapKeyHashCompare { - public: - size_t hash(const PluginMapKey& key) const { - return std::hash{}(key.value); - } - - bool equal(const PluginMapKey& lhs, const PluginMapKey& rhs) const { - return lhs.value == rhs.value; - } -}; - -using PluginMapBackend = - ConcurrentHashMapBackend; - -class PluginMapBackendImpl : public PluginMapBackend { - public: - std::optional Find(const PluginMapKey& key) const override { - if (value_ && value_->first.value == key.value) { - return value_->second; - } - return std::nullopt; - } - - void Insert(const PluginMapKey& key, const std::string& value) override { - value_ = std::make_pair(key, value); - } - - void Erase(const PluginMapKey& key) override { - if (value_ && value_->first.value == key.value) { - value_ = std::nullopt; - } - } - - size_t Size() const override { - return value_ ? 1 : 0; - } - - private: - std::optional> value_; -}; - -int32_t plugin_map_backend_create_count = 0; -const bool plugin_map_backend_registered = - ConcurrentBackendFactory::Register([]() { - ++plugin_map_backend_create_count; - return std::make_unique(); - }); - -} // namespace TEST(ConcurrentHashMapTest, TestSimple) { ConcurrentHashMap hash_map; @@ -211,17 +155,4 @@ TEST(ConcurrentHashMapTest, TestMultiThreadInsertAndFindAndDelete) { } } -TEST(ConcurrentHashMapTest, TestStaticallyRegisteredBackendSelection) { - ASSERT_TRUE(plugin_map_backend_registered); - ConcurrentHashMap hash_map; - hash_map.Insert(PluginMapKey{1}, "plugin"); - ASSERT_EQ(hash_map.Find(PluginMapKey{1}), "plugin"); - -#ifdef PAIMON_USE_TBB - ASSERT_EQ(plugin_map_backend_create_count, 0); -#else - ASSERT_EQ(plugin_map_backend_create_count, 1); -#endif -} - } // namespace paimon::test diff --git a/src/paimon/common/utils/std_concurrent_bounded_queue.h b/src/paimon/common/utils/std_concurrent_bounded_queue.h new file mode 100644 index 000000000..5a009cff4 --- /dev/null +++ b/src/paimon/common/utils/std_concurrent_bounded_queue.h @@ -0,0 +1,92 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace paimon { + +template +class ConcurrentBoundedQueue { + public: + using size_type = std::ptrdiff_t; + using value_type = T; + using reference = T&; + using const_reference = const T&; + + ConcurrentBoundedQueue() = default; + ~ConcurrentBoundedQueue() = default; + + ConcurrentBoundedQueue(const ConcurrentBoundedQueue&) = delete; + ConcurrentBoundedQueue& operator=(const ConcurrentBoundedQueue&) = delete; + ConcurrentBoundedQueue(ConcurrentBoundedQueue&&) = delete; + ConcurrentBoundedQueue& operator=(ConcurrentBoundedQueue&&) = delete; + + void set_capacity(size_type capacity) { + { + std::unique_lock lock(mutex_); + capacity_ = + capacity < 0 ? std::numeric_limits::max() : static_cast(capacity); + } + capacity_available_.notify_all(); + } + + void push(const T& value) { + T copied_value = value; + push(std::move(copied_value)); + } + + void push(T&& value) { + std::unique_lock lock(mutex_); + capacity_available_.wait(lock, [this]() { return queue_.size() < capacity_; }); + queue_.push(std::move(value)); + } + + bool try_pop(T& value) { + { + std::unique_lock lock(mutex_); + if (queue_.empty()) { + return false; + } + value = std::move(queue_.front()); + queue_.pop(); + } + capacity_available_.notify_one(); + return true; + } + + bool empty() const { + std::unique_lock lock(mutex_); + return queue_.empty(); + } + + private: + std::queue queue_; + size_t capacity_ = std::numeric_limits::max(); + mutable std::mutex mutex_; + std::condition_variable capacity_available_; +}; + +} // namespace paimon diff --git a/src/paimon/common/utils/std_concurrent_hash_map.h b/src/paimon/common/utils/std_concurrent_hash_map.h new file mode 100644 index 000000000..b0ad856ca --- /dev/null +++ b/src/paimon/common/utils/std_concurrent_hash_map.h @@ -0,0 +1,129 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "paimon/common/utils/murmurhash_utils.h" + +namespace paimon { + +template +class DefaultHashCompare { + public: + size_t hash(const Key& key) const { + return std::hash{}(key); + } + + bool equal(const Key& lhs, const Key& rhs) const { + return lhs == rhs; + } +}; + +namespace detail { + +template +class HashCompareHasher { + public: + size_t operator()(const Key& key) const { + return HashCompare{}.hash(key); + } +}; + +template +class HashCompareEqual { + public: + bool operator()(const Key& lhs, const Key& rhs) const { + return HashCompare{}.equal(lhs, rhs); + } +}; + +} // namespace detail + +template > +class ConcurrentHashMap { + private: + using HashMap = std::unordered_map, + detail::HashCompareEqual>; + + public: + ConcurrentHashMap() = default; + ~ConcurrentHashMap() = default; + + ConcurrentHashMap(const ConcurrentHashMap&) = delete; + void operator=(const ConcurrentHashMap&) = delete; + ConcurrentHashMap(ConcurrentHashMap&&) = delete; + ConcurrentHashMap& operator=(ConcurrentHashMap&&) = delete; + + std::optional Find(const Key& key) const { + std::shared_lock lock(mutex_); + typename HashMap::const_iterator iter = hash_map_.find(key); + if (iter != hash_map_.end()) { + return iter->second; + } + return std::nullopt; + } + + void Insert(const Key& key, const T& value) { + std::unique_lock lock(mutex_); + hash_map_.insert_or_assign(key, value); + } + + void Erase(const Key& key) { + std::unique_lock lock(mutex_); + hash_map_.erase(key); + } + + size_t Size() const { + std::shared_lock lock(mutex_); + return hash_map_.size(); + } + + private: + HashMap hash_map_; + mutable std::shared_mutex mutex_; +}; + +class VectorStringHashCompare { + public: + size_t hash(const std::vector& key) const { + int32_t ret = MurmurHashUtils::DEFAULT_SEED; + for (const auto& s : key) { + ret = MurmurHashUtils::HashUnsafeBytes(reinterpret_cast(s.data()), + /*offset=*/0, s.size(), ret); + } + return ret; + } + + bool equal(const std::vector& a, const std::vector& b) const { + return a == b; + } +}; + +} // namespace paimon diff --git a/src/paimon/core/io/async_key_value_producer_and_consumer.cpp b/src/paimon/core/io/async_key_value_producer_and_consumer.cpp index 682f91592..a609d8923 100644 --- a/src/paimon/core/io/async_key_value_producer_and_consumer.cpp +++ b/src/paimon/core/io/async_key_value_producer_and_consumer.cpp @@ -40,8 +40,8 @@ AsyncKeyValueProducerAndConsumer::AsyncKeyValueProducerAndConsumer( pool_(pool), sort_merge_reader_(std::move(sort_merge_reader)), create_consumer_(std::move(create_consumer)) { - kv_queue_.SetCapacity(consumer_thread_num * 2); - result_queue_.SetCapacity(RESULT_BATCH_COUNT); + kv_queue_.set_capacity(consumer_thread_num * 2); + result_queue_.set_capacity(RESULT_BATCH_COUNT); } template @@ -94,9 +94,9 @@ Result AsyncKeyValueProducerAndConsumer::NextBatch() { } R result; - while (!result_queue_.TryPop(result)) { + while (!result_queue_.try_pop(result)) { PAIMON_RETURN_NOT_OK(CheckStatusAndCleanUp()); - if (consumer_finished_count_ == consumer_thread_num_ && result_queue_.Empty()) { + if (consumer_finished_count_ == consumer_thread_num_ && result_queue_.empty()) { // all consume thread finished next_batch_finished_ = true; return R(); @@ -124,7 +124,7 @@ Status AsyncKeyValueProducerAndConsumer::ProduceLoop() { } batch.push_back(std::move(iterator->Next())); if (static_cast(batch.size()) >= batch_size_) { - kv_queue_.Push(std::move(batch)); + kv_queue_.push(std::move(batch)); batch = std::vector(); batch.reserve(batch_size_); } @@ -132,10 +132,10 @@ Status AsyncKeyValueProducerAndConsumer::ProduceLoop() { } // Push remaining rows if (!batch.empty()) { - kv_queue_.Push(std::move(batch)); + kv_queue_.push(std::move(batch)); } // Push empty batch as EOF signal - kv_queue_.Push(std::vector()); + kv_queue_.push(std::vector()); return Status::OK(); } @@ -156,7 +156,7 @@ void AsyncKeyValueProducerAndConsumer::CleanUp() { template void AsyncKeyValueProducerAndConsumer::CleanUpQueue() { R read_batch; - while (result_queue_.TryPop(read_batch)) { + while (result_queue_.try_pop(read_batch)) { if constexpr (std::is_same_v) { if (!BatchReader::IsEofBatch(read_batch)) { ReaderUtils::ReleaseReadBatch(std::move(read_batch)); @@ -169,7 +169,7 @@ void AsyncKeyValueProducerAndConsumer::CleanUpQueue() { } std::vector kv_batch; - while (kv_queue_.TryPop(kv_batch)) { + while (kv_queue_.try_pop(kv_batch)) { } } @@ -206,17 +206,17 @@ template Status AsyncKeyValueConsumer::ConsumeLoop() { while (!consume_finished_) { std::vector key_value_vec; - if (!kv_queue_.TryPop(key_value_vec)) { + if (!kv_queue_.try_pop(key_value_vec)) { usleep(100); continue; } if (key_value_vec.empty()) { // Empty batch is EOF signal; re-push for other consumers - kv_queue_.Push(std::move(key_value_vec)); + kv_queue_.push(std::move(key_value_vec)); break; } PAIMON_ASSIGN_OR_RAISE(R result, key_value_consumer_->NextBatch(key_value_vec)); - result_queue_.Push(std::move(result)); + result_queue_.push(std::move(result)); } consumer_finished_count_++; return Status::OK(); From f481ffa8cd1a58926f3f0accf9c3e35868eea2b0 Mon Sep 17 00:00:00 2001 From: zhongzheyun Date: Thu, 27 Aug 2026 20:07:37 +0800 Subject: [PATCH 3/3] fix: address non-TBB concurrent container issues --- src/paimon/CMakeLists.txt | 1 + .../common/utils/concurrent_hash_map_test.cpp | 26 ++++ .../common/utils/std_concurrent_hash_map.h | 23 ++- .../async_key_value_producer_and_consumer.cpp | 30 +++- .../async_key_value_producer_and_consumer.h | 3 + ...c_key_value_producer_and_consumer_test.cpp | 132 ++++++++++++++++++ 6 files changed, 211 insertions(+), 4 deletions(-) create mode 100644 src/paimon/core/io/async_key_value_producer_and_consumer_test.cpp diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index e0b252800..bb47f23c0 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -756,6 +756,7 @@ if(PAIMON_BUILD_TESTS) core/index/pk/primary_key_index_definitions_test.cpp core/index/pksorted/pk_sorted_bucket_index_state_test.cpp core/index/index_file_handler_test.cpp + core/io/async_key_value_producer_and_consumer_test.cpp core/io/compact_increment_test.cpp core/io/infer_shredding_file_writer_test.cpp core/io/concat_key_value_record_reader_test.cpp diff --git a/src/paimon/common/utils/concurrent_hash_map_test.cpp b/src/paimon/common/utils/concurrent_hash_map_test.cpp index 95a2f7fa9..249ddafeb 100644 --- a/src/paimon/common/utils/concurrent_hash_map_test.cpp +++ b/src/paimon/common/utils/concurrent_hash_map_test.cpp @@ -21,6 +21,7 @@ #include +#include #include #include #include @@ -31,6 +32,23 @@ namespace paimon::test { +class StatefulHashCompare { + public: + StatefulHashCompare() : seed_(next_seed_++) {} + + size_t hash(const int32_t& key) const { + return std::hash{}(key) + seed_; + } + + bool equal(const int32_t& lhs, const int32_t& rhs) const { + return lhs == rhs; + } + + private: + inline static std::atomic next_seed_ = 0; + size_t seed_; +}; + TEST(ConcurrentHashMapTest, TestSimple) { ConcurrentHashMap hash_map; ASSERT_EQ(hash_map.Find(10), std::nullopt); @@ -72,6 +90,14 @@ TEST(ConcurrentHashMapTest, TestVectorStringHashCompare) { ASSERT_EQ(hash_map.Size(), 4); } +TEST(ConcurrentHashMapTest, TestStatefulHashCompare) { + ConcurrentHashMap hash_map; + hash_map.Insert(1, "a"); + ASSERT_EQ(hash_map.Find(1).value(), "a"); + hash_map.Erase(1); + ASSERT_EQ(hash_map.Find(1), std::nullopt); +} + TEST(ConcurrentHashMapTest, TestMultiThreadInsertAndFindAndDelete) { int32_t map_size = 1000; auto insert_task = [&](ConcurrentHashMap& hash_map) { diff --git a/src/paimon/common/utils/std_concurrent_hash_map.h b/src/paimon/common/utils/std_concurrent_hash_map.h index b0ad856ca..bd91d7f22 100644 --- a/src/paimon/common/utils/std_concurrent_hash_map.h +++ b/src/paimon/common/utils/std_concurrent_hash_map.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -51,17 +52,29 @@ namespace detail { template class HashCompareHasher { public: + explicit HashCompareHasher(std::shared_ptr hash_compare) + : hash_compare_(std::move(hash_compare)) {} + size_t operator()(const Key& key) const { - return HashCompare{}.hash(key); + return hash_compare_->hash(key); } + + private: + std::shared_ptr hash_compare_; }; template class HashCompareEqual { public: + explicit HashCompareEqual(std::shared_ptr hash_compare) + : hash_compare_(std::move(hash_compare)) {} + bool operator()(const Key& lhs, const Key& rhs) const { - return HashCompare{}.equal(lhs, rhs); + return hash_compare_->equal(lhs, rhs); } + + private: + std::shared_ptr hash_compare_; }; } // namespace detail @@ -73,7 +86,10 @@ class ConcurrentHashMap { detail::HashCompareEqual>; public: - ConcurrentHashMap() = default; + ConcurrentHashMap() + : hash_compare_(std::make_shared()), + hash_map_(0, detail::HashCompareHasher(hash_compare_), + detail::HashCompareEqual(hash_compare_)) {} ~ConcurrentHashMap() = default; ConcurrentHashMap(const ConcurrentHashMap&) = delete; @@ -106,6 +122,7 @@ class ConcurrentHashMap { } private: + std::shared_ptr hash_compare_; HashMap hash_map_; mutable std::shared_mutex mutex_; }; diff --git a/src/paimon/core/io/async_key_value_producer_and_consumer.cpp b/src/paimon/core/io/async_key_value_producer_and_consumer.cpp index a609d8923..96e74f839 100644 --- a/src/paimon/core/io/async_key_value_producer_and_consumer.cpp +++ b/src/paimon/core/io/async_key_value_producer_and_consumer.cpp @@ -21,6 +21,7 @@ #include #include +#include #include #include "arrow/c/abi.h" @@ -143,14 +144,35 @@ template void AsyncKeyValueProducerAndConsumer::CleanUp() { consume_finished_ = true; next_batch_finished_ = true; + + // A producer or consumer may already be blocked in a bounded queue push. Keep draining both + // queues until every worker has stopped producing, otherwise a single drain can miss an item + // pushed by a worker after it is woken up. + while (true) { + bool producer_finished = + !producer_future_.valid() || + producer_future_.wait_for(std::chrono::microseconds(0)) == std::future_status::ready; + bool consumers_finished = true; + for (const std::unique_ptr>& consumer : consumers_) { + if (!consumer->IsFinished()) { + consumers_finished = false; + break; + } + } + if (producer_finished && consumers_finished) { + break; + } + CleanUpQueue(); + std::this_thread::yield(); + } CleanUpQueue(); + if (producer_future_.valid()) { [[maybe_unused]] Status status = producer_future_.get(); } for (auto& consumer : consumers_) { consumer->CleanUp(); } - CleanUpQueue(); } template @@ -202,6 +224,12 @@ Status AsyncKeyValueConsumer::GetStatus() const { return Status::OK(); } +template +bool AsyncKeyValueConsumer::IsFinished() const { + return !consumer_future_.valid() || + consumer_future_.wait_for(std::chrono::microseconds(0)) == std::future_status::ready; +} + template Status AsyncKeyValueConsumer::ConsumeLoop() { while (!consume_finished_) { diff --git a/src/paimon/core/io/async_key_value_producer_and_consumer.h b/src/paimon/core/io/async_key_value_producer_and_consumer.h index 99ae53d5b..bdfa8f0a0 100644 --- a/src/paimon/core/io/async_key_value_producer_and_consumer.h +++ b/src/paimon/core/io/async_key_value_producer_and_consumer.h @@ -115,6 +115,9 @@ class AsyncKeyValueConsumer { void CleanUp(); private: + friend class AsyncKeyValueProducerAndConsumer; + + bool IsFinished() const; Status ConsumeLoop(); private: diff --git a/src/paimon/core/io/async_key_value_producer_and_consumer_test.cpp b/src/paimon/core/io/async_key_value_producer_and_consumer_test.cpp new file mode 100644 index 000000000..fc2f7d340 --- /dev/null +++ b/src/paimon/core/io/async_key_value_producer_and_consumer_test.cpp @@ -0,0 +1,132 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/io/async_key_value_producer_and_consumer.h" + +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +class ConsumerBarrier { + public: + explicit ConsumerBarrier(size_t consumer_count) : consumer_count_(consumer_count) {} + + void Wait() { + std::unique_lock lock(mutex_); + ++arrived_count_; + if (arrived_count_ == consumer_count_) { + condition_.notify_all(); + } else { + condition_.wait(lock, [this]() { return arrived_count_ >= consumer_count_; }); + } + } + + private: + const size_t consumer_count_; + size_t arrived_count_ = 0; + std::mutex mutex_; + std::condition_variable condition_; +}; + +class TestKeyValueConsumer : public RowToArrowArrayConverter { + public: + explicit TestKeyValueConsumer(std::shared_ptr barrier) + : RowToArrowArrayConverter(/*reserve_count=*/0, std::vector(), nullptr, + nullptr), + barrier_(std::move(barrier)) {} + + Result NextBatch(const std::vector&) override { + barrier_->Wait(); + return KeyValueBatch(); + } + + private: + std::shared_ptr barrier_; +}; + +class TestSortMergeReader : public SortMergeReader { + public: + explicit TestSortMergeReader(size_t row_count) : rows_(row_count) {} + + class Iterator : public SortMergeReader::Iterator { + public: + explicit Iterator(TestSortMergeReader* reader) : reader_(reader) {} + + Result HasNext() override { + return reader_->next_row_ < reader_->rows_.size(); + } + + KeyValue&& Next() override { + return std::move(reader_->rows_[reader_->next_row_++]); + } + + private: + TestSortMergeReader* reader_; + }; + + Result> NextBatch() override { + if (iterator_created_) { + return std::unique_ptr(); + } + iterator_created_ = true; + return std::make_unique(this); + } + + void Close() override {} + + std::shared_ptr GetReaderMetrics() const override { + return nullptr; + } + + private: + std::vector rows_; + size_t next_row_ = 0; + bool iterator_created_ = false; +}; + +} // namespace + +TEST(AsyncKeyValueProducerAndConsumerTest, TestEarlyCloseWithBlockedConsumers) { + constexpr int32_t kConsumerCount = 32; + std::shared_ptr barrier = std::make_shared(kConsumerCount); + auto create_consumer = + [barrier]() -> Result>> { + std::unique_ptr> consumer = + std::make_unique(barrier); + return consumer; + }; + + AsyncKeyValueProducerAndConsumer producer_and_consumer( + std::make_unique(kConsumerCount * 4), create_consumer, + /*batch_size=*/1, kConsumerCount, /*pool=*/nullptr); + + ASSERT_OK_AND_ASSIGN(KeyValueBatch first_batch, producer_and_consumer.NextBatch()); + ASSERT_FALSE(first_batch.batch); + producer_and_consumer.Close(); +} + +} // namespace paimon::test