diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 2ac9ad32d..2a140a3af 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -11,6 +11,7 @@ set(TD_TEST_SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/message_entities.cpp ${CMAKE_CURRENT_SOURCE_DIR}/mtproto.cpp ${CMAKE_CURRENT_SOURCE_DIR}/poll.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/query_merger.cpp ${CMAKE_CURRENT_SOURCE_DIR}/secret.cpp ${CMAKE_CURRENT_SOURCE_DIR}/secure_storage.cpp ${CMAKE_CURRENT_SOURCE_DIR}/set_with_position.cpp diff --git a/test/query_merger.cpp b/test/query_merger.cpp new file mode 100644 index 000000000..cdda2390f --- /dev/null +++ b/test/query_merger.cpp @@ -0,0 +1,96 @@ +// +// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2023 +// +// Distributed under the Boost Software License, Version 1.0. (See accompanying +// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt) +// +#include "td/telegram/QueryMerger.h" + +#include "td/actor/actor.h" +#include "td/actor/ConcurrentScheduler.h" +#include "td/actor/SleepActor.h" + +#include "td/utils/common.h" +#include "td/utils/FlatHashSet.h" +#include "td/utils/Random.h" +#include "td/utils/tests.h" + +#include + +class TestQueryMerger : public td::Actor { + void start_up() override { + query_merger_.set_merge_function([this](td::vector query_ids, td::Promise &&promise) { + ASSERT_TRUE(!query_ids.empty()); + ASSERT_EQ(query_ids.size(), td::min(next_query_ids_.size(), MAX_MERGED_QUERY_COUNT)); + for (auto query_id : query_ids) { + auto next_query_id = next_query_ids_.front(); + next_query_ids_.pop(); + ASSERT_EQ(query_id, next_query_id); + } + current_query_count_++; + ASSERT_TRUE(current_query_count_ <= MAX_CONCURRENT_QUERY_COUNT); + if (!next_query_ids_.empty()) { + ASSERT_EQ(current_query_count_, MAX_CONCURRENT_QUERY_COUNT); + } + td::create_actor("CompleteMergeQuery", 0.02, + td::PromiseCreator::lambda([this, query_ids, promise = std::move(promise)]( + td::Result result) mutable { + for (auto query_id : query_ids) { + LOG(INFO) << "Complete " << query_id; + bool is_erased = pending_query_ids_.erase(query_id); + ASSERT_TRUE(is_erased); + } + current_query_count_--; + promise.set_result(std::move(result)); + })) + .release(); + yield(); + }); + loop(); + } + + void loop() { + std::size_t query_count = 0; + std::size_t added_queries = td::Random::fast(1, 3); + while (query_count++ < added_queries && total_query_count_++ < MAX_QUERY_COUNT) { + td::int64 query_id = td::Random::fast(1, 20); + if (pending_query_ids_.insert(query_id).second) { + next_query_ids_.push(query_id); + } + query_merger_.add_query(query_id, + td::PromiseCreator::lambda([this, query_id](td::Result result) mutable { + completed_query_count_++; + if (completed_query_count_ == MAX_QUERY_COUNT) { + ASSERT_EQ(current_query_count_, 0u); + ASSERT_TRUE(next_query_ids_.empty()); + ASSERT_TRUE(pending_query_ids_.empty()); + td::Scheduler::instance()->finish(); + } else { + yield(); + } + })); + } + } + + static constexpr std::size_t MAX_CONCURRENT_QUERY_COUNT = 5; + static constexpr std::size_t MAX_MERGED_QUERY_COUNT = 3; + static constexpr std::size_t MAX_QUERY_COUNT = 1000; + + td::QueryMerger query_merger_{"QueryMerger", MAX_CONCURRENT_QUERY_COUNT, MAX_MERGED_QUERY_COUNT}; + std::size_t current_query_count_ = 0; + std::size_t total_query_count_ = 0; + std::size_t completed_query_count_ = 0; + + std::queue next_query_ids_; + td::FlatHashSet pending_query_ids_; +}; + +TEST(QueryMerger, stress) { + td::ConcurrentScheduler sched(0, 0); + sched.create_actor_unsafe(0, "TestQueryMerger").release(); + sched.start(); + while (sched.run_main(10)) { + // empty + } + sched.finish(); +}