From 3b4d55233c1cc21920f728882eabf6003d6ffabd Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Wed, 26 Aug 2020 12:24:06 +0300
Subject: [PATCH] Improve protected parts calculation.

GitOrigin-RevId: f066b0680a11b6c376eead756f9d67880eb6d17e
---
 td/telegram/files/FileLoader.cpp      | 8 ++++++--
 td/telegram/files/ResourceManager.cpp | 2 +-
 td/telegram/files/ResourceManager.h   | 2 ++
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/td/telegram/files/FileLoader.cpp b/td/telegram/files/FileLoader.cpp
index e656badf0..2485d4abe 100644
--- a/td/telegram/files/FileLoader.cpp
+++ b/td/telegram/files/FileLoader.cpp
@@ -68,8 +68,11 @@ void FileLoader::update_local_file_location(const LocalFileLocation &local) {
 void FileLoader::update_downloaded_part(int64 offset, int64 limit) {
   if (parts_manager_.get_streaming_offset() != offset) {
     auto begin_part_id = parts_manager_.set_streaming_offset(offset, limit);
-    auto end_part_id = begin_part_id + static_cast<int32>(part_map_.size()) * 2;
-    VLOG(files) << "Protect parts " << begin_part_id << " ... " << end_part_id;
+    auto new_end_part_id = limit <= 0 ? parts_manager_.get_part_count()
+                                      : static_cast<int32>((offset + limit - 1) / parts_manager_.get_part_size()) + 1;
+    auto max_parts = static_cast<int32>(ResourceManager::MAX_RESOURCE_LIMIT / parts_manager_.get_part_size());
+    auto end_part_id = begin_part_id + td::min(max_parts, new_end_part_id - begin_part_id);
+    VLOG(files) << "Protect parts " << begin_part_id << " ... " << end_part_id - 1;
     for (auto &it : part_map_) {
       if (!(begin_part_id <= it.second.first.id && it.second.first.id < end_part_id)) {
         VLOG(files) << "Cancel part " << it.second.first.id;
@@ -224,6 +227,7 @@ void FileLoader::tear_down() {
   ordered_parts_.clear([](auto &&part) { part.second->clear(); });
   send_closure(std::move(delay_dispatcher_), &DelayDispatcher::close_silent);
 }
+
 void FileLoader::update_estimated_limit() {
   if (stop_flag_) {
     return;
diff --git a/td/telegram/files/ResourceManager.cpp b/td/telegram/files/ResourceManager.cpp
index 39c9fcf28..a66b0fc81 100644
--- a/td/telegram/files/ResourceManager.cpp
+++ b/td/telegram/files/ResourceManager.cpp
@@ -132,7 +132,7 @@ void ResourceManager::loop() {
     return;
   }
   auto active_limit = resource_state_.active_limit();
-  resource_state_.update_limit(2 * 1024 * (1 << 10) - active_limit);
+  resource_state_.update_limit(MAX_RESOURCE_LIMIT - active_limit);
   LOG(INFO) << tag("unused", resource_state_.unused());
 
   if (mode_ == Mode::Greedy) {
diff --git a/td/telegram/files/ResourceManager.h b/td/telegram/files/ResourceManager.h
index 3eb9fe535..eaeba97c2 100644
--- a/td/telegram/files/ResourceManager.h
+++ b/td/telegram/files/ResourceManager.h
@@ -28,6 +28,8 @@ class ResourceManager : public Actor {
 
   void register_worker(ActorShared<FileLoaderActor> callback, int8 priority);
 
+  static constexpr int64 MAX_RESOURCE_LIMIT = 1 << 21;
+
  private:
   Mode mode_;
   using NodeId = uint64;