Lookup lower distances if there's space

src/main/java/org/ethereum/beacon/discovery/storage/KBuckets.java

@@ -23,7 +23,7 @@ public class KBuckets {
    * Minimum distance we create a bucket for. 0 is our local node record and negative distances
    * aren't allowed.
    */
-  private static final int MIN_BUCKET = 1;
+  public static final int MINIMUM_BUCKET = 1;
 
   /** Maximum distance we create a bucket for. This is enough to cover all 32 byte node IDs. */
   public static final int MAXIMUM_BUCKET = 256;
@@ -130,7 +130,7 @@ public synchronized void performMaintenance() {
   }
 
   private Optional<KBucket> getOrCreateBucket(final int distance) {
-    if (distance > MAXIMUM_BUCKET || distance < MIN_BUCKET) {
+    if (distance > MAXIMUM_BUCKET || distance < MINIMUM_BUCKET) {
       // Distance too great, ignore.
       return Optional.empty();
     }

src/main/java/org/ethereum/beacon/discovery/task/DiscoveryTaskManager.java

@@ -4,6 +4,9 @@
 
 package org.ethereum.beacon.discovery.task;
 
+import static com.google.common.base.Preconditions.checkArgument;
+
+import com.google.common.annotations.VisibleForTesting;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -26,6 +29,7 @@ public class DiscoveryTaskManager {
   public static final Duration DEFAULT_LIVE_CHECK_INTERVAL = Duration.ofSeconds(1);
   public static final Duration DEFAULT_RECURSIVE_LOOKUP_INTERVAL = Duration.ofSeconds(10);
   private static final int RECURSIVE_SEARCH_QUERY_LIMIT = 15;
+  private static final int LOOKUP_REQUEST_LIMIT = 4;
   private final Bytes homeNodeId;
   private final Scheduler scheduler;
   private final RecursiveLookupTasks recursiveLookupTasks;
@@ -96,14 +100,38 @@ private CompletableFuture<Collection<NodeRecord>> performSearchForNewPeers() {
         .execute();
   }
 
-  private CompletableFuture<Collection<NodeRecord>> findNodes(
-      final NodeRecord nodeRecord, final int targetDistance) {
+  @VisibleForTesting
+  static List<Integer> lookupDistances(int targetDistance) {
+    checkArgument(
+        targetDistance >= KBuckets.MINIMUM_BUCKET && targetDistance <= KBuckets.MAXIMUM_BUCKET,
+        "invalid target distance: %s",
+        targetDistance);
+
     final List<Integer> distances = new ArrayList<>();
-    for (int distance = targetDistance;
-        distance <= Math.min(KBuckets.MAXIMUM_BUCKET, targetDistance + 3);
-        distance++) {
-      distances.add(distance);
+    distances.add(targetDistance);
+
+    // We want to prioritize higher distances in the request over lower ones
+    // because lower distances are less likely to contain nodes.
+    for (int offset = 1;
+        targetDistance + offset <= KBuckets.MAXIMUM_BUCKET
+            && distances.size() < LOOKUP_REQUEST_LIMIT;
+        offset++) {
+      distances.add(targetDistance + offset);
+    }
+
+    // If there's any space left, add lower distances.
+    for (int offset = 1;
+        targetDistance - offset >= KBuckets.MINIMUM_BUCKET
+            && distances.size() < LOOKUP_REQUEST_LIMIT;
+        offset++) {
+      distances.add(targetDistance - offset);
     }
-    return recursiveLookupTasks.add(nodeRecord, distances);
+
+    return distances;
+  }
+
+  private CompletableFuture<Collection<NodeRecord>> findNodes(
+      final NodeRecord nodeRecord, final int targetDistance) {
+    return recursiveLookupTasks.add(nodeRecord, lookupDistances(targetDistance));
   }
 }

src/test/java/org/ethereum/beacon/discovery/task/DiscoveryTaskManagerTest.java

@@ -0,0 +1,65 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.ethereum.beacon.discovery.task;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import java.util.List;
+import org.ethereum.beacon.discovery.storage.KBuckets;
+import org.junit.jupiter.api.Test;
+
+public class DiscoveryTaskManagerTest {
+
+  @Test
+  void lookupDistancesShouldReturnFourValues() {
+    for (int distance = KBuckets.MINIMUM_BUCKET; distance <= KBuckets.MAXIMUM_BUCKET; distance++) {
+      assertThat(DiscoveryTaskManager.lookupDistances(distance)).hasSize(4);
+    }
+  }
+
+  @Test
+  void lookupDistancesShouldRejectTargetBelowMin() {
+    assertThatThrownBy(() -> DiscoveryTaskManager.lookupDistances(KBuckets.MINIMUM_BUCKET - 1))
+        .hasMessageContaining("invalid target distance: 0")
+        .isInstanceOf(IllegalArgumentException.class);
+  }
+
+  @Test
+  void lookupDistancesMinTarget() {
+    assertThat(DiscoveryTaskManager.lookupDistances(1)).isEqualTo(List.of(1, 2, 3, 4));
+  }
+
+  @Test
+  void lookupDistancesMiddleTarget() {
+    assertThat(DiscoveryTaskManager.lookupDistances(128)).isEqualTo(List.of(128, 129, 130, 131));
+  }
+
+  @Test
+  void lookupDistancesTargetDistance253() {
+    assertThat(DiscoveryTaskManager.lookupDistances(253)).isEqualTo(List.of(253, 254, 255, 256));
+  }
+
+  @Test
+  void lookupDistancesTargetDistance254() {
+    assertThat(DiscoveryTaskManager.lookupDistances(254)).isEqualTo(List.of(254, 255, 256, 253));
+  }
+
+  @Test
+  void lookupDistancesTargetDistance255() {
+    assertThat(DiscoveryTaskManager.lookupDistances(255)).isEqualTo(List.of(255, 256, 254, 253));
+  }
+
+  @Test
+  void lookupDistancesMaxTarget() {
+    assertThat(DiscoveryTaskManager.lookupDistances(256)).isEqualTo(List.of(256, 255, 254, 253));
+  }
+
+  @Test
+  void lookupDistancesShouldRejectTargetAboveMax() {
+    assertThatThrownBy(() -> DiscoveryTaskManager.lookupDistances(KBuckets.MAXIMUM_BUCKET + 1))
+        .hasMessageContaining("invalid target distance: 257")
+        .isInstanceOf(IllegalArgumentException.class);
+  }
+}