From 59b80b6d2165af2ee128af9b80ab16d9fe1fb341 Mon Sep 17 00:00:00 2001
From: iceseer <iceseer@gmail.com>
Date: Thu, 15 Aug 2024 10:32:29 +0300
Subject: [PATCH] force authority discovery for the given peer set

Signed-off-by: iceseer <iceseer@gmail.com>
---
 core/authority_discovery/query/query.hpp      |  5 ++++
 core/authority_discovery/query/query_impl.cpp | 13 +++++++----
 core/authority_discovery/query/query_impl.hpp |  1 +
 .../validator/impl/parachain_processor.cpp    | 23 ++++++++++++++++++-
 4 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/core/authority_discovery/query/query.hpp b/core/authority_discovery/query/query.hpp
index fe5eae810d..368a482755 100644
--- a/core/authority_discovery/query/query.hpp
+++ b/core/authority_discovery/query/query.hpp
@@ -29,5 +29,10 @@ namespace kagome::authority_discovery {
      */
     virtual std::optional<primitives::AuthorityDiscoveryId> get(
         const libp2p::peer::PeerId &peer_id) const = 0;
+
+    /**
+     * Updates authority discovery set with the givens.
+     */
+    virtual outcome::result<void> update(std::vector<primitives::AuthorityDiscoveryId> &&authorities) = 0;
   };
 }  // namespace kagome::authority_discovery
diff --git a/core/authority_discovery/query/query_impl.cpp b/core/authority_discovery/query/query_impl.cpp
index fd45592639..3508ea5fa2 100644
--- a/core/authority_discovery/query/query_impl.cpp
+++ b/core/authority_discovery/query/query_impl.cpp
@@ -135,18 +135,23 @@ namespace kagome::authority_discovery {
   }
 
   outcome::result<void> QueryImpl::update() {
-    std::unique_lock lock{mutex_};
     OUTCOME_TRY(
         authorities,
         authority_discovery_api_->authorities(block_tree_->bestBlock().hash));
+    return update(std::move(authorities));
+  }
+  
+  outcome::result<void> QueryImpl::update(std::vector<primitives::AuthorityDiscoveryId> &&authorities) {
+    OUTCOME_TRY(local_keys,
+                key_store_->sr25519().getPublicKeys(
+                    crypto::KeyTypes::AUTHORITY_DISCOVERY));
+
+    std::unique_lock lock{mutex_};
     for (auto &id : authorities) {
       if (not hash_to_auth_.contains(id)) {
         hash_to_auth_.emplace(crypto::sha256(id), id);
       }
     }
-    OUTCOME_TRY(local_keys,
-                key_store_->sr25519().getPublicKeys(
-                    crypto::KeyTypes::AUTHORITY_DISCOVERY));
     auto has = [](const std::vector<primitives::AuthorityDiscoveryId> &keys,
                   const primitives::AuthorityDiscoveryId &key) {
       return std::find(keys.begin(), keys.end(), key) != keys.end();
diff --git a/core/authority_discovery/query/query_impl.hpp b/core/authority_discovery/query/query_impl.hpp
index a4c6c55506..6e2f4a3940 100644
--- a/core/authority_discovery/query/query_impl.hpp
+++ b/core/authority_discovery/query/query_impl.hpp
@@ -69,6 +69,7 @@ namespace kagome::authority_discovery {
         const std::vector<libp2p::protocol::kademlia::Value> &) override;
 
     outcome::result<void> update();
+    outcome::result<void> update(std::vector<primitives::AuthorityDiscoveryId> &&authorities) override;
 
    private:
     struct Authority {
diff --git a/core/parachain/validator/impl/parachain_processor.cpp b/core/parachain/validator/impl/parachain_processor.cpp
index 83a3eae671..8e8f076b32 100644
--- a/core/parachain/validator/impl/parachain_processor.cpp
+++ b/core/parachain/validator/impl/parachain_processor.cpp
@@ -1113,11 +1113,32 @@ namespace kagome::parachain {
           session_info->validator_groups,
           grid::shuffle(session_info->discovery_keys.size(), randomness),
           *global_v_index);
+      Groups groups{session_info->validator_groups, minimum_backing_votes};
 
+      std::unordered_set<primitives::AuthorityDiscoveryId> peers;
+      if (validator_index) {
+        auto our_group = groups.byValidatorIndex(*validator_index);
+        if (our_group) {
+          const auto &group = session_info->validator_groups[*our_group];
+          for (const auto vi : group) {
+            peers.insert(session_info->discovery_keys[vi]);
+          }
+        }
+      }
+      for (const auto &view : grid_view) {
+        for (const auto vi : view.sending) {
+          peers.insert(session_info->discovery_keys[vi]);
+        }
+        for (const auto vi : view.receiving) {
+          peers.insert(session_info->discovery_keys[vi]);
+        }
+      }
+
+      std::ignore = query_audi_->update(std::vector<primitives::AuthorityDiscoveryId>{ peers.begin(), peers.end() });
       return RefCache<SessionIndex, PerSessionState>::RefObj(
           session_index,
           *session_info,
-          Groups{session_info->validator_groups, minimum_backing_votes},
+          std::move(groups),
           std::move(grid_view),
           validator_index,
           peer_use_count_);