Limit imported msg queue size

crypto/block/block.tlb

@@ -829,6 +829,7 @@ top_block_descr#d5 proof_for:BlockIdExt signatures:(Maybe ^BlockSignatures)
 //  COLLATED DATA
 //
 top_block_descr_set#4ac789f3 collection:(HashmapE 96 ^TopBlockDescr) = TopBlockDescrSet;
+neighbor_msg_queue_limits#7e549333 neighbors:(HashmapE 96 int32) = NeighborMsgQueueLimits;
 
 //
 //  VALIDATOR MISBEHAVIOR COMPLAINTS

crypto/block/output-queue-merger.cpp

@@ -134,34 +134,33 @@ bool OutputQueueMerger::MsgKeyValue::split(MsgKeyValue& second) {
   return true;
 }
 
-bool OutputQueueMerger::add_root(int src, Ref<vm::Cell> outmsg_root) {
+void OutputQueueMerger::add_root(int src, Ref<vm::Cell> outmsg_root, td::int32 msg_limit) {
   if (outmsg_root.is_null()) {
-    return true;
+    return;
   }
   //block::gen::HashmapAug{352, block::gen::t_EnqueuedMsg, block::gen::t_uint64}.print_ref(std::cerr, outmsg_root);
   auto kv = std::make_unique<MsgKeyValue>(src, std::move(outmsg_root));
   if (kv->replace_by_prefix(common_pfx.cbits(), common_pfx_len)) {
     heap.push_back(std::move(kv));
   }
-  return true;
-}
-
-OutputQueueMerger::OutputQueueMerger(ton::ShardIdFull _queue_for, std::vector<block::McShardDescr> _neighbors)
-    : queue_for(_queue_for), neighbors(std::move(_neighbors)), eof(false), failed(false) {
-  init();
+  if ((int)src_remaining_msgs_.size() < src + 1) {
+    src_remaining_msgs_.resize(src + 1);
+  }
+  src_remaining_msgs_[src] = msg_limit;
 }
 
-void OutputQueueMerger::init() {
+OutputQueueMerger::OutputQueueMerger(ton::ShardIdFull queue_for, std::vector<OutputQueueMerger::Neighbor> neighbors)
+    : eof(false), failed(false) {
   common_pfx.bits().store_int(queue_for.workchain, 32);
   int l = queue_for.pfx_len();
   td::bitstring::bits_store_long_top(common_pfx.bits() + 32, queue_for.shard, l);
   common_pfx_len = 32 + l;
   int i = 0;
-  for (block::McShardDescr& neighbor : neighbors) {
-    if (!neighbor.is_disabled()) {
-      LOG(DEBUG) << "adding " << (neighbor.outmsg_root.is_null() ? "" : "non-") << "empty output queue for neighbor #"
-                 << i << " (" << neighbor.blk_.to_str() << ")";
-      add_root(i++, neighbor.outmsg_root);
+  for (Neighbor& neighbor : neighbors) {
+    if (!neighbor.disabled_) {
+      LOG(DEBUG) << "adding " << (neighbor.outmsg_root_.is_null() ? "" : "non-") << "empty output queue for neighbor #"
+                 << i << " (" << neighbor.block_id_.to_str() << ")";
+      add_root(i++, neighbor.outmsg_root_, neighbor.msg_limit_);
     } else {
       LOG(DEBUG) << "skipping output queue for disabled neighbor #" << i;
       i++;
@@ -200,6 +199,10 @@ bool OutputQueueMerger::load() {
   unsigned long long lt = heap[0]->lt;
   std::size_t orig_size = msg_list.size();
   do {
+    if (src_remaining_msgs_[heap[0]->source] == 0) {
+      std::pop_heap(heap.begin(), heap.end(), MsgKeyValue::greater);
+      continue;
+    }
     while (heap[0]->is_fork()) {
       auto other = std::make_unique<MsgKeyValue>();
       if (!heap[0]->split(*other)) {
@@ -215,6 +218,17 @@ bool OutputQueueMerger::load() {
     heap.pop_back();
   } while (!heap.empty() && heap[0]->lt <= lt);
   std::sort(msg_list.begin() + orig_size, msg_list.end(), MsgKeyValue::less);
+  size_t j = orig_size;
+  for (size_t i = orig_size; i < msg_list.size(); ++i) {
+    td::int32 &remaining = src_remaining_msgs_[msg_list[i]->source];
+    if (remaining != 0) {
+      if (remaining > 0) {
+        --remaining;
+      }
+      msg_list[j++] = std::move(msg_list[i]);
+    }
+  }
+  msg_list.resize(j);
   return true;
 }
 

crypto/block/output-queue-merger.h

@@ -52,12 +52,20 @@ struct OutputQueueMerger {
     bool split(MsgKeyValue& second);
   };
   //
-  ton::ShardIdFull queue_for;
   std::vector<std::unique_ptr<MsgKeyValue>> msg_list;
-  std::vector<block::McShardDescr> neighbors;
 
  public:
-  OutputQueueMerger(ton::ShardIdFull _queue_for, std::vector<block::McShardDescr> _neighbors);
+  struct Neighbor {
+    ton::BlockIdExt block_id_;
+    td::Ref<vm::Cell> outmsg_root_;
+    bool disabled_;
+    td::int32 msg_limit_;  // -1 - unlimited
+    Neighbor(ton::BlockIdExt block_id, td::Ref<vm::Cell> outmsg_root, bool disabled = false, td::int32 msg_limit = -1)
+        : block_id_(block_id), outmsg_root_(std::move(outmsg_root)), disabled_(disabled), msg_limit_(msg_limit) {
+    }
+  };
+
+  OutputQueueMerger(ton::ShardIdFull queue_for, std::vector<Neighbor> neighbors);
   bool is_eof() const {
     return eof;
   }
@@ -70,10 +78,10 @@ struct OutputQueueMerger {
   int common_pfx_len;
   std::vector<std::unique_ptr<MsgKeyValue>> heap;
   std::size_t pos{0};
+  std::vector<td::int32> src_remaining_msgs_;
   bool eof;
   bool failed;
-  void init();
-  bool add_root(int src, Ref<vm::Cell> outmsg_root);
+  void add_root(int src, Ref<vm::Cell> outmsg_root, td::int32 msg_limit);
   bool load();
 };
 

tl/generate/scheme/ton_api.tl

@@ -415,7 +415,7 @@ tonNode.success = tonNode.Success;
 tonNode.archiveNotFound = tonNode.ArchiveInfo;
 tonNode.archiveInfo id:long = tonNode.ArchiveInfo;
 
-tonNode.outMsgQueueProof queue_proof:bytes block_state_proof:bytes = tonNode.OutMsgQueueProof;
+tonNode.outMsgQueueProof queue_proof:bytes block_state_proof:bytes msg_count:int = tonNode.OutMsgQueueProof;
 tonNode.outMsgQueueProofEmpty = tonNode.OutMsgQueueProof;
 
 tonNode.forgetPeer = tonNode.ForgetPeer;

validator/impl/collator-impl.h

@@ -133,6 +133,8 @@ class Collator final : public td::actor::Actor {
   std::unique_ptr<block::ConfigInfo> config_;
   std::unique_ptr<block::ShardConfig> shard_conf_;
   std::map<BlockSeqno, Ref<MasterchainStateQ>> aux_mc_states_;
+  std::vector<block::OutputQueueMerger::Neighbor> neighbor_queues_;
+  vm::Dictionary neighbor_msg_queues_limits_{32 + 64};
   std::vector<block::McShardDescr> neighbors_;
   std::unique_ptr<block::OutputQueueMerger> nb_out_msgs_;
   std::vector<ton::StdSmcAddress> special_smcs;

validator/impl/collator.cpp

@@ -77,7 +77,7 @@ Collator::Collator(ShardIdFull shard, bool is_hardfork, BlockIdExt min_mastercha
 }
 
 void Collator::start_up() {
-  LOG(DEBUG) << "Collator for shard " << shard_.to_str() << " started";
+  LOG(INFO) << "Collator for shard " << shard_.to_str() << " started";
   LOG(DEBUG) << "Previous block #1 is " << prev_blocks.at(0).to_str();
   if (prev_blocks.size() > 1) {
     LOG(DEBUG) << "Previous block #2 is " << prev_blocks.at(1).to_str();
@@ -690,7 +690,15 @@ void Collator::got_neighbor_msg_queue(unsigned i, td::Result<Ref<OutMsgQueueProo
     fatal_error("cannot unpack neighbor output queue info");
     return;
   }
-  descr.set_queue_root(qinfo.out_queue->prefetch_ref(0));
+  auto queue_root = qinfo.out_queue->prefetch_ref(0);
+  descr.set_queue_root(queue_root);
+  neighbor_queues_.emplace_back(descr.top_block_id(), queue_root, descr.is_disabled(), res->msg_count_);
+  if (res->msg_count_ != -1) {
+    td::BitArray<96> key;
+    key.bits().store_int(block_id.id.workchain, 32);
+    (key.bits() + 32).store_uint(block_id.id.shard, 64);
+    neighbor_msg_queues_limits_.set_builder(key, vm::CellBuilder().store_long(res->msg_count_, 32));
+  }
   // comment the next two lines in the future when the output queues become huge
   //  CHECK(block::gen::t_OutMsgQueueInfo.validate_ref(1000000, outq_descr->root_cell()));
   //  CHECK(block::tlb::t_OutMsgQueueInfo.validate_ref(1000000, outq_descr->root_cell()));
@@ -1723,7 +1731,7 @@ bool Collator::do_collate() {
   // 1.3. create OutputQueueMerger from adjusted neighbors
   CHECK(!nb_out_msgs_);
   LOG(DEBUG) << "creating OutputQueueMerger";
-  nb_out_msgs_ = std::make_unique<block::OutputQueueMerger>(shard_, neighbors_);
+  nb_out_msgs_ = std::make_unique<block::OutputQueueMerger>(shard_, neighbor_queues_);
   // 1.4. compute created / minted / recovered
   if (!init_value_create()) {
     return fatal_error("cannot compute the value to be created / minted / recovered");
@@ -4013,18 +4021,25 @@ bool Collator::create_collated_data() {
     auto cell = collate_shard_block_descr_set();
     if (cell.is_null()) {
       return true;
-      return fatal_error("cannot collate the collection of used shard block descriptions");
+      // return fatal_error("cannot collate the collection of used shard block descriptions");
     }
     collated_roots_.push_back(std::move(cell));
   }
+  // 2. Message count for neighbors' out queues
+  if (!neighbor_msg_queues_limits_.is_empty()) {
+    vm::CellBuilder cb;
+    cb.store_long(block::gen::t_NeighborMsgQueueLimits.cons_tag[0], 32);
+    cb.store_maybe_ref(neighbor_msg_queues_limits_.get_root_cell());
+    collated_roots_.push_back(cb.finalize_novm());
+  }
   if (!full_collated_data_) {
     return true;
   }
-  // 2. Proofs for hashes of states: previous states + neighbors
+  // 3. Proofs for hashes of states: previous states + neighbors
   for (const auto& p : block_state_proofs_) {
     collated_roots_.push_back(p.second);
   }
-  // 3. Previous state proof (only shadchains)
+  // 4. Previous state proof (only shadchains)
   std::map<td::Bits256, Ref<vm::Cell>> proofs;
   if (!is_masterchain()) {
     if (!prepare_msg_queue_proof()) {
@@ -4038,7 +4053,7 @@ bool Collator::create_collated_data() {
     }
     proofs[prev_state_root_->get_hash().bits()] = std::move(state_proof);
   }
-  // 4. Proofs for message queues
+  // 5. Proofs for message queues
   for (vm::MerkleProofBuilder &mpb : neighbor_proof_builders_) {
     auto r_proof = mpb.extract_proof();
     if (r_proof.is_error()) {