Don't keep exact remote location for bots.

This commit is contained in:
levlam 2023-05-10 15:29:38 +03:00
parent 6f94c363f7
commit de22159b7f
3 changed files with 34 additions and 17 deletions

View File

@ -3809,6 +3809,10 @@ void Td::init_file_manager() {
FileReferenceManager::reload_photo(std::move(source), std::move(promise)); FileReferenceManager::reload_photo(std::move(source), std::move(promise));
} }
bool keep_exact_remote_location() final {
return !td_->auth_manager_->is_bot();
}
ActorShared<> create_reference() final { ActorShared<> create_reference() final {
return td_->create_reference(); return td_->create_reference();
} }

View File

@ -841,8 +841,9 @@ void FileManager::init_actor() {
FileManager::~FileManager() { FileManager::~FileManager() {
Scheduler::instance()->destroy_on_scheduler(G()->get_gc_scheduler_id(), remote_location_info_, file_hash_to_file_id_, Scheduler::instance()->destroy_on_scheduler(G()->get_gc_scheduler_id(), remote_location_info_, file_hash_to_file_id_,
local_location_to_file_id_, generate_location_to_file_id_, remote_location_to_file_id_, local_location_to_file_id_,
pmc_id_to_file_node_id_, file_id_info_, empty_file_ids_, file_nodes_); generate_location_to_file_id_, pmc_id_to_file_node_id_, file_id_info_,
empty_file_ids_, file_nodes_);
} }
string FileManager::fix_file_extension(Slice file_name, Slice file_type, Slice file_extension) { string FileManager::fix_file_extension(Slice file_name, Slice file_type, Slice file_extension) {
@ -1300,21 +1301,27 @@ Result<FileId> FileManager::register_file(FileData &&data, FileLocationSource fi
} }
}; };
bool new_remote = false; bool new_remote = false;
FileId *new_remote_file_id = nullptr;
int32 remote_key = 0; int32 remote_key = 0;
if (file_view.has_remote_location()) { if (file_view.has_remote_location()) {
RemoteInfo info{file_view.remote_location(), file_location_source, file_id}; if (context_->keep_exact_remote_location()) {
remote_key = remote_location_info_.add(info); RemoteInfo info{file_view.remote_location(), file_location_source, file_id};
auto &stored_info = remote_location_info_.get(remote_key); remote_key = remote_location_info_.add(info);
if (stored_info.file_id_ == file_id) { auto &stored_info = remote_location_info_.get(remote_key);
get_file_id_info(file_id)->pin_flag_ = true; if (stored_info.file_id_ == file_id) {
new_remote = true; get_file_id_info(file_id)->pin_flag_ = true;
} else { new_remote = true;
to_merge.push_back(stored_info.file_id_); } else {
if (merge_choose_remote_location(file_view.remote_location(), file_location_source, stored_info.remote_, to_merge.push_back(stored_info.file_id_);
stored_info.file_location_source_) == 0) { if (merge_choose_remote_location(file_view.remote_location(), file_location_source, stored_info.remote_,
stored_info.remote_ = file_view.remote_location(); stored_info.file_location_source_) == 0) {
stored_info.file_location_source_ = file_location_source; stored_info.remote_ = file_view.remote_location();
stored_info.file_location_source_ = file_location_source;
}
} }
} else {
new_remote_file_id = register_location(file_view.remote_location(), remote_location_to_file_id_);
new_remote = new_remote_file_id != nullptr;
} }
} }
FileId *new_local_file_id = nullptr; FileId *new_local_file_id = nullptr;
@ -1344,6 +1351,9 @@ Result<FileId> FileManager::register_file(FileData &&data, FileLocationSource fi
try_flush_node(get_file_node(file_id), "register_file"); try_flush_node(get_file_node(file_id), "register_file");
auto main_file_id = get_file_node(file_id)->main_file_id_; auto main_file_id = get_file_node(file_id)->main_file_id_;
if (main_file_id != file_id) { if (main_file_id != file_id) {
if (new_remote_file_id != nullptr) {
*new_remote_file_id = main_file_id;
}
if (new_local_file_id != nullptr) { if (new_local_file_id != nullptr) {
*new_local_file_id = main_file_id; *new_local_file_id = main_file_id;
} }
@ -1352,7 +1362,7 @@ Result<FileId> FileManager::register_file(FileData &&data, FileLocationSource fi
} }
try_forget_file_id(file_id); try_forget_file_id(file_id);
} }
if ((new_local_file_id != nullptr) || (new_generate_file_id != nullptr)) { if (new_cnt > 0) {
get_file_id_info(main_file_id)->pin_flag_ = true; get_file_id_info(main_file_id)->pin_flag_ = true;
} }
@ -3221,7 +3231,7 @@ Result<FileId> FileManager::check_input_file_id(FileType type, Result<FileId> re
} }
int32 remote_id = file_id.get_remote(); int32 remote_id = file_id.get_remote();
if (remote_id == 0) { if (remote_id == 0 && context_->keep_exact_remote_location()) {
RemoteInfo info{file_view.remote_location(), FileLocationSource::FromUser, file_id}; RemoteInfo info{file_view.remote_location(), FileLocationSource::FromUser, file_id};
remote_id = remote_location_info_.add(info); remote_id = remote_location_info_.add(info);
if (remote_location_info_.get(remote_id).file_id_ == file_id) { if (remote_location_info_.get(remote_id).file_id_ == file_id) {
@ -4062,7 +4072,7 @@ std::pair<FileManager::Query, bool> FileManager::finish_query(QueryId query_id)
} }
FullRemoteFileLocation *FileManager::get_remote(int32 key) { FullRemoteFileLocation *FileManager::get_remote(int32 key) {
if (key == 0) { if (key == 0 || !context_->keep_exact_remote_location()) {
return nullptr; return nullptr;
} }
return &remote_location_info_.get(key).remote_; return &remote_location_info_.get(key).remote_;

View File

@ -402,6 +402,8 @@ class FileManager final : public FileLoadManager::Callback {
virtual void reload_photo(PhotoSizeSource source, Promise<Unit> promise) = 0; virtual void reload_photo(PhotoSizeSource source, Promise<Unit> promise) = 0;
virtual bool keep_exact_remote_location() = 0;
virtual ActorShared<> create_reference() = 0; virtual ActorShared<> create_reference() = 0;
Context() = default; Context() = default;
@ -587,6 +589,7 @@ class FileManager final : public FileLoadManager::Callback {
WaitFreeHashMap<string, FileId> file_hash_to_file_id_; WaitFreeHashMap<string, FileId> file_hash_to_file_id_;
std::map<FullRemoteFileLocation, FileId> remote_location_to_file_id_;
std::map<FullLocalFileLocation, FileId> local_location_to_file_id_; std::map<FullLocalFileLocation, FileId> local_location_to_file_id_;
std::map<FullGenerateFileLocation, FileId> generate_location_to_file_id_; std::map<FullGenerateFileLocation, FileId> generate_location_to_file_id_;
std::map<FileDbId, int32> pmc_id_to_file_node_id_; std::map<FileDbId, int32> pmc_id_to_file_node_id_;