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;
+  if ((int)src_remaining_msgs_.size() < src + 1) {
+    src_remaining_msgs_.resize(src + 1);
+  }
+  src_remaining_msgs_[src] = msg_limit;
 }
 
-OutputQueueMerger::OutputQueueMerger(ton::ShardIdFull _queue_for, std::vector<block::McShardDescr> _neighbors)
-    : queue_for(_queue_for), neighbors(std::move(_neighbors)), eof(false), failed(false) {
-  init();
-}
-
-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();
 };