Improve dht lookup in overlays (#1104)

Continue dht lookup even if value was found

dht/dht-in.hpp

@@ -179,6 +179,7 @@ class DhtMemberImpl : public DhtMember {
   void get_value(DhtKey key, td::Promise<DhtValue> result) override {
     get_value_in(key.compute_key_id(), std::move(result));
   }
+  void get_value_many(DhtKey key, std::function<void(DhtValue)> callback, td::Promise<td::Unit> promise) override;
 
   void alarm() override {
     alarm_timestamp() = td::Timestamp::in(1.0);

dht/dht-query.cpp

@@ -210,8 +210,11 @@ void DhtQueryFindValue::on_result(td::Result<td::BufferSlice> 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<td::BufferSlice> 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,

dht/dht-query.hpp

@@ -126,16 +126,11 @@ class DhtQueryFindNodes : public DhtQuery {
 };
 
 class DhtQueryFindValue : public DhtQuery {
- private:
-  td::Promise<DhtValue> 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<DhtMember> node, td::actor::ActorId<adnl::Adnl> adnl,
-                    td::Promise<DhtValue> 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<DhtMember> node, td::actor::ActorId<adnl::Adnl> 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<DhtMember> node, td::actor::ActorId<adnl::Adnl> adnl,
+                          td::Promise<DhtValue> 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<DhtValue> 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<DhtMember> node, td::actor::ActorId<adnl::Adnl> adnl,
+                         std::function<void(DhtValue)> callback, td::Promise<td::Unit> 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<void(DhtValue)> callback_;
+  td::Promise<td::Unit> promise_;
+  bool found_ = false;
 };
 
 class DhtQueryStore : public td::actor::Actor {

dht/dht.cpp

@@ -470,7 +470,7 @@ void DhtMemberImpl::get_value_in(DhtKeyId key, td::Promise<DhtValue> result) {
                                        network_id = network_id_, id = id_,
                                        client_only = client_only_](td::Result<DhtNode> R) mutable {
     R.ensure();
-    td::actor::create_actor<DhtQueryFindValue>("FindValueQuery", key, print_id, id, std::move(list), k, a, network_id,
+    td::actor::create_actor<DhtQueryFindValueSingle>("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<DhtValue> result) {
   get_self_node(std::move(P));
 }
 
+void DhtMemberImpl::get_value_many(DhtKey key, std::function<void(DhtValue)> callback, td::Promise<td::Unit> 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<DhtNode> R) mutable {
+        R.ensure();
+        td::actor::create_actor<DhtQueryFindValueMany>("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<td::Unit> promise) {
   auto client_short = client.compute_short_id();
   td::uint32 ttl = (td::uint32)td::Clocks::system() + 300;

dht/dht.h

@@ -53,6 +53,7 @@ class Dht : public td::actor::Actor {
 
   virtual void set_value(DhtValue key_value, td::Promise<td::Unit> result) = 0;
   virtual void get_value(DhtKey key, td::Promise<DhtValue> result) = 0;
+  virtual void get_value_many(DhtKey key, std::function<void(DhtValue)> callback, td::Promise<td::Unit> promise) = 0;
 
   virtual void register_reverse_connection(adnl::AdnlNodeIdFull client, td::Promise<td::Unit> promise) = 0;
   virtual void request_reverse_ping(adnl::AdnlNode target, adnl::AdnlNodeIdShort client,