From baed5606fbcaad1f2b525512d6768ef4d6b2f612 Mon Sep 17 00:00:00 2001 From: levlam Date: Mon, 6 Nov 2023 20:37:17 +0300 Subject: [PATCH] Prevent unloading of replied polls. --- td/telegram/MessageContent.cpp | 18 ++++++++++++++++++ td/telegram/MessageContent.h | 4 ++++ td/telegram/MessagesManager.cpp | 5 +++++ td/telegram/PollManager.cpp | 30 +++++++++++++++++++++++++++--- td/telegram/PollManager.h | 6 ++++++ td/telegram/RepliedMessageInfo.cpp | 12 ++++++++++++ td/telegram/RepliedMessageInfo.h | 4 ++++ 7 files changed, 76 insertions(+), 3 deletions(-) diff --git a/td/telegram/MessageContent.cpp b/td/telegram/MessageContent.cpp index dd2eea716..a72878aa4 100644 --- a/td/telegram/MessageContent.cpp +++ b/td/telegram/MessageContent.cpp @@ -4961,6 +4961,24 @@ void unregister_message_content(Td *td, const MessageContent *content, MessageFu } } +void register_reply_message_content(Td *td, const MessageContent *content) { + switch (content->get_type()) { + case MessageContentType::Poll: + return td->poll_manager_->register_reply_poll(static_cast(content)->poll_id); + default: + return; + } +} + +void unregister_reply_message_content(Td *td, const MessageContent *content) { + switch (content->get_type()) { + case MessageContentType::Poll: + return td->poll_manager_->unregister_reply_poll(static_cast(content)->poll_id); + default: + return; + } +} + template static tl_object_ptr secret_to_telegram(FromT &from); diff --git a/td/telegram/MessageContent.h b/td/telegram/MessageContent.h index fb150ff74..d04248e3f 100644 --- a/td/telegram/MessageContent.h +++ b/td/telegram/MessageContent.h @@ -205,6 +205,10 @@ void reregister_message_content(Td *td, const MessageContent *old_content, const void unregister_message_content(Td *td, const MessageContent *content, MessageFullId message_full_id, const char *source); +void register_reply_message_content(Td *td, const MessageContent *content); + +void unregister_reply_message_content(Td *td, const MessageContent *content); + unique_ptr get_secret_message_content( Td *td, string message_text, unique_ptr file, tl_object_ptr &&media_ptr, diff --git a/td/telegram/MessagesManager.cpp b/td/telegram/MessagesManager.cpp index 1a471a72c..c946f342c 100644 --- a/td/telegram/MessagesManager.cpp +++ b/td/telegram/MessagesManager.cpp @@ -27356,6 +27356,8 @@ bool MessagesManager::can_register_message_reply(const Message *m) const { } void MessagesManager::register_message_reply(DialogId dialog_id, const Message *m) { + m->replied_message_info.register_content(td_); + if (!can_register_message_reply(m)) { return; } @@ -27378,6 +27380,7 @@ void MessagesManager::register_message_reply(DialogId dialog_id, const Message * } void MessagesManager::reregister_message_reply(DialogId dialog_id, const Message *m) { + // reply itself wan't changed, so there is nothing to reregister if (!can_register_message_reply(m)) { return; } @@ -27406,6 +27409,8 @@ void MessagesManager::reregister_message_reply(DialogId dialog_id, const Message } void MessagesManager::unregister_message_reply(DialogId dialog_id, const Message *m) { + m->replied_message_info.unregister_content(td_); + if (!can_register_message_reply(m)) { return; } diff --git a/td/telegram/PollManager.cpp b/td/telegram/PollManager.cpp index 0df725d41..0e39141e4 100644 --- a/td/telegram/PollManager.cpp +++ b/td/telegram/PollManager.cpp @@ -274,7 +274,8 @@ void PollManager::tear_down() { PollManager::~PollManager() { Scheduler::instance()->destroy_on_scheduler(G()->get_gc_scheduler_id(), polls_, server_poll_messages_, - other_poll_messages_, poll_voters_, loaded_from_database_polls_); + other_poll_messages_, reply_poll_counts_, poll_voters_, + loaded_from_database_polls_); } void PollManager::on_update_poll_timeout_callback(void *poll_manager_ptr, int64 poll_id_int) { @@ -725,13 +726,36 @@ void PollManager::unregister_poll(PollId poll_id, MessageFullId message_full_id, } } +void PollManager::register_reply_poll(PollId poll_id) { + CHECK(have_poll(poll_id)); + CHECK(!is_local_poll_id(poll_id)); + LOG(INFO) << "Register replied " << poll_id; + reply_poll_counts_[poll_id]++; + if (!G()->close_flag()) { + unload_poll_timeout_.cancel_timeout(poll_id.get()); + } +} + +void PollManager::unregister_reply_poll(PollId poll_id) { + CHECK(have_poll(poll_id)); + CHECK(!is_local_poll_id(poll_id)); + LOG(INFO) << "Unregister replied " << poll_id; + auto &count = reply_poll_counts_[poll_id]; + CHECK(count > 0); + count--; + if (count == 0) { + reply_poll_counts_.erase(poll_id); + schedule_poll_unload(poll_id); + } +} + bool PollManager::can_unload_poll(PollId poll_id) { if (G()->close_flag()) { return false; } if (is_local_poll_id(poll_id) || server_poll_messages_.count(poll_id) != 0 || - other_poll_messages_.count(poll_id) != 0 || pending_answers_.count(poll_id) != 0 || - being_closed_polls_.count(poll_id) != 0) { + other_poll_messages_.count(poll_id) != 0 || reply_poll_counts_.count(poll_id) != 0 || + pending_answers_.count(poll_id) != 0 || being_closed_polls_.count(poll_id) != 0) { return false; } diff --git a/td/telegram/PollManager.h b/td/telegram/PollManager.h index 751623b98..356ddbd00 100644 --- a/td/telegram/PollManager.h +++ b/td/telegram/PollManager.h @@ -57,6 +57,10 @@ class PollManager final : public Actor { void unregister_poll(PollId poll_id, MessageFullId message_full_id, const char *source); + void register_reply_poll(PollId poll_id); + + void unregister_reply_poll(PollId poll_id); + bool get_poll_is_closed(PollId poll_id) const; bool get_poll_is_anonymous(PollId poll_id) const; @@ -235,6 +239,8 @@ class PollManager final : public Actor { WaitFreeHashMap, PollIdHash> server_poll_messages_; WaitFreeHashMap, PollIdHash> other_poll_messages_; + WaitFreeHashMap reply_poll_counts_; + struct PendingPollAnswer { vector options_; vector> promises_; diff --git a/td/telegram/RepliedMessageInfo.cpp b/td/telegram/RepliedMessageInfo.cpp index a6702df4d..24d5deb64 100644 --- a/td/telegram/RepliedMessageInfo.cpp +++ b/td/telegram/RepliedMessageInfo.cpp @@ -377,6 +377,18 @@ MessageFullId RepliedMessageInfo::get_reply_message_full_id(DialogId owner_dialo return {dialog_id_.is_valid() ? dialog_id_ : owner_dialog_id, message_id_}; } +void RepliedMessageInfo::register_content(Td *td) const { + if (content_ != nullptr) { + register_reply_message_content(td, content_.get()); + } +} + +void RepliedMessageInfo::unregister_content(Td *td) const { + if (content_ != nullptr) { + unregister_reply_message_content(td, content_.get()); + } +} + bool operator==(const RepliedMessageInfo &lhs, const RepliedMessageInfo &rhs) { if (!(lhs.message_id_ == rhs.message_id_ && lhs.dialog_id_ == rhs.dialog_id_ && lhs.origin_date_ == rhs.origin_date_ && lhs.origin_ == rhs.origin_ && lhs.quote_ == rhs.quote_ && diff --git a/td/telegram/RepliedMessageInfo.h b/td/telegram/RepliedMessageInfo.h index 49231d976..776a67a5e 100644 --- a/td/telegram/RepliedMessageInfo.h +++ b/td/telegram/RepliedMessageInfo.h @@ -102,6 +102,10 @@ class RepliedMessageInfo { MessageFullId get_reply_message_full_id(DialogId owner_dialog_id) const; + void register_content(Td *td) const; + + void unregister_content(Td *td) const; + template void store(StorerT &storer) const;