FileManager: download big files with unknown size

GitOrigin-RevId: 3f4246f2d61127015138a32f2e93c8f932a81c50

td/telegram/files/FileBitmask.cpp

@@ -21,6 +21,20 @@ Bitmask::Bitmask(Ones, int64 count) : data_(narrow_cast<size_t>((count + 7) / 8)
   }
 }
 
+Bitmask Bitmask::compress(int k) {
+  Bitmask res;
+  for (int64 i = 0; i * k < size(); i++) {
+    bool f = true;
+    for (int64 j = 0; j < k && f; j++) {
+      f &= get(i * k + j);
+    }
+    if (f) {
+      res.set(i);
+    }
+  }
+  return res;
+}
+
 std::string Bitmask::encode(int32 prefix_count) {
   // remove zeroes in the end to make encoding deterministic
   td::Slice data(data_);

td/telegram/files/FileBitmask.h

@@ -30,6 +30,8 @@ class Bitmask {
   void set(int64 offset_part);
   int64 size() const;
 
+  Bitmask compress(int k);
+
  private:
   std::string data_;
 };

td/telegram/files/FileManager.cpp

@@ -892,6 +892,37 @@ Status FileManager::check_local_location(FileNodePtr node) {
   return status;
 }
 
+bool FileManager::try_fix_parital_local_location(FileNodePtr node) {
+  LOG(INFO) << "Trying to fix partial local location";
+  if (node->local_.type() != LocalFileLocation::Type::Partial) {
+    LOG(INFO) << "   failed - not a partial location";
+    return false;
+  }
+  auto partial = node->local_.partial();
+  if (!partial.iv_.empty()) {
+    // can't recalc iv_
+    LOG(INFO) << "   failed - partial location has nonempty iv";
+    return false;
+  }
+  if (partial.part_size_ >= 512 * (1 << 10)) {
+    LOG(INFO) << "   failed - too big part_size already: " << partial.part_size_;
+    return false;
+  }
+  auto old_part_size = partial.part_size_;
+  int new_part_size = 512 * (1 << 10);
+  auto k = new_part_size / old_part_size;
+  Bitmask mask(Bitmask::Decode(), partial.ready_bitmask_);
+  auto new_mask = mask.compress(k);
+
+  partial.part_size_ = new_part_size;
+  partial.ready_bitmask_ = new_mask.encode();
+
+  auto ready_size = new_mask.get_total_size(partial.part_size_, node->size_);
+  node->set_local_location(LocalFileLocation(partial), ready_size, -1, -1);
+  LOG(INFO) << "   ok: increase part_size " << old_part_size << "->" << new_part_size;
+  return true;
+}
+
 FileManager::FileIdInfo *FileManager::get_file_id_info(FileId file_id) {
   LOG_CHECK(0 <= file_id.get() && file_id.get() < static_cast<int32>(file_id_info_.size()))
       << file_id << " " << file_id_info_.size();
@@ -3413,11 +3444,18 @@ void FileManager::on_error_impl(FileNodePtr node, FileManager::Query::Type type,
   if (begins_with(status.message(), "FILE_DOWNLOAD_RESTART")) {
     if (ends_with(status.message(), "WITH_FILE_REFERENCE")) {
       node->download_was_update_file_reference_ = true;
+      run_download(node);
+      return;
+    } else if (ends_with(status.message(), "INCREASE_PART_SIZE")) {
+      if (try_fix_parital_local_location(node)) {
+        run_download(node);
+        return;
+      }
     } else {
       node->can_search_locally_ = false;
+      run_download(node);
+      return;
     }
-    run_download(node);
-    return;
   }
 
   if (!was_active) {

td/telegram/files/FileManager.h

@@ -564,6 +564,7 @@ class FileManager : public FileLoadManager::Callback {
   FileId register_pmc_file_data(FileData &&data);
 
   Status check_local_location(FileNodePtr node);
+  bool try_fix_parital_local_location(FileNodePtr node);
   Status check_local_location(FullLocalFileLocation &location, int64 &size);
   void try_flush_node_full(FileNodePtr node, bool new_remote, bool new_local, bool new_generate, FileDbId other_pmc_id);
   void try_flush_node(FileNodePtr node, const char *source);

td/telegram/files/PartsManager.cpp

@@ -266,6 +266,10 @@ Result<Part> PartsManager::start_part() {
     if (unknown_size_flag_) {
       part_count_++;
       if (part_count_ > MAX_PART_COUNT) {
+        if (!is_upload_) {
+          // Caller will try to increase part size if it is possible
+          return Status::Error("FILE_DOWNLOAD_RESTART_INCREASE_PART_SIZE");
+        }
         return Status::Error("Too big file with unknown size");
       }
       part_status_.push_back(PartStatus::Empty);