Fix speed limits in storage

SpyCheese · Jun 28, 2023 · 53be89f · 53be89f
1 parent 33ea89c
commit 53be89f
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 9 deletions.
diff --git a/storage/PeerActor.cpp b/storage/PeerActor.cpp
@@ -25,6 +25,7 @@
 #include "td/utils/overloaded.h"
 #include "td/utils/Random.h"
 #include "vm/boc.h"
+#include "common/delay.h"
 
 namespace ton {
 
@@ -119,9 +120,9 @@ void PeerActor::on_get_piece_result(PartId piece_id, td::Result<td::BufferSlice>
     return std::move(res);
   }();
   if (res.is_error()) {
-    LOG(DEBUG) << "getPiece " << piece_id << "query: " << res.error();
+    LOG(DEBUG) << "getPiece " << piece_id << " query: " << res.error();
   } else {
-    LOG(DEBUG) << "getPiece " << piece_id << "query: OK";
+    LOG(DEBUG) << "getPiece " << piece_id << " query: OK";
   }
   state_->node_queries_results_.add_element(std::make_pair(piece_id, std::move(res)));
   notify_node();
@@ -343,11 +344,20 @@ void PeerActor::loop_node_get_piece() {
       }
       auto piece_size =
           std::min<td::uint64>(torrent_info_->piece_size, torrent_info_->file_size - part * torrent_info_->piece_size);
-      td::actor::send_closure(state_->speed_limiters_.download, &SpeedLimiter::enqueue, (double)piece_size,
-                              td::Timestamp::in(3.0), [part, SelfId = actor_id(this)](td::Result<td::Unit> R) {
-                                td::actor::send_closure(SelfId, &PeerActor::node_get_piece_query_ready, part,
-                                                        std::move(R));
-                              });
+      td::Timestamp timeout = td::Timestamp::in(3.0);
+      td::actor::send_closure(
+          state_->speed_limiters_.download, &SpeedLimiter::enqueue, (double)piece_size, timeout,
+          [=, SelfId = actor_id(this)](td::Result<td::Unit> R) {
+            if (R.is_ok()) {
+              td::actor::send_closure(SelfId, &PeerActor::node_get_piece_query_ready, part, std::move(R));
+            } else {
+              delay_action(
+                  [=, R = std::move(R)]() mutable {
+                    td::actor::send_closure(SelfId, &PeerActor::node_get_piece_query_ready, part, std::move(R));
+                  },
+                  timeout);
+            }
+          });
     }
   }
 

diff --git a/storage/SpeedLimiter.cpp b/storage/SpeedLimiter.cpp
@@ -16,6 +16,7 @@
 */
 
 #include "SpeedLimiter.h"
+#include "common/errorcode.h"
 
 namespace ton {
 
@@ -41,11 +42,11 @@ void SpeedLimiter::enqueue(double size, td::Timestamp timeout, td::Promise<td::U
     return;
   }
   if (max_speed_ == 0.0) {
-    promise.set_error(td::Status::Error("Speed limit is 0"));
+    promise.set_error(td::Status::Error(ErrorCode::timeout, "Speed limit is 0"));
     return;
   }
   if (timeout < unlock_at_) {
-    promise.set_error(td::Status::Error("Timeout caused by speed limit"));
+    promise.set_error(td::Status::Error(ErrorCode::timeout, "Timeout caused by speed limit"));
     return;
   }
   if (queue_.empty() && unlock_at_.is_in_past()) {