From 9661676646f059e1103d2b3cab4bd33712de53ea Mon Sep 17 00:00:00 2001 From: SpyCheese Date: Thu, 15 Aug 2024 15:26:35 +0300 Subject: [PATCH] Improve dht lookup in overlays (#1104) Continue dht lookup even if value was found --- dht/dht-in.hpp | 1 + dht/dht-query.cpp | 34 +++++++++++++++++++++++--- dht/dht-query.hpp | 51 +++++++++++++++++++++++++++++++++------ dht/dht.cpp | 17 ++++++++++++- dht/dht.h | 1 + overlay/overlay.cpp | 49 +++++++++++++++++++------------------ overlay/overlay.hpp | 3 ++- tdutils/td/utils/Status.h | 6 +++++ 8 files changed, 127 insertions(+), 35 deletions(-) diff --git a/dht/dht-in.hpp b/dht/dht-in.hpp index c2d20455f..0d668d438 100644 --- a/dht/dht-in.hpp +++ b/dht/dht-in.hpp @@ -179,6 +179,7 @@ class DhtMemberImpl : public DhtMember { void get_value(DhtKey key, td::Promise result) override { get_value_in(key.compute_key_id(), std::move(result)); } + void get_value_many(DhtKey key, std::function callback, td::Promise promise) override; void alarm() override { alarm_timestamp() = td::Timestamp::in(1.0); diff --git a/dht/dht-query.cpp b/dht/dht-query.cpp index b84ef8c37..3d43b1069 100644 --- a/dht/dht-query.cpp +++ b/dht/dht-query.cpp @@ -210,8 +210,11 @@ void DhtQueryFindValue::on_result(td::Result R, adnl::AdnlNodeI send_get_nodes = true; return; } - promise_.set_value(std::move(value)); - need_stop = true; + if (on_value_found(std::move(value))) { + send_get_nodes = true; + } else { + need_stop = true; + } }, [&](ton_api::dht_valueNotFound &v) { add_nodes(DhtNodesList{std::move(v.nodes_), our_network_id()}); @@ -244,7 +247,32 @@ void DhtQueryFindValue::on_result_nodes(td::Result R, adnl::Adn } void DhtQueryFindValue::finish(DhtNodesList list) { - promise_.set_error(td::Status::Error(ErrorCode::notready, "dht key not found")); +} + +bool DhtQueryFindValueSingle::on_value_found(DhtValue value) { + promise_.set_value(std::move(value)); + found_ = true; + return false; +} + +void DhtQueryFindValueSingle::tear_down() { + if (!found_) { + promise_.set_error(td::Status::Error(ErrorCode::notready, "dht key not found")); + } +} + +bool DhtQueryFindValueMany::on_value_found(DhtValue value) { + callback_(std::move(value)); + found_ = true; + return true; +} + +void DhtQueryFindValueMany::tear_down() { + if (found_) { + promise_.set_value(td::Unit()); + } else { + promise_.set_error(td::Status::Error(ErrorCode::notready, "dht key not found")); + } } DhtQueryStore::DhtQueryStore(DhtValue key_value, DhtMember::PrintId print_id, adnl::AdnlNodeIdShort src, diff --git a/dht/dht-query.hpp b/dht/dht-query.hpp index e47403618..e305f107a 100644 --- a/dht/dht-query.hpp +++ b/dht/dht-query.hpp @@ -126,16 +126,11 @@ class DhtQueryFindNodes : public DhtQuery { }; class DhtQueryFindValue : public DhtQuery { - private: - td::Promise promise_; - public: DhtQueryFindValue(DhtKeyId key, DhtMember::PrintId print_id, adnl::AdnlNodeIdShort src, DhtNodesList list, td::uint32 k, td::uint32 a, td::int32 our_network_id, DhtNode self, bool client_only, - td::actor::ActorId node, td::actor::ActorId adnl, - td::Promise promise) - : DhtQuery(key, print_id, src, k, a, our_network_id, std::move(self), client_only, node, adnl) - , promise_(std::move(promise)) { + td::actor::ActorId node, td::actor::ActorId adnl) + : DhtQuery(key, print_id, src, k, a, our_network_id, std::move(self), client_only, node, adnl) { add_nodes(std::move(list)); } void send_one_query(adnl::AdnlNodeIdShort id) override; @@ -146,6 +141,48 @@ class DhtQueryFindValue : public DhtQuery { std::string get_name() const override { return "find value"; } + + virtual bool on_value_found(DhtValue value) = 0; +}; + +class DhtQueryFindValueSingle : public DhtQueryFindValue { + public: + DhtQueryFindValueSingle(DhtKeyId key, DhtMember::PrintId print_id, adnl::AdnlNodeIdShort src, DhtNodesList list, + td::uint32 k, td::uint32 a, td::int32 our_network_id, DhtNode self, bool client_only, + td::actor::ActorId node, td::actor::ActorId adnl, + td::Promise promise) + : DhtQueryFindValue(key, print_id, src, std::move(list), k, a, our_network_id, std::move(self), client_only, node, + adnl) + , promise_(std::move(promise)) { + add_nodes(std::move(list)); + } + bool on_value_found(DhtValue value) override; + void tear_down() override; + + private: + td::Promise promise_; + bool found_ = false; +}; + +class DhtQueryFindValueMany : public DhtQueryFindValue { + public: + DhtQueryFindValueMany(DhtKeyId key, DhtMember::PrintId print_id, adnl::AdnlNodeIdShort src, DhtNodesList list, + td::uint32 k, td::uint32 a, td::int32 our_network_id, DhtNode self, bool client_only, + td::actor::ActorId node, td::actor::ActorId adnl, + std::function callback, td::Promise promise) + : DhtQueryFindValue(key, print_id, src, std::move(list), k, a, our_network_id, std::move(self), client_only, node, + adnl) + , callback_(std::move(callback)) + , promise_(std::move(promise)) { + add_nodes(std::move(list)); + } + bool on_value_found(DhtValue value) override; + void tear_down() override; + + private: + std::function callback_; + td::Promise promise_; + bool found_ = false; }; class DhtQueryStore : public td::actor::Actor { diff --git a/dht/dht.cpp b/dht/dht.cpp index 8d7b02b7d..b774b5f91 100644 --- a/dht/dht.cpp +++ b/dht/dht.cpp @@ -470,7 +470,7 @@ void DhtMemberImpl::get_value_in(DhtKeyId key, td::Promise result) { network_id = network_id_, id = id_, client_only = client_only_](td::Result R) mutable { R.ensure(); - td::actor::create_actor("FindValueQuery", key, print_id, id, std::move(list), k, a, network_id, + td::actor::create_actor("FindValueQuery", key, print_id, id, std::move(list), k, a, network_id, R.move_as_ok(), client_only, SelfId, adnl, std::move(promise)) .release(); }); @@ -478,6 +478,21 @@ void DhtMemberImpl::get_value_in(DhtKeyId key, td::Promise result) { get_self_node(std::move(P)); } +void DhtMemberImpl::get_value_many(DhtKey key, std::function callback, td::Promise promise) { + DhtKeyId key_id = key.compute_key_id(); + auto P = td::PromiseCreator::lambda( + [key = key_id, callback = std::move(callback), promise = std::move(promise), SelfId = actor_id(this), + print_id = print_id(), adnl = adnl_, list = get_nearest_nodes(key_id, k_ * 2), k = k_, a = a_, + network_id = network_id_, id = id_, client_only = client_only_](td::Result R) mutable { + R.ensure(); + td::actor::create_actor("FindValueManyQuery", key, print_id, id, std::move(list), k, a, + network_id, R.move_as_ok(), client_only, SelfId, adnl, + std::move(callback), std::move(promise)) + .release(); + }); + get_self_node(std::move(P)); +} + void DhtMemberImpl::register_reverse_connection(adnl::AdnlNodeIdFull client, td::Promise promise) { auto client_short = client.compute_short_id(); td::uint32 ttl = (td::uint32)td::Clocks::system() + 300; diff --git a/dht/dht.h b/dht/dht.h index b9c65c8a7..5abff94a1 100644 --- a/dht/dht.h +++ b/dht/dht.h @@ -53,6 +53,7 @@ class Dht : public td::actor::Actor { virtual void set_value(DhtValue key_value, td::Promise result) = 0; virtual void get_value(DhtKey key, td::Promise result) = 0; + virtual void get_value_many(DhtKey key, std::function callback, td::Promise promise) = 0; virtual void register_reverse_connection(adnl::AdnlNodeIdFull client, td::Promise promise) = 0; virtual void request_reverse_ping(adnl::AdnlNode target, adnl::AdnlNodeIdShort client, diff --git a/overlay/overlay.cpp b/overlay/overlay.cpp index f964ccc5c..43b3b7e5a 100644 --- a/overlay/overlay.cpp +++ b/overlay/overlay.cpp @@ -283,11 +283,14 @@ void OverlayImpl::alarm() { } if (next_dht_query_ && next_dht_query_.is_in_past()) { next_dht_query_ = td::Timestamp::never(); - auto P = td::PromiseCreator::lambda([SelfId = actor_id(this)](td::Result res) { - td::actor::send_closure(SelfId, &OverlayImpl::receive_dht_nodes, std::move(res), true); - }); - td::actor::send_closure(dht_node_, &dht::Dht::get_value, dht::DhtKey{overlay_id_.pubkey_hash(), "nodes", 0}, - std::move(P)); + std::function callback = [SelfId = actor_id(this)](dht::DhtValue value) { + td::actor::send_closure(SelfId, &OverlayImpl::receive_dht_nodes, std::move(value)); + }; + td::Promise on_finish = [SelfId = actor_id(this)](td::Result R) { + td::actor::send_closure(SelfId, &OverlayImpl::dht_lookup_finished, R.move_as_status()); + }; + td::actor::send_closure(dht_node_, &dht::Dht::get_value_many, dht::DhtKey{overlay_id_.pubkey_hash(), "nodes", 0}, + std::move(callback), std::move(on_finish)); } if (update_db_at_.is_in_past()) { if (peers_.size() > 0) { @@ -311,30 +314,30 @@ void OverlayImpl::alarm() { } } -void OverlayImpl::receive_dht_nodes(td::Result res, bool dummy) { +void OverlayImpl::receive_dht_nodes(dht::DhtValue v) { CHECK(public_); - if (res.is_ok()) { - auto v = res.move_as_ok(); - auto R = fetch_tl_object(v.value().clone(), true); - if (R.is_ok()) { - auto r = R.move_as_ok(); - VLOG(OVERLAY_INFO) << this << ": received " << r->nodes_.size() << " nodes from overlay"; - VLOG(OVERLAY_EXTRA_DEBUG) << this << ": nodes: " << ton_api::to_string(r); - std::vector nodes; - for (auto &n : r->nodes_) { - auto N = OverlayNode::create(n); - if (N.is_ok()) { - nodes.emplace_back(N.move_as_ok()); - } + auto R = fetch_tl_object(v.value().clone(), true); + if (R.is_ok()) { + auto r = R.move_as_ok(); + VLOG(OVERLAY_INFO) << this << ": received " << r->nodes_.size() << " nodes from overlay"; + VLOG(OVERLAY_EXTRA_DEBUG) << this << ": nodes: " << ton_api::to_string(r); + std::vector nodes; + for (auto &n : r->nodes_) { + auto N = OverlayNode::create(n); + if (N.is_ok()) { + nodes.emplace_back(N.move_as_ok()); } - add_peers(std::move(nodes)); - } else { - VLOG(OVERLAY_WARNING) << this << ": incorrect value in DHT for overlay nodes: " << R.move_as_error(); } + add_peers(std::move(nodes)); } else { - VLOG(OVERLAY_NOTICE) << this << ": can not get value from DHT: " << res.move_as_error(); + VLOG(OVERLAY_WARNING) << this << ": incorrect value in DHT for overlay nodes: " << R.move_as_error(); } +} +void OverlayImpl::dht_lookup_finished(td::Status S) { + if (S.is_error()) { + VLOG(OVERLAY_NOTICE) << this << ": can not get value from DHT: " << S; + } if (!(next_dht_store_query_ && next_dht_store_query_.is_in_past())) { finish_dht_query(); return; diff --git a/overlay/overlay.hpp b/overlay/overlay.hpp index 90fcc43d7..8fb3d91d7 100644 --- a/overlay/overlay.hpp +++ b/overlay/overlay.hpp @@ -166,7 +166,8 @@ class OverlayImpl : public Overlay { certs_[key] = std::move(cert); } - void receive_dht_nodes(td::Result res, bool dummy); + void receive_dht_nodes(dht::DhtValue v); + void dht_lookup_finished(td::Status S); void update_dht_nodes(OverlayNode node); void update_neighbours(td::uint32 nodes_to_change); diff --git a/tdutils/td/utils/Status.h b/tdutils/td/utils/Status.h index 9c4f1b7df..8bc210dba 100644 --- a/tdutils/td/utils/Status.h +++ b/tdutils/td/utils/Status.h @@ -555,6 +555,12 @@ class Result { }; return status_.move_as_error_suffix(suffix); } + Status move_as_status() TD_WARN_UNUSED_RESULT { + if (status_.is_error()) { + return move_as_error(); + } + return Status::OK(); + } const T &ok() const { LOG_CHECK(status_.is_ok()) << status_; return value_;