Teach DownloadState how to download split states

DanShaders · DanShaders · commit 3a93b225e288 · 2025-06-20T09:37:55.000-04:00
This way, we propagate hardcoding UnsplitStateType one layer up.
diff --git a/validator/full-node-shard.cpp b/validator/full-node-shard.cpp
@@ -990,8 +990,8 @@ void FullNodeShardImpl::download_persistent_state(BlockIdExt id, BlockIdExt mast
                                                   td::Timestamp timeout, td::Promise<td::BufferSlice> promise) {
   auto &b = choose_neighbour();
   td::actor::create_actor<DownloadState>(PSTRING() << "downloadstatereq" << id.id.to_str(), id, masterchain_block_id,
-                                         adnl_id_, overlay_id_, b.adnl_id, priority, timeout, validator_manager_,
-                                         rldp2_, overlays_, adnl_, client_, std::move(promise))
+                                         UnsplitStateType{}, adnl_id_, overlay_id_, b.adnl_id, priority, timeout,
+                                         validator_manager_, rldp2_, overlays_, adnl_, client_, std::move(promise))
       .release();
 }
 
diff --git a/validator/interfaces/persistent-state.h b/validator/interfaces/persistent-state.h
@@ -61,6 +61,14 @@ auto persistent_state_from_v2_query(auto const &query) {
   return std::tuple{block, mc_block, PersistentStateType{SplitAccountStateType{effective_shard}}};
 }
 
+inline ShardId persistent_state_to_effective_shard(ShardIdFull const &shard, PersistentStateType const &type) {
+  ShardId result = 0;
+  type.visit(td::overloaded([](UnsplitStateType) {},
+                            [&](SplitAccountStateType type) { result = type.effective_shard_id; },
+                            [&](SplitPersistentStateType) { result = shard.shard; }));
+  return result;
+}
+
 inline std::string persistent_state_type_to_string(ShardIdFull const &shard, PersistentStateType const &state) {
   std::string result;
   state.visit(td::overloaded([&](UnsplitStateType) { result = "unsplit"; },
diff --git a/validator/net/download-state.cpp b/validator/net/download-state.cpp
@@ -28,15 +28,17 @@ namespace validator {
 
 namespace fullnode {
 
-DownloadState::DownloadState(BlockIdExt block_id, BlockIdExt masterchain_block_id, adnl::AdnlNodeIdShort local_id,
-                             overlay::OverlayIdShort overlay_id, adnl::AdnlNodeIdShort download_from,
-                             td::uint32 priority, td::Timestamp timeout,
+DownloadState::DownloadState(BlockIdExt block_id, BlockIdExt masterchain_block_id, PersistentStateType type,
+                             adnl::AdnlNodeIdShort local_id, overlay::OverlayIdShort overlay_id,
+                             adnl::AdnlNodeIdShort download_from, td::uint32 priority, td::Timestamp timeout,
                              td::actor::ActorId<ValidatorManagerInterface> validator_manager,
                              td::actor::ActorId<adnl::AdnlSenderInterface> rldp,
                              td::actor::ActorId<overlay::Overlays> overlays, td::actor::ActorId<adnl::Adnl> adnl,
                              td::actor::ActorId<adnl::AdnlExtClient> client, td::Promise<td::BufferSlice> promise)
     : block_id_(block_id)
     , masterchain_block_id_(masterchain_block_id)
+    , type_(type)
+    , effective_shard_(persistent_state_to_effective_shard(block_id_.shard_full(), type))
     , local_id_(local_id)
     , overlay_id_(overlay_id)
     , download_from_(download_from)
@@ -48,6 +50,7 @@ DownloadState::DownloadState(BlockIdExt block_id, BlockIdExt masterchain_block_i
     , adnl_(adnl)
     , client_(client)
     , promise_(std::move(promise)) {
+  CHECK(masterchain_block_id_.is_valid() || effective_shard_ == 0);
 }
 
 void DownloadState::abort_query(td::Status reason) {
@@ -74,7 +77,7 @@ void DownloadState::start_up() {
   alarm_timestamp() = timeout_;
 
   td::actor::send_closure(validator_manager_, &ValidatorManagerInterface::get_persistent_state, block_id_,
-                          masterchain_block_id_, UnsplitStateType{},
+                          masterchain_block_id_, type_,
                           [SelfId = actor_id(this), block_id = block_id_](td::Result<td::BufferSlice> R) {
                             if (R.is_error()) {
                               td::actor::send_closure(SelfId, &DownloadState::get_block_handle);
@@ -124,22 +127,39 @@ void DownloadState::got_block_handle(BlockHandle handle) {
 
 void DownloadState::got_node_to_download(adnl::AdnlNodeIdShort node) {
   download_from_ = node;
-  LOG(WARNING) << "downloading state " << block_id_.to_str() << " from " << download_from_;
+  LOG(WARNING) << "downloading state " << block_id_.to_str() << " ("
+               << persistent_state_type_to_string(block_id_.shard_full(), type_) << ") from " << download_from_;
 
-  auto P = td::PromiseCreator::lambda([SelfId = actor_id(this)](td::Result<td::BufferSlice> R) mutable {
-    if (R.is_error()) {
-      td::actor::send_closure(SelfId, &DownloadState::abort_query, R.move_as_error());
+  td::Promise<td::BufferSlice> P;
+  td::BufferSlice query;
+
+  if (effective_shard_ == 0) {
+    P = td::PromiseCreator::lambda([SelfId = actor_id(this)](td::Result<td::BufferSlice> R) mutable {
+      if (R.is_error()) {
+        td::actor::send_closure(SelfId, &DownloadState::abort_query, R.move_as_error());
+      } else {
+        td::actor::send_closure(SelfId, &DownloadState::got_block_state_description, R.move_as_ok());
+      }
+    });
+
+    if (masterchain_block_id_.is_valid()) {
+      query = create_serialize_tl_object<ton_api::tonNode_preparePersistentState>(
+          create_tl_block_id(block_id_), create_tl_block_id(masterchain_block_id_));
     } else {
-      td::actor::send_closure(SelfId, &DownloadState::got_block_state_description, R.move_as_ok());
+      query = create_serialize_tl_object<ton_api::tonNode_prepareZeroState>(create_tl_block_id(block_id_));
     }
-  });
-
-  td::BufferSlice query;
-  if (masterchain_block_id_.is_valid()) {
-    query = create_serialize_tl_object<ton_api::tonNode_preparePersistentState>(
-        create_tl_block_id(block_id_), create_tl_block_id(masterchain_block_id_));
   } else {
-    query = create_serialize_tl_object<ton_api::tonNode_prepareZeroState>(create_tl_block_id(block_id_));
+    P = td::PromiseCreator::lambda([SelfId = actor_id(this)](td::Result<td::BufferSlice> R) mutable {
+      if (R.is_error()) {
+        td::actor::send_closure(SelfId, &DownloadState::abort_query, R.move_as_error());
+      } else {
+        td::actor::send_closure(SelfId, &DownloadState::got_state_size, R.move_as_ok());
+      }
+    });
+
+    query = create_serialize_tl_object<ton_api::tonNode_getPersistentStateSizeV2>(
+        create_tl_object<ton_api::tonNode_persistentStateIdV2>(
+            create_tl_block_id(block_id_), create_tl_block_id(masterchain_block_id_), effective_shard_));
   }
 
   if (client_.empty()) {
@@ -195,6 +215,25 @@ void DownloadState::got_block_state_description(td::BufferSlice data) {
           }));
 }
 
+void DownloadState::got_state_size(td::BufferSlice size_or_not_found) {
+  auto F = fetch_tl_object<ton_api::tonNode_PersistentStateSize>(std::move(size_or_not_found), true);
+  if (F.is_error()) {
+    abort_query(F.move_as_error());
+    return;
+  }
+  prev_logged_timer_ = td::Timer();
+
+  ton_api::downcast_call(*F.move_as_ok().get(),
+                         td::overloaded(
+                             [&](ton_api::tonNode_persistentStateSizeNotFound &f) {
+                               abort_query(td::Status::Error(ErrorCode::notready, "state not found"));
+                             },
+                             [&](ton_api::tonNode_persistentStateSize &f) {
+                               total_size_ = f.size_;
+                               got_block_state_part(td::BufferSlice{}, 0);
+                             }));
+}
+
 void DownloadState::request_total_size() {
   auto P = td::PromiseCreator::lambda([SelfId = actor_id(this)](td::Result<td::BufferSlice> R) {
     if (R.is_error()) {
@@ -207,8 +246,15 @@ void DownloadState::request_total_size() {
     td::actor::send_closure(SelfId, &DownloadState::got_total_size, res.ok()->size_);
   });
 
-  td::BufferSlice query = create_serialize_tl_object<ton_api::tonNode_getPersistentStateSize>(
-      create_tl_block_id(block_id_), create_tl_block_id(masterchain_block_id_));
+  td::BufferSlice query;
+  if (effective_shard_ == 0) {
+    query = create_serialize_tl_object<ton_api::tonNode_getPersistentStateSize>(
+        create_tl_block_id(block_id_), create_tl_block_id(masterchain_block_id_));
+  } else {
+    query = create_serialize_tl_object<ton_api::tonNode_getPersistentStateSizeV2>(
+        create_tl_object<ton_api::tonNode_persistentStateIdV2>(
+            create_tl_block_id(block_id_), create_tl_block_id(masterchain_block_id_), effective_shard_));
+  }
   if (client_.empty()) {
     td::actor::send_closure(overlays_, &overlay::Overlays::send_query_via, download_from_, local_id_, overlay_id_,
                             "get size", std::move(P), td::Timestamp::in(3.0), std::move(query),
@@ -275,8 +321,16 @@ void DownloadState::got_block_state_part(td::BufferSlice data, td::uint32 reques
     }
   });
 
-  td::BufferSlice query = create_serialize_tl_object<ton_api::tonNode_downloadPersistentStateSlice>(
-      create_tl_block_id(block_id_), create_tl_block_id(masterchain_block_id_), sum_, part_size);
+  td::BufferSlice query;
+  if (effective_shard_ == 0) {
+    query = create_serialize_tl_object<ton_api::tonNode_downloadPersistentStateSlice>(
+        create_tl_block_id(block_id_), create_tl_block_id(masterchain_block_id_), sum_, part_size);
+  } else {
+    query = create_serialize_tl_object<ton_api::tonNode_downloadPersistentStateSliceV2>(
+        create_tl_object<ton_api::tonNode_persistentStateIdV2>(
+            create_tl_block_id(block_id_), create_tl_block_id(masterchain_block_id_), effective_shard_),
+        sum_, part_size);
+  }
   if (client_.empty()) {
     td::actor::send_closure(overlays_, &overlay::Overlays::send_query_via, download_from_, local_id_, overlay_id_,
                             "download state", std::move(P), td::Timestamp::in(20.0), std::move(query),
diff --git a/validator/net/download-state.hpp b/validator/net/download-state.hpp
@@ -33,9 +33,10 @@ namespace fullnode {
 
 class DownloadState : public td::actor::Actor {
  public:
-  DownloadState(BlockIdExt block_id, BlockIdExt masterchain_block_id, adnl::AdnlNodeIdShort local_id,
-                overlay::OverlayIdShort overlay_id, adnl::AdnlNodeIdShort download_from, td::uint32 priority,
-                td::Timestamp timeout, td::actor::ActorId<ValidatorManagerInterface> validator_manager,
+  DownloadState(BlockIdExt block_id, BlockIdExt masterchain_block_id, PersistentStateType type,
+                adnl::AdnlNodeIdShort local_id, overlay::OverlayIdShort overlay_id, adnl::AdnlNodeIdShort download_from,
+                td::uint32 priority, td::Timestamp timeout,
+                td::actor::ActorId<ValidatorManagerInterface> validator_manager,
                 td::actor::ActorId<adnl::AdnlSenderInterface> rldp, td::actor::ActorId<overlay::Overlays> overlays,
                 td::actor::ActorId<adnl::Adnl> adnl, td::actor::ActorId<adnl::AdnlExtClient> client,
                 td::Promise<td::BufferSlice> promise);
@@ -49,6 +50,7 @@ class DownloadState : public td::actor::Actor {
   void got_block_handle(BlockHandle handle);
   void got_node_to_download(adnl::AdnlNodeIdShort node);
   void got_block_state_description(td::BufferSlice data_description);
+  void got_state_size(td::BufferSlice size_or_not_found);
   void request_total_size();
   void got_total_size(td::uint64 size);
   void got_block_state_part(td::BufferSlice data, td::uint32 requested_size);
@@ -57,6 +59,8 @@ class DownloadState : public td::actor::Actor {
  private:
   BlockIdExt block_id_;
   BlockIdExt masterchain_block_id_;
+  PersistentStateType type_;
+  ShardId effective_shard_;
   adnl::AdnlNodeIdShort local_id_;
   overlay::OverlayIdShort overlay_id_;
 
