FileManager: store file source in db (partial imlementation)

GitOrigin-RevId: 860e92368e226401e57072e8d64df48c263029f6

td/telegram/FileReferenceManager.cpp

@@ -93,6 +93,14 @@ void FileReferenceManager::remove_file_source(NodeId node_id, FileSourceId file_
   nodes_[node_id].file_source_ids.remove(file_source_id);
 }
 
+std::vector<FileSourceId> FileReferenceManager::get_some_file_sources(NodeId node_id) {
+  auto it = nodes_.find(node_id);
+  if (it == nodes_.end()) {
+    return {};
+  }
+  return it->second.file_source_ids.get_some_elements();
+}
+
 void FileReferenceManager::merge(NodeId to_node_id, NodeId from_node_id) {
   auto from_it = nodes_.find(from_node_id);
   if (from_it == nodes_.end()) {
@@ -177,17 +185,19 @@ void FileReferenceManager::send_query(Destination dest, FileSourceId file_source
           send_closure(file_reference_manager, &FileReferenceManager::on_query_result, dest, file_source_id,
                        std::move(status), 0);
         });
-    if (result.is_error()) {
-      new_promise.set_result(std::move(result));
-    }
-    send_lambda(file_manager, [file_manager, dest, new_promise = std::move(new_promise)]() mutable {
+
+    send_lambda(file_manager, [file_manager, dest, result = std::move(result), file_source_id,
+                               new_promise = std::move(new_promise)]() mutable {
       auto view = file_manager.get_actor_unsafe()->get_file_view(dest.node_id);
       CHECK(!view.empty());
-      if (view.has_active_remote_location()) {
-        new_promise.set_value({});
-      } else {
-        new_promise.set_error(Status::Error("No active remote location"));
+      if (result.is_ok() && !view.has_active_remote_location()) {
+        result = Status::Error("No active remote location");
       }
+      if (result.is_error() && result.error().code() != 429 && result.error().code() < 500) {
+        VLOG(file_references) << "Invalid " << file_source_id << " " << result.error();
+        file_manager.get_actor_unsafe()->remove_file_source(dest.node_id, file_source_id);
+      }
+      new_promise.set_result(std::move(result));
     });
   });
   auto index = static_cast<size_t>(file_source_id.get()) - 1;
@@ -253,10 +263,6 @@ FileReferenceManager::Destination FileReferenceManager::on_query_result(Destinat
     }
     node.query = {};
   }
-  if (status.is_error() && status.error().code() != 429 && status.error().code() < 500 && !G()->close_flag()) {
-    VLOG(file_references) << "Invalid " << file_source_id << " " << status;
-    remove_file_source(dest.node_id, file_source_id);
-  }
 
   run_node(dest.node_id);
   return dest;

td/telegram/FileReferenceManager.h

@@ -45,6 +45,8 @@ class FileReferenceManager : public Actor {
 
   void add_file_source(NodeId node_id, FileSourceId file_source_id);
 
+  std::vector<FileSourceId> get_some_file_sources(NodeId node_id);
+
   void remove_file_source(NodeId node_id, FileSourceId file_source_id);
 
   void merge(NodeId to_node_id, NodeId from_node_id);

td/telegram/SetWithPosition.h

@@ -18,6 +18,26 @@ namespace td {
 template <class T>
 class FastSetWithPosition {
  public:
+  std::vector<T> get_some_elements() const {
+    std::vector<T> res;
+    using std::prev;
+    using std::next;
+    res.reserve(5);
+    if (!checked_.empty()) {
+      res.push_back(*begin(checked_));
+      res.push_back(*prev(end(checked_)));
+    }
+    if (!not_checked_.empty()) {
+      res.push_back(*begin(not_checked_));
+      res.push_back(*prev(end(not_checked_)));
+    }
+    std::sort(res.begin(), res.end());
+    res.erase(std::unique(res.begin(), res.end()), res.end());
+    if (res.size() > 2) {
+      res.erase(next(res.begin()), prev(res.end()));
+    }
+    return res;
+  }
   void add(T x) {
     if (checked_.count(x) != 0) {
       return;
@@ -89,6 +109,15 @@ class FastSetWithPosition {
 template <class T>
 class SetWithPosition {
  public:
+  std::vector<T> get_some_elements() const {
+    if (fast_) {
+      return fast_->get_some_elements();
+    }
+    if (has_value_) {
+      return {value_};
+    }
+    return {};
+  }
   void add(T x) {
     if (fast_) {
       fast_->add(x);

td/telegram/files/FileData.h

@@ -9,6 +9,7 @@
 #include "td/telegram/DialogId.h"
 #include "td/telegram/files/FileEncryptionKey.h"
 #include "td/telegram/files/FileLocation.h"
+#include "td/telegram/files/FileSourceId.h"
 
 #include "td/utils/common.h"
 #include "td/utils/format.h"
@@ -29,6 +30,7 @@ class FileData {
   string remote_name_;
   string url_;
   FileEncryptionKey encryption_key_;
+  std::vector<FileSourceId> sources_;
 
   template <class StorerT>
   void store(StorerT &storer) const {
@@ -36,10 +38,13 @@ class FileData {
     bool has_owner_dialog_id = owner_dialog_id_.is_valid();
     bool has_expected_size = size_ == 0 && expected_size_ != 0;
     bool encryption_key_is_secure = encryption_key_.is_secure();
+    bool has_sources = !sources_.empty();
     BEGIN_STORE_FLAGS();
     STORE_FLAG(has_owner_dialog_id);
     STORE_FLAG(has_expected_size);
     STORE_FLAG(encryption_key_is_secure);
+    //TODO: uncomment
+    //STORE_FLAG(has_sources);
     END_STORE_FLAGS();
 
     if (has_owner_dialog_id) {
@@ -58,6 +63,10 @@ class FileData {
     store(remote_name_, storer);
     store(url_, storer);
     store(encryption_key_, storer);
+    if (has_sources) {
+      // TODO: uncomment
+      // store(sources_, storer);
+    }
   }
   template <class ParserT>
   void parse(ParserT &parser) {
@@ -65,10 +74,12 @@ class FileData {
     bool has_owner_dialog_id;
     bool has_expected_size;
     bool encryption_key_is_secure;
+    bool has_sources;
     BEGIN_PARSE_FLAGS();
     PARSE_FLAG(has_owner_dialog_id);
     PARSE_FLAG(has_expected_size);
     PARSE_FLAG(encryption_key_is_secure);
+    PARSE_FLAG(has_sources);
     END_PARSE_FLAGS_GENERIC();
 
     if (has_owner_dialog_id) {
@@ -93,6 +104,10 @@ class FileData {
     parse(url_, parser);
     encryption_key_.parse(encryption_key_is_secure ? FileEncryptionKey::Type::Secure : FileEncryptionKey::Type::Secret,
                           parser);
+    if (has_sources) {
+      //TODO: uncomment
+      // parse(sources_, parser);
+    }
   }
 };
 
@@ -112,6 +127,7 @@ inline StringBuilder &operator<<(StringBuilder &sb, const FileData &file_data) {
   if (file_data.remote_.type() == RemoteFileLocation::Type::Full) {
     sb << " remote " << file_data.remote_.full();
   }
+  sb << format::as_array(file_data.sources_);
   return sb << "]";
 }
 

td/telegram/files/FileManager.cpp

@@ -1361,6 +1361,8 @@ void FileManager::add_file_source(FileId file_id, FileSourceId file_source_id) {
   CHECK(file_source_id.is_valid());
   send_closure(G()->file_reference_manager(), &FileReferenceManager::add_file_source, node->main_file_id_,
                file_source_id);
+  node->on_pmc_changed();
+  try_flush_node(node, "add_file_source");
 }
 
 void FileManager::remove_file_source(FileId file_id, FileSourceId file_source_id) {
@@ -1372,6 +1374,8 @@ void FileManager::remove_file_source(FileId file_id, FileSourceId file_source_id
   CHECK(file_source_id.is_valid());
   send_closure(G()->file_reference_manager(), &FileReferenceManager::remove_file_source, node->main_file_id_,
                file_source_id);
+  node->on_pmc_changed();
+  try_flush_node(node, "remove_file_source");
 }
 
 void FileManager::change_files_source(FileSourceId file_source_id, const vector<FileId> &old_file_ids,
@@ -1386,13 +1390,13 @@ void FileManager::change_files_source(FileSourceId file_source_id, const vector<
   for (auto file_id : old_main_file_ids) {
     auto it = new_main_file_ids.find(file_id);
     if (it == new_main_file_ids.end()) {
-      send_closure(G()->file_reference_manager(), &FileReferenceManager::remove_file_source, file_id, file_source_id);
+      remove_file_source(file_id, file_source_id);
     } else {
       new_main_file_ids.erase(it);
     }
   }
   for (auto file_id : new_main_file_ids) {
-    send_closure(G()->file_reference_manager(), &FileReferenceManager::add_file_source, file_id, file_source_id);
+    add_file_source(file_id, file_source_id);
   }
 }
 
@@ -1509,6 +1513,7 @@ void FileManager::flush_to_pmc(FileNodePtr node, bool new_remote, bool new_local
   data.encryption_key_ = node->encryption_key_;
   data.url_ = node->url_;
   data.owner_dialog_id_ = node->owner_dialog_id_;
+  data.sources_ = G()->file_reference_manager().get_actor_unsafe()->get_some_file_sources(view.file_id());
 
   file_db_->set_file_data(node->pmc_id_, data, (create_flag || new_remote), (create_flag || new_local),
                           (create_flag || new_generate));