File: handle FILE_PART_INVALID error, fix expected_size usage

GitOrigin-RevId: e9c583993737194785e0f36742507724810e163c

td/telegram/files/FileBitmask.cpp

@@ -50,12 +50,21 @@ int64 Bitmask::get_ready_prefix_size(int64 offset, int64 part_size, int64 file_s
   return res;
 }
 
-int64 Bitmask::get_total_size(int64 part_size) const {
+int64 Bitmask::get_total_size(int64 part_size, int64 file_size) const {
   int64 res = 0;
   for (int64 i = 0; i < size(); i++) {
-    res += static_cast<int64>(get(i));
+    if (get(i)) {
+      auto from = i * part_size;
+      auto to = from + part_size;
+      if (file_size != 0 && file_size < to) {
+        to = file_size;
+      }
+      if (from < to) {
+        res += to - from;
+      }
+    }
   }
-  return res * part_size;
+  return res;
 }
 
 bool Bitmask::get(int64 offset_part) const {

td/telegram/files/FileBitmask.h

@@ -21,7 +21,7 @@ class Bitmask {
   Bitmask(Ones, int64 count);
   std::string encode() const;
   int64 get_ready_prefix_size(int64 offset, int64 part_size, int64 file_size) const;
-  int64 get_total_size(int64 part_size) const;
+  int64 get_total_size(int64 part_size, int64 file_size) const;
   bool get(int64 offset_part) const;
 
   int64 get_ready_parts(int64 offset_part) const;

td/telegram/files/FileLoadManager.cpp

@@ -60,7 +60,7 @@ void FileLoadManager::download(QueryId id, const FullRemoteFileLocation &remote_
 }
 
 void FileLoadManager::upload(QueryId id, const LocalFileLocation &local_location,
-                             const RemoteFileLocation &remote_location, int64 size,
+                             const RemoteFileLocation &remote_location, int64 expected_size,
                              const FileEncryptionKey &encryption_key, int8 priority, vector<int> bad_parts) {
   if (stop_flag_) {
     return;
@@ -71,7 +71,7 @@ void FileLoadManager::upload(QueryId id, const LocalFileLocation &local_location
   CHECK(node);
   node->query_id_ = id;
   auto callback = make_unique<FileUploaderCallback>(actor_shared(this, node_id));
-  node->loader_ = create_actor<FileUploader>("Uploader", local_location, remote_location, size, encryption_key,
+  node->loader_ = create_actor<FileUploader>("Uploader", local_location, remote_location, expected_size, encryption_key,
                                              std::move(bad_parts), std::move(callback));
   send_closure(upload_resource_manager_, &ResourceManager::register_worker,
                ActorShared<FileLoaderActor>(node->loader_.get(), static_cast<uint64>(-1)), priority);

td/telegram/files/FileLoadManager.h

@@ -49,7 +49,7 @@ class FileLoadManager final : public Actor {
   void download(QueryId id, const FullRemoteFileLocation &remote_location, const LocalFileLocation &local, int64 size,
                 string name, const FileEncryptionKey &encryption_key, bool search_file, int64 offset, int8 priority);
   void upload(QueryId id, const LocalFileLocation &local_location, const RemoteFileLocation &remote_location,
-              int64 size, const FileEncryptionKey &encryption_key, int8 priority, vector<int> bad_parts);
+              int64 expected_size, const FileEncryptionKey &encryption_key, int8 priority, vector<int> bad_parts);
   void upload_by_hash(QueryId id, const FullLocalFileLocation &local_location, int64 size, int8 priority);
   void update_priority(QueryId id, int8 priority);
   void from_bytes(QueryId id, FileType type, BufferSlice bytes, string name);

td/telegram/files/FileLocation.h

@@ -33,6 +33,8 @@
 
 namespace td {
 
+constexpr int64 SMALL_FILE_MAX_SIZE = 10 * (1 << 20);
+
 enum class FileType : int8 {
   Thumbnail,
   ProfilePhoto,

td/telegram/files/FileManager.cpp

@@ -91,7 +91,7 @@ void FileNode::init_ready_size() {
   }
   auto bitmask = Bitmask(Bitmask::Decode{}, local_.partial().ready_bitmask_);
   local_ready_prefix_size_ = bitmask.get_ready_prefix_size(0, local_.partial().part_size_, size_);
-  local_ready_size_ = bitmask.get_total_size(local_.partial().part_size_);
+  local_ready_size_ = bitmask.get_total_size(local_.partial().part_size_, size_);
 }
 
 void FileNode::set_download_offset(int64 download_offset) {
@@ -343,11 +343,18 @@ int64 FileView::size() const {
   return node_->size_;
 }
 
-int64 FileView::expected_size() const {
+int64 FileView::expected_size(bool may_guess) const {
   if (node_->size_ != 0) {
     return node_->size_;
   }
-  return node_->expected_size_;
+  int64 current_size = local_total_size();  // TODO: this is not the best approximation
+  if (node_->expected_size_ != 0) {
+    return max(current_size, node_->expected_size_);
+  }
+  if (may_guess && node_->local_.type() == LocalFileLocation::Type::Partial) {
+    current_size *= 3;
+  }
+  return current_size;
 }
 
 bool FileView::is_downloading() const {
@@ -1911,9 +1918,9 @@ void FileManager::run_upload(FileNodePtr node, std::vector<int> bad_parts) {
 
   QueryId id = queries_container_.create(Query{file_id, Query::Upload});
   node->upload_id_ = id;
-  send_closure(file_load_manager_, &FileLoadManager::upload, id, node->local_, node->remote_, node->size_,
+  send_closure(file_load_manager_, &FileLoadManager::upload, id, node->local_, node->remote_,
-               node->encryption_key_, narrow_cast<int8>(bad_parts.empty() ? -priority : priority),
+               file_view.expected_size(true), node->encryption_key_,
-               std::move(bad_parts));
+               narrow_cast<int8>(bad_parts.empty() ? -priority : priority), std::move(bad_parts));
 
   LOG(INFO) << "File " << file_id << " upload request has sent to FileLoadManager";
 }
@@ -2578,6 +2585,23 @@ void FileManager::on_error_impl(FileNodePtr node, FileManager::Query::Type type,
     }
   }
 
+  if (status.message() == "FILE_PART_INVALID") {
+    bool has_partial_small_location =
+        node->remote_.type() == RemoteFileLocation::Type::Partial && !node->remote_.partial().is_big_;
+    auto expected_size = FileView(node).expected_size(true);
+    bool should_be_big_location = expected_size > SMALL_FILE_MAX_SIZE;
+
+    node->set_remote_location(RemoteFileLocation(), FileLocationSource::None, 0);
+    if (has_partial_small_location && should_be_big_location) {
+      run_upload(node, {});
+      return;
+    }
+
+    LOG(WARNING) << "Failed to upload file: unexpected " << status << " " << has_partial_small_location << " "
+                 << should_be_big_location << " "
+                 << "expected size: " << expected_size;
+  }
+
   if (begins_with(status.message(), "FILE_GENERATE_LOCATION_INVALID")) {
     node->set_generate_location(nullptr);
   }

td/telegram/files/FileManager.h

@@ -220,7 +220,7 @@ class FileView {
   }
 
   int64 size() const;
-  int64 expected_size() const;
+  int64 expected_size(bool may_guess = false) const;
   bool is_downloading() const;
   int64 download_offset() const;
   int64 downloaded_prefix(int64 offset) const;

td/telegram/files/FileUploader.cpp

@@ -61,7 +61,7 @@ Result<FileLoader::FileInfo> FileUploader::init() {
     offset = partial.ready_part_count_;
   } else {
     file_id_ = Random::secure_int64();
-    big_flag_ = expected_size_ > 10 * (1 << 20);
+    big_flag_ = expected_size_ > SMALL_FILE_MAX_SIZE;
   }
 
   std::vector<bool> ok(offset, true);