diff --git a/td/telegram/PollManager.cpp b/td/telegram/PollManager.cpp index 2562b7e83..4d969eb72 100644 --- a/td/telegram/PollManager.cpp +++ b/td/telegram/PollManager.cpp @@ -329,18 +329,16 @@ void PollManager::notify_on_poll_update(PollId poll_id) { return; } - auto server_it = server_poll_messages_.find(poll_id); - if (server_it != server_poll_messages_.end()) { - for (const auto &full_message_id : server_it->second) { + if (server_poll_messages_.count(poll_id) > 0) { + server_poll_messages_[poll_id].foreach([&](const FullMessageId &full_message_id) { td_->messages_manager_->on_external_update_message_content(full_message_id); - } + }); } - auto other_it = other_poll_messages_.find(poll_id); - if (other_it != other_poll_messages_.end()) { - for (const auto &full_message_id : other_it->second) { + if (other_poll_messages_.count(poll_id) > 0) { + other_poll_messages_[poll_id].foreach([&](const FullMessageId &full_message_id) { td_->messages_manager_->on_external_update_message_content(full_message_id); - } + }); } } @@ -647,14 +645,12 @@ PollId PollManager::create_poll(string &&question, vector &&options, boo void PollManager::register_poll(PollId poll_id, FullMessageId full_message_id, const char *source) { CHECK(have_poll(poll_id)); if (full_message_id.get_message_id().is_scheduled() || !full_message_id.get_message_id().is_server()) { - bool is_inserted = other_poll_messages_[poll_id].insert(full_message_id).second; - LOG_CHECK(is_inserted) << source << ' ' << poll_id << ' ' << full_message_id; + other_poll_messages_[poll_id].insert(full_message_id); unload_poll_timeout_.cancel_timeout(poll_id.get()); return; } LOG(INFO) << "Register " << poll_id << " from " << full_message_id << " from " << source; - bool is_inserted = server_poll_messages_[poll_id].insert(full_message_id).second; - LOG_CHECK(is_inserted) << source << ' ' << poll_id << ' ' << full_message_id; + server_poll_messages_[poll_id].insert(full_message_id); auto poll = get_poll(poll_id); CHECK(poll != nullptr); if (!td_->auth_manager_->is_bot() && !is_local_poll_id(poll_id) && @@ -1265,12 +1261,11 @@ void PollManager::on_update_poll_timeout(PollId poll_id) { return; } - auto it = server_poll_messages_.find(poll_id); - if (it == server_poll_messages_.end()) { + if (server_poll_messages_.count(poll_id) == 0) { return; } - auto full_message_id = *it->second.begin(); + auto full_message_id = server_poll_messages_[poll_id].get_random(); LOG(INFO) << "Fetching results of " << poll_id << " from " << full_message_id; auto query_promise = PromiseCreator::lambda([poll_id, generation = current_generation_, actor_id = actor_id(this)]( Result> &&result) { @@ -1379,14 +1374,13 @@ void PollManager::on_online() { return; } - for (auto &it : server_poll_messages_) { - auto poll_id = it.first; + server_poll_messages_.foreach([&](const PollId &poll_id, WaitFreeHashSet &) { if (update_poll_timeout_.has_timeout(poll_id.get())) { auto timeout = Random::fast(3, 30); LOG(INFO) << "Schedule updating of " << poll_id << " in " << timeout; update_poll_timeout_.set_timeout_in(poll_id.get(), timeout); } - } + }); } PollId PollManager::dup_poll(PollId poll_id) { diff --git a/td/telegram/PollManager.h b/td/telegram/PollManager.h index 9b318bec6..fb38dfaac 100644 --- a/td/telegram/PollManager.h +++ b/td/telegram/PollManager.h @@ -25,6 +25,7 @@ #include "td/utils/Promise.h" #include "td/utils/Status.h" #include "td/utils/WaitFreeHashMap.h" +#include "td/utils/WaitFreeHashSet.h" #include @@ -227,8 +228,8 @@ class PollManager final : public Actor { ActorShared<> parent_; WaitFreeHashMap, PollIdHash> polls_; - FlatHashMap, PollIdHash> server_poll_messages_; - FlatHashMap, PollIdHash> other_poll_messages_; + WaitFreeHashMap, PollIdHash> server_poll_messages_; + WaitFreeHashMap, PollIdHash> other_poll_messages_; struct PendingPollAnswer { vector options_; diff --git a/tdutils/td/utils/WaitFreeHashSet.h b/tdutils/td/utils/WaitFreeHashSet.h index 6573a6d80..c7e6ad808 100644 --- a/tdutils/td/utils/WaitFreeHashSet.h +++ b/tdutils/td/utils/WaitFreeHashSet.h @@ -96,6 +96,22 @@ class WaitFreeHashSet { } } + KeyT get_random() const { + if (wait_free_storage_ != nullptr) { + for (size_t i = 0; i < MAX_STORAGE_COUNT; i++) { + if (!wait_free_storage_->sets_[i].empty()) { + return wait_free_storage_->sets_[i].get_random(); + } + } + // no need to explicitly return KeyT() + } + + if (default_set_.empty()) { + return KeyT(); + } + return *default_set_.begin(); + } + size_t calc_size() const { if (wait_free_storage_ == nullptr) { return default_set_.size();