Storage stat cache (#1724)

Co-authored-by: SpyCheese <[email protected]>

Author: EmelyanenkoK

crypto/block/account-storage-stat.h

@@ -57,6 +57,10 @@ class AccountStorageStat {
     return root.is_null() ? td::Bits256::zero() : td::Bits256{root->get_hash().bits()};
   }
 
+  bool is_dict_ready() const {
+    return dict_up_to_date_;
+  }
+
   void apply_child_stat(AccountStorageStat &&child);
 
   static constexpr int errorcode_limits_exceeded = 999;

crypto/block/transaction.cpp

@@ -561,6 +561,44 @@ bool Account::init_new(ton::UnixTime now) {
   return true;
 }
 
+/**
+ * Initializes account_storage_stat of the account using the existing dict_root.
+ * This is not strictly necessary, as the storage stat is recomputed in Transaction.
+ * However, it can be used to optimize cell usage.
+ * This requires storage_dict_hash to be set, as it guarantees that the stored storage_used was computed recently
+ * (in older versions it included extra currency balance, in newer versions it does not).
+ *
+ * @param dict_root Root of the storage dictionary.
+ *
+ * @returns Status of the operation.
+ */
+td::Status Account::init_account_storage_stat(Ref<vm::Cell> dict_root) {
+  if (storage.is_null()) {
+    if (dict_root.not_null()) {
+      return td::Status::Error("storage is null, but dict_root is not null");
+    }
+    account_storage_stat = {};
+    return td::Status::OK();
+  }
+  if (!storage_dict_hash) {
+    return td::Status::Error("cannot init storage dict: storage_dict_hash is not set");
+  }
+  // Root of AccountStorage is not counted in AccountStorageStat
+  if (storage_used.cells < 1 || storage_used.bits < storage->size()) {
+    return td::Status::Error(PSTRING() << "storage_used is too small: cells=" << storage_used.cells
+                                       << " bits=" << storage_used.bits << " storage_root_bits=" << storage->size());
+  }
+  AccountStorageStat new_stat(std::move(dict_root), storage->prefetch_all_refs(), storage_used.cells - 1,
+                              storage_used.bits - storage->size());
+  TRY_RESULT(root_hash, new_stat.get_dict_hash());
+  if (storage_dict_hash.value() != root_hash) {
+    return td::Status::Error(PSTRING() << "invalid storage dict hash: computed " << root_hash.to_hex() << ", found "
+                                       << storage_dict_hash.value().to_hex());
+  }
+  account_storage_stat = std::move(new_stat);
+  return td::Status::OK();
+}
+
 /**
  * Resets the fixed prefix length of the account.
  *

crypto/block/transaction.h

@@ -298,6 +298,7 @@ struct Account {
   bool set_address(ton::WorkchainId wc, td::ConstBitPtr new_addr);
   bool unpack(Ref<vm::CellSlice> account, ton::UnixTime now, bool special);
   bool init_new(ton::UnixTime now);
+  td::Status init_account_storage_stat(Ref<vm::Cell> dict_root);
   bool deactivate();
   bool recompute_tmp_addr(Ref<vm::CellSlice>& tmp_addr, int fixed_prefix_length, td::ConstBitPtr orig_addr_rewrite) const;
   td::RefInt256 compute_storage_fees(ton::UnixTime now, const std::vector<block::StoragePrices>& pricing) const;

tdutils/td/utils/LRUCache.h

@@ -26,8 +26,7 @@ namespace td {
 template <typename K, typename V>
 class LRUCache {
  public:
-  explicit LRUCache(size_t max_size) : max_size_(max_size) {
-    CHECK(max_size_ > 0);
+  explicit LRUCache(uint64 max_size) : max_size_(max_size) {
   }
   LRUCache(const LRUCache&) = delete;
   LRUCache& operator=(const LRUCache&) = delete;
@@ -49,13 +48,14 @@ class LRUCache {
     return cache_.contains(key);
   }
 
-  bool put(const K& key, V value, bool update = true) {
+  bool put(const K& key, V value, bool update = true, uint64 weight = 1) {
     bool added = false;
     auto it = cache_.find(key);
     if (it == cache_.end()) {
       update = true;
-      it = cache_.insert(std::make_unique<Entry>(key, std::move(value))).first;
+      it = cache_.insert(std::make_unique<Entry>(key, std::move(value), weight)).first;
       added = true;
+      total_weight_ += weight;
     } else {
       (*it)->value = std::move(value);
       if (update) {
@@ -69,11 +69,12 @@ class LRUCache {
     return added;
   }
 
-  V& get(const K& key, bool update = true) {
+  V& get(const K& key, bool update = true, uint64 weight = 1) {
     auto it = cache_.find(key);
     if (it == cache_.end()) {
       update = true;
-      it = cache_.insert(std::make_unique<Entry>(key)).first;
+      it = cache_.insert(std::make_unique<Entry>(key, weight)).first;
+      total_weight_ += weight;
     } else if (update) {
       (*it)->remove();
     }
@@ -87,12 +88,13 @@ class LRUCache {
 
  private:
   struct Entry : ListNode {
-    explicit Entry(K key) : key(std::move(key)) {
+    Entry(K key, uint64 weight) : key(std::move(key)), weight(weight) {
     }
-    Entry(K key, V value) : key(std::move(key)), value(std::move(value)) {
+    Entry(K key, V value, uint64 weight) : key(std::move(key)), value(std::move(value)), weight(weight) {
     }
     K key;
     V value;
+    uint64 weight;
   };
   struct Cmp {
     using is_transparent = void;
@@ -108,13 +110,15 @@ class LRUCache {
   };
   std::set<std::unique_ptr<Entry>, Cmp> cache_;
   ListNode lru_;
-  size_t max_size_;
+  uint64 max_size_;
+  uint64 total_weight_ = 0;
 
   void cleanup() {
-    while (cache_.size() > max_size_) {
+    while (total_weight_ > max_size_ && cache_.size() > 1) {
       auto to_remove = (Entry*)lru_.get();
       CHECK(to_remove);
       to_remove->remove();
+      total_weight_ -= to_remove->weight;
       cache_.erase(cache_.find(to_remove->key));
     }
   }

validator/CMakeLists.txt

@@ -59,6 +59,7 @@ set(VALIDATOR_HEADERS
   import-db-slice-local.hpp
   queue-size-counter.hpp
   validator-telemetry.hpp
+  storage-stat-cache.hpp
 
   manager-disk.h
   manager-disk.hpp
@@ -87,6 +88,7 @@ set(VALIDATOR_SOURCE
   validator-options.cpp
   queue-size-counter.cpp
   validator-telemetry.cpp
+  storage-stat-cache.cpp
 
   downloaders/wait-block-data.cpp
   downloaders/wait-block-state.cpp

validator/impl/collator-impl.h

@@ -227,6 +227,9 @@ class Collator final : public td::actor::Actor {
   bool deferring_messages_enabled_ = false;
   bool store_out_msg_queue_size_ = false;
 
+  std::function<td::Ref<vm::Cell>(const td::Bits256&)> storage_stat_cache_;
+  std::vector<std::pair<td::Ref<vm::Cell>, td::uint32>> storage_stat_cache_update_;
+
   td::PerfWarningTimer perf_timer_;
   //
   block::Account* lookup_account(td::ConstBitPtr addr) const;
@@ -247,6 +250,7 @@ class Collator final : public td::actor::Actor {
   void after_get_shard_state(int idx, td::Result<Ref<ShardState>> res);
   void after_get_block_data(int idx, td::Result<Ref<BlockData>> res);
   void after_get_shard_blocks(td::Result<std::vector<Ref<ShardTopBlockDescription>>> res);
+  void after_get_storage_stat_cache(td::Result<std::function<td::Ref<vm::Cell>(const td::Bits256&)>> res);
   bool preprocess_prev_mc_state();
   bool register_mc_state(Ref<MasterchainStateQ> other_mc_state);
   bool request_aux_mc_state(BlockSeqno seqno, Ref<MasterchainStateQ>& state);

validator/impl/collator.cpp

@@ -277,7 +277,16 @@ void Collator::start_up() {
           td::actor::send_closure_later(std::move(self), &Collator::after_get_shard_blocks, std::move(res));
         });
   }
-  // 6. set timeout
+  // 6. get storage stat cache
+  ++pending;
+  LOG(DEBUG) << "sending get_storage_stat_cache() query to Manager";
+  td::actor::send_closure_later(
+      manager, &ValidatorManager::get_storage_stat_cache,
+      [self = get_self()](td::Result<std::function<td::Ref<vm::Cell>(const td::Bits256&)>> res) {
+        LOG(DEBUG) << "got answer to get_storage_stat_cache() query";
+        td::actor::send_closure_later(std::move(self), &Collator::after_get_storage_stat_cache, std::move(res));
+      });
+  // 7. set timeout
   alarm_timestamp() = timeout;
   CHECK(pending);
 }
@@ -686,6 +695,22 @@ void Collator::after_get_shard_blocks(td::Result<std::vector<Ref<ShardTopBlockDe
   check_pending();
 }
 
+/**
+ * Callback function called after retrieving storage stat cache.
+ *
+ * @param res The retrieved storage stat cache.
+ */
+void Collator::after_get_storage_stat_cache(td::Result<std::function<td::Ref<vm::Cell>(const td::Bits256&)>> res) {
+  --pending;
+  if (res.is_error()) {
+    LOG(INFO) << "after_get_storage_stat_cache : " << res.error();
+  } else {
+    LOG(DEBUG) << "after_get_storage_stat_cache : OK";
+    storage_stat_cache_ = res.move_as_ok();
+  }
+  check_pending();
+}
+
 /**
  * Unpacks the last masterchain state and initializes the Collator object with the extracted configuration.
  *
@@ -2448,6 +2473,20 @@ std::unique_ptr<block::Account> Collator::make_account_from(td::ConstBitPtr addr
     return nullptr;
   }
   ptr->block_lt = start_lt;
+  if (storage_stat_cache_ && ptr->storage_dict_hash) {
+    auto dict_root = storage_stat_cache_(ptr->storage_dict_hash.value());
+    if (dict_root.not_null()) {
+      auto S = ptr->init_account_storage_stat(dict_root);
+      if (S.is_error()) {
+        fatal_error(S.move_as_error_prefix(PSTRING() << "failed to init storage stat from cache for account "
+                                                     << addr.to_hex(256) << ": "));
+        return nullptr;
+      }
+      LOG(DEBUG) << "Inited storage stat from cache for account " << addr.to_hex(256) << " (" << ptr->storage_used.cells
+                 << " cells)";
+      storage_stat_cache_update_.emplace_back(dict_root, ptr->storage_used.cells);
+    }
+  }
   return ptr;
 }
 
@@ -2589,6 +2628,10 @@ bool Collator::combine_account_transactions() {
           }
         }
       }
+      if (acc.storage_dict_hash && acc.account_storage_stat && acc.account_storage_stat.value().is_dict_ready()) {
+        storage_stat_cache_update_.emplace_back(acc.account_storage_stat.value().get_dict_root().move_as_ok(),
+                                                acc.storage_used.cells);
+      }
     } else {
       if (acc.total_state->get_hash() != acc.orig_total_state->get_hash()) {
         return fatal_error(std::string{"total state of account "} + z.first.to_hex() +
@@ -5685,6 +5728,7 @@ bool Collator::create_block_candidate() {
   stats_.cat_lt_delta = block_limit_status_->limits.classify_lt(block_limit_status_->cur_lt);
   td::actor::send_closure(manager, &ValidatorManager::record_collate_query_stats, block_candidate->id, work_time,
                           cpu_work_time, std::move(stats_));
+  td::actor::send_closure(manager, &ValidatorManager::update_storage_stat_cache, std::move(storage_stat_cache_update_));
   return true;
 }
 

validator/impl/validate-query.cpp

@@ -381,6 +381,15 @@ void ValidateQuery::start_up() {
       return;
     }
   }
+  // 5. get storage stat cache
+  ++pending;
+  LOG(DEBUG) << "sending get_storage_stat_cache() query to Manager";
+  td::actor::send_closure_later(
+      manager, &ValidatorManager::get_storage_stat_cache,
+      [self = get_self()](td::Result<std::function<td::Ref<vm::Cell>(const td::Bits256&)>> res) {
+        LOG(DEBUG) << "got answer to get_storage_stat_cache() query";
+        td::actor::send_closure_later(std::move(self), &ValidateQuery::after_get_storage_stat_cache, std::move(res));
+      });
   // ...
   CHECK(pending);
 }
@@ -749,6 +758,26 @@ void ValidateQuery::got_mc_handle(td::Result<BlockHandle> res) {
       });
 }
 
+/**
+ * Callback function called after retrieving storage stat cache.
+ *
+ * @param res The retrieved storage stat cache.
+ */
+void ValidateQuery::after_get_storage_stat_cache(td::Result<std::function<td::Ref<vm::Cell>(const td::Bits256&)>> res) {
+  --pending;
+  if (res.is_error()) {
+    LOG(INFO) << "after_get_storage_stat_cache : " << res.error();
+  } else {
+    LOG(DEBUG) << "after_get_storage_stat_cache : OK";
+    storage_stat_cache_ = res.move_as_ok();
+  }
+  if (!pending) {
+    if (!try_validate()) {
+      fatal_error("cannot validate new block");
+    }
+  }
+}
+
 /**
  * Callback function called after retrieving the shard state for a previous block.
  *
@@ -5133,6 +5162,20 @@ std::unique_ptr<block::Account> ValidateQuery::make_account_from(td::ConstBitPtr
     return nullptr;
   }
   ptr->block_lt = start_lt_;
+  if (storage_stat_cache_ && ptr->storage_dict_hash) {
+    auto dict_root = storage_stat_cache_(ptr->storage_dict_hash.value());
+    if (dict_root.not_null()) {
+      auto S = ptr->init_account_storage_stat(dict_root);
+      if (S.is_error()) {
+        fatal_error(S.move_as_error_prefix(PSTRING() << "failed to init storage stat from cache for account "
+                                                     << addr.to_hex(256) << ": "));
+        return nullptr;
+      }
+      LOG(DEBUG) << "Inited storage stat from cache for account " << addr.to_hex(256) << " (" << ptr->storage_used.cells
+                 << " cells)";
+      storage_stat_cache_update_.emplace_back(dict_root, ptr->storage_used.cells);
+    }
+  }
   return ptr;
 }
 
@@ -5746,6 +5789,11 @@ bool ValidateQuery::check_account_transactions(const StdSmcAddress& acc_addr, Re
       })) {
     return reject_query("at least one Transaction of account "s + acc_addr.to_hex() + " is invalid");
   }
+  if (account.storage_dict_hash && account.account_storage_stat &&
+      account.account_storage_stat.value().is_dict_ready()) {
+    storage_stat_cache_update_.emplace_back(account.account_storage_stat.value().get_dict_root().move_as_ok(),
+                                            account.storage_used.cells);
+  }
   if (is_masterchain() && account.libraries_changed()) {
     return scan_account_libraries(account.orig_library, account.library, acc_addr);
   } else {
@@ -6923,6 +6971,7 @@ bool ValidateQuery::save_candidate() {
 
   td::actor::send_closure(manager, &ValidatorManager::set_block_candidate, id_, block_candidate.clone(),
                           validator_set_->get_catchain_seqno(), validator_set_->get_validator_set_hash(), std::move(P));
+  td::actor::send_closure(manager, &ValidatorManager::update_storage_stat_cache, std::move(storage_stat_cache_update_));
   return true;
 }
 

validator/impl/validate-query.hpp

@@ -236,6 +236,9 @@ class ValidateQuery : public td::actor::Actor {
   std::set<StdSmcAddress> account_expected_defer_all_messages_;
   td::uint64 old_out_msg_queue_size_ = 0, new_out_msg_queue_size_ = 0;
 
+  std::function<td::Ref<vm::Cell>(const td::Bits256&)> storage_stat_cache_;
+  std::vector<std::pair<td::Ref<vm::Cell>, td::uint32>> storage_stat_cache_update_;
+
   bool msg_metadata_enabled_ = false;
   bool deferring_messages_enabled_ = false;
   bool store_out_msg_queue_size_ = false;
@@ -289,6 +292,7 @@ class ValidateQuery : public td::actor::Actor {
   void after_get_latest_mc_state(td::Result<std::pair<Ref<MasterchainState>, BlockIdExt>> res);
   void after_get_mc_state(td::Result<Ref<ShardState>> res);
   void got_mc_handle(td::Result<BlockHandle> res);
+  void after_get_storage_stat_cache(td::Result<std::function<td::Ref<vm::Cell>(const td::Bits256&)>> res);
   void after_get_shard_state(int idx, td::Result<Ref<ShardState>> res);
   bool process_mc_state(Ref<MasterchainState> mc_state);
   bool try_unpack_mc_state();

validator/interfaces/validator-manager.h

@@ -221,6 +221,13 @@ class ValidatorManager : public ValidatorManagerInterface {
 
   virtual void add_persistent_state_description(td::Ref<PersistentStateDescription> desc) = 0;
 
+  virtual void get_storage_stat_cache(td::Promise<std::function<td::Ref<vm::Cell>(const td::Bits256&)>> promise) {
+    promise.set_error(td::Status::Error("not implemented"));
+  }
+  virtual void update_storage_stat_cache(std::vector<std::pair<td::Ref<vm::Cell>, td::uint32>> data) {
+    // not implemented
+  }
+
   static bool is_persistent_state(UnixTime ts, UnixTime prev_ts) {
     return ts / (1 << 17) != prev_ts / (1 << 17);
   }