diff --git a/td/telegram/net/NetQueryDispatcher.cpp b/td/telegram/net/NetQueryDispatcher.cpp index 3af0e41cf..c0d76cba4 100644 --- a/td/telegram/net/NetQueryDispatcher.cpp +++ b/td/telegram/net/NetQueryDispatcher.cpp @@ -296,26 +296,7 @@ void NetQueryDispatcher::try_fix_migrate(NetQueryPtr &net_query) { for (auto &prefix : prefixes) { if (msg.substr(0, prefix.size()) == prefix) { int32 new_main_dc_id = to_integer(msg.substr(prefix.size())); - if (!DcId::is_valid(new_main_dc_id)) { - LOG(FATAL) << "Receive " << prefix << " to wrong dc " << new_main_dc_id; - } - if (new_main_dc_id != main_dc_id_.load(std::memory_order_relaxed)) { - // Very rare event. Mutex is ok. - std::lock_guard guard(main_dc_id_mutex_); - if (new_main_dc_id != main_dc_id_) { - LOG(INFO) << "Update: " << tag("main_dc_id", main_dc_id_.load(std::memory_order_relaxed)); - if (is_dc_inited(main_dc_id_.load(std::memory_order_relaxed))) { - send_closure_later(dcs_[main_dc_id_ - 1].main_session_, &SessionMultiProxy::update_main_flag, false); - } - main_dc_id_ = new_main_dc_id; - if (is_dc_inited(main_dc_id_.load(std::memory_order_relaxed))) { - send_closure_later(dcs_[main_dc_id_ - 1].main_session_, &SessionMultiProxy::update_main_flag, true); - } - send_closure_later(dc_auth_manager_, &DcAuthManager::update_main_dc, - DcId::internal(main_dc_id_.load(std::memory_order_relaxed))); - G()->td_db()->get_binlog_pmc()->set("main_dc_id", to_string(main_dc_id_.load(std::memory_order_relaxed))); - } - } + set_main_dc_id(new_main_dc_id); if (!net_query->dc_id().is_main()) { LOG(ERROR) << msg << " from query to non-main dc " << net_query->dc_id(); @@ -328,4 +309,32 @@ void NetQueryDispatcher::try_fix_migrate(NetQueryPtr &net_query) { } } +void NetQueryDispatcher::set_main_dc_id(int32 new_main_dc_id) { + if (!DcId::is_valid(new_main_dc_id)) { + LOG(ERROR) << "Receive wrong DC " << new_main_dc_id; + return; + } + if (new_main_dc_id == main_dc_id_.load(std::memory_order_relaxed)) { + return; + } + + // Very rare event. Mutex is ok. + std::lock_guard guard(main_dc_id_mutex_); + if (new_main_dc_id == main_dc_id_) { + return; + } + + LOG(INFO) << "Update: " << tag("main_dc_id", main_dc_id_.load(std::memory_order_relaxed)); + if (is_dc_inited(main_dc_id_.load(std::memory_order_relaxed))) { + send_closure_later(dcs_[main_dc_id_ - 1].main_session_, &SessionMultiProxy::update_main_flag, false); + } + main_dc_id_ = new_main_dc_id; + if (is_dc_inited(main_dc_id_.load(std::memory_order_relaxed))) { + send_closure_later(dcs_[main_dc_id_ - 1].main_session_, &SessionMultiProxy::update_main_flag, true); + } + send_closure_later(dc_auth_manager_, &DcAuthManager::update_main_dc, + DcId::internal(main_dc_id_.load(std::memory_order_relaxed))); + G()->td_db()->get_binlog_pmc()->set("main_dc_id", to_string(main_dc_id_.load(std::memory_order_relaxed))); +} + } // namespace td diff --git a/td/telegram/net/NetQueryDispatcher.h b/td/telegram/net/NetQueryDispatcher.h index 8f78eb5c6..c721e9e6f 100644 --- a/td/telegram/net/NetQueryDispatcher.h +++ b/td/telegram/net/NetQueryDispatcher.h @@ -56,6 +56,8 @@ class NetQueryDispatcher { return DcId::internal(main_dc_id_.load()); } + void set_main_dc_id(int32 new_main_dc_id); + private: std::atomic stop_flag_{false}; bool need_destroy_auth_key_{false};