Prevent unloading of replied polls.

This commit is contained in:
levlam 2023-11-06 20:37:17 +03:00
parent da461f5da2
commit baed5606fb
7 changed files with 76 additions and 3 deletions

View File

@ -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<const MessagePoll *>(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<const MessagePoll *>(content)->poll_id);
default:
return;
}
}
template <class ToT, class FromT>
static tl_object_ptr<ToT> secret_to_telegram(FromT &from);

View File

@ -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<MessageContent> get_secret_message_content(
Td *td, string message_text, unique_ptr<EncryptedFile> file,
tl_object_ptr<secret_api::DecryptedMessageMedia> &&media_ptr,

View File

@ -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;
}

View File

@ -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;
}

View File

@ -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<PollId, WaitFreeHashSet<MessageFullId, MessageFullIdHash>, PollIdHash> server_poll_messages_;
WaitFreeHashMap<PollId, WaitFreeHashSet<MessageFullId, MessageFullIdHash>, PollIdHash> other_poll_messages_;
WaitFreeHashMap<PollId, int32, PollIdHash> reply_poll_counts_;
struct PendingPollAnswer {
vector<string> options_;
vector<Promise<Unit>> promises_;

View File

@ -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_ &&

View File

@ -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 <class StorerT>
void store(StorerT &storer) const;