From 3a203e6674fb12cb68f5817cd6800f705d312bb9 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Tue, 17 Dec 2019 02:35:58 +0300
Subject: [PATCH] Add NetQueryDispatcher::set_main_dc_id.

GitOrigin-RevId: 80e660bbaffa0c00a2b85d9c0f73c543397aa3d0
---
 td/telegram/net/NetQueryDispatcher.cpp | 49 +++++++++++++++-----------
 td/telegram/net/NetQueryDispatcher.h   |  2 ++
 2 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/td/telegram/net/NetQueryDispatcher.cpp b/td/telegram/net/NetQueryDispatcher.cpp
index 3af0e41c..c0d76cba 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<int32>(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<std::mutex> 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<std::mutex> 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 8f78eb5c..c721e9e6 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<bool> stop_flag_{false};
   bool need_destroy_auth_key_{false};