From a3769caffb89b047f6df64425b01a6ed17851086 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20K=C3=B6cher?= <git@kchr.de>
Date: Thu, 5 Jan 2023 22:11:27 +0100
Subject: [PATCH] Kademlia: Speed-up the record fetching (#13081)

Before libp2p 0.50.0 we used a quorum of one to fetch records from the DHT. In the pr that upgraded
to libp2p 0.50.0 we accidentally changed this behavior. This pr brings back the old behavior of
using a qorum of one and thus, a faster discovery. After finding the first value, we directly finish
the query. There was also another behavior change in libp2p, they stopped automatic caching on
remote nodes. This pr also brings back the remote caching on nodes that are nearest to the key from
our point of view of the network.

The pr that changed the behavior in libp2p: https://github.com/libp2p/rust-libp2p/pull/2712
---
 client/network/src/discovery.rs | 75 ++++++++++++++++++++++-----------
 1 file changed, 50 insertions(+), 25 deletions(-)

diff --git a/client/network/src/discovery.rs b/client/network/src/discovery.rs
index 16dbc7662c192..7abc723c4047e 100644
--- a/client/network/src/discovery.rs
+++ b/client/network/src/discovery.rs
@@ -249,7 +249,7 @@ impl DiscoveryConfig {
 				NonZeroUsize::new(MAX_KNOWN_EXTERNAL_ADDRESSES)
 					.expect("value is a constant; constant is non-zero; qed."),
 			),
-			outbound_query_records: Vec::new(),
+			records_to_publish: Default::default(),
 		}
 	}
 }
@@ -287,8 +287,12 @@ pub struct DiscoveryBehaviour {
 	allow_non_globals_in_dht: bool,
 	/// A cache of discovered external addresses. Only used for logging purposes.
 	known_external_addresses: LruHashSet<Multiaddr>,
-	/// A cache of outbound query records.
-	outbound_query_records: Vec<(record::Key, Vec<u8>)>,
+	/// Records to publish per QueryId.
+	///
+	/// After finishing a Kademlia query, libp2p will return us a list of the closest peers that
+	/// did not return the record(in `FinishedWithNoAdditionalRecord`). We will then put the record
+	/// to these peers.
+	records_to_publish: HashMap<QueryId, Record>,
 }
 
 impl DiscoveryBehaviour {
@@ -692,33 +696,54 @@ impl NetworkBehaviour for DiscoveryBehaviour {
 					KademliaEvent::OutboundQueryProgressed {
 						result: QueryResult::GetRecord(res),
 						stats,
-						step,
+						id,
 						..
 					} => {
 						let ev = match res {
-							Ok(ok) =>
-								if let GetRecordOk::FoundRecord(r) = ok {
-									self.outbound_query_records
-										.push((r.record.key, r.record.value));
+							Ok(GetRecordOk::FoundRecord(r)) => {
+								debug!(
+									target: "sub-libp2p",
+									"Libp2p => Found record ({:?}) with value: {:?}",
+									r.record.key,
+									r.record.value,
+								);
+
+								// Let's directly finish the query, as we are only interested in a
+								// quorum of 1.
+								if let Some(kad) = self.kademlia.as_mut() {
+									if let Some(mut query) = kad.query_mut(&id) {
+										query.finish();
+									}
+								}
+
+								self.records_to_publish.insert(id, r.record.clone());
+
+								DiscoveryOut::ValueFound(
+									vec![(r.record.key, r.record.value)],
+									stats.duration().unwrap_or_default(),
+								)
+							},
+							Ok(GetRecordOk::FinishedWithNoAdditionalRecord {
+								cache_candidates,
+							}) => {
+								if cache_candidates.is_empty() {
 									continue
-								} else {
-									debug!(
-										target: "sub-libp2p",
-										"Libp2p => Query progressed to {:?} step (last: {:?})",
-										step.count,
-										step.last,
-									);
-									if step.last {
-										let records =
-											self.outbound_query_records.drain(..).collect();
-										DiscoveryOut::ValueFound(
-											records,
-											stats.duration().unwrap_or_default(),
-										)
-									} else {
-										continue
+								}
+
+								// Put the record to the `cache_candidates` that are nearest to the
+								// record key from our point of view of the network.
+								if let Some(record) = self.records_to_publish.remove(&id) {
+									if let Some(kad) = self.kademlia.as_mut() {
+										kad.put_record_to(
+											record,
+											cache_candidates.into_iter().map(|v| v.1),
+											Quorum::One,
+										);
 									}
-								},
+								}
+
+								continue
+							},
 							Err(e @ libp2p::kad::GetRecordError::NotFound { .. }) => {
 								trace!(
 									target: "sub-libp2p",