refactor: extract peer selection logic

src/lib/select_peer.ts

@@ -1,4 +1,8 @@
+import type { PeerId } from "@libp2p/interface-peer-id";
 import type { Peer, PeerStore } from "@libp2p/interface-peer-store";
+import debug from "debug";
+
+const log = debug("waku:select-peer");
 
 /**
  * Returns a pseudo-random peer that supports the given protocol.
@@ -29,3 +33,45 @@ export async function getPeersForProtocol(
   });
   return peers;
 }
+
+export async function selectPeerForProtocol(
+  peerStore: PeerStore,
+  protocols: string[],
+  peerId?: PeerId
+): Promise<{ peer: Peer; protocol: string } | undefined> {
+  let peer;
+  if (peerId) {
+    peer = await peerStore.get(peerId);
+    if (!peer) {
+      log(
+        `Failed to retrieve connection details for provided peer in peer store: ${peerId.toString()}`
+      );
+      return;
+    }
+  } else {
+    const peers = await getPeersForProtocol(peerStore, protocols);
+    peer = selectRandomPeer(peers);
+    if (!peer) {
+      log("Failed to find known peer that registers protocols", protocols);
+      return;
+    }
+  }
+
+  let protocol;
+  for (const codec of protocols) {
+    if (peer.protocols.includes(codec)) {
+      protocol = codec;
+      // Do not break as we want to keep the last value
+    }
+  }
+  log(`Using codec ${protocol}`);
+  if (!protocol) {
+    log(
+      `Peer does not register required protocols: ${peer.id.toString()}`,
+      protocols
+    );
+    return;
+  }
+
+  return { peer, protocol };
+}

src/lib/waku_filter/index.ts

@@ -11,7 +11,11 @@ import type { Libp2p } from "libp2p";
 import { WakuMessage as WakuMessageProto } from "../../proto/message";
 import { DefaultPubSubTopic } from "../constants";
 import { selectConnection } from "../select_connection";
-import { getPeersForProtocol, selectRandomPeer } from "../select_peer";
+import {
+  getPeersForProtocol,
+  selectPeerForProtocol,
+  selectRandomPeer,
+} from "../select_peer";
 import { hexToBytes } from "../utils";
 import { DecryptionMethod, WakuMessage } from "../waku_message";
 
@@ -228,23 +232,15 @@ export class WakuFilter {
   }
 
   private async getPeer(peerId?: PeerId): Promise<Peer> {
-    let peer;
+    const res = await selectPeerForProtocol(
-    if (peerId) {
+      this.libp2p.peerStore,
-      peer = await this.libp2p.peerStore.get(peerId);
+      [FilterCodec],
-      if (!peer) {
+      peerId
-        throw new Error(
-          `Failed to retrieve connection details for provided peer in peer store: ${peerId.toString()}`
     );
+    if (!res) {
+      throw new Error(`Failed to select peer for ${FilterCodec}`);
     }
-    } else {
+    return res.peer;
-      peer = await this.randomPeer();
-      if (!peer) {
-        throw new Error(
-          "Failed to find known peer that registers waku filter protocol"
-        );
-      }
-    }
-    return peer;
   }
 
   /**

src/lib/waku_light_push/index.ts

@@ -10,7 +10,11 @@ import { Uint8ArrayList } from "uint8arraylist";
 import { PushResponse } from "../../proto/light_push";
 import { DefaultPubSubTopic } from "../constants";
 import { selectConnection } from "../select_connection";
-import { getPeersForProtocol, selectRandomPeer } from "../select_peer";
+import {
+  getPeersForProtocol,
+  selectPeerForProtocol,
+  selectRandomPeer,
+} from "../select_peer";
 import { WakuMessage } from "../waku_message";
 
 import { PushRPC } from "./push_rpc";
@@ -51,16 +55,16 @@ export class WakuLightPush {
     message: WakuMessage,
     opts?: PushOptions
   ): Promise<PushResponse | null> {
-    let peer;
+    const res = await selectPeerForProtocol(
-    if (opts?.peerId) {
+      this.libp2p.peerStore,
-      peer = await this.libp2p.peerStore.get(opts.peerId);
+      [LightPushCodec],
-      if (!peer) throw "Peer is unknown";
+      opts?.peerId
-    } else {
+    );
-      peer = await this.randomPeer();
+
+    if (!res) {
+      throw new Error("Failed to get a peer");
     }
-    if (!peer) throw "No peer available";
+    const { peer } = res;
-    if (!peer.protocols.includes(LightPushCodec))
-      throw "Peer does not register waku light push protocol";
 
     const connections = this.libp2p.connectionManager.getConnections(peer.id);
     const connection = selectConnection(connections);

src/lib/waku_store/index.ts

@@ -11,7 +11,7 @@ import * as protoV2Beta4 from "../../proto/store_v2beta4";
 import { HistoryResponse } from "../../proto/store_v2beta4";
 import { DefaultPubSubTopic, StoreCodecs } from "../constants";
 import { selectConnection } from "../select_connection";
-import { getPeersForProtocol, selectRandomPeer } from "../select_peer";
+import { getPeersForProtocol, selectPeerForProtocol } from "../select_peer";
 import { hexToBytes } from "../utils";
 import {
   DecryptionMethod,
@@ -152,29 +152,17 @@ export class WakuStore {
       ...options,
     });
 
-    let peer;
+    const res = await selectPeerForProtocol(
-    if (opts.peerId) {
+      this.libp2p.peerStore,
-      peer = await this.libp2p.peerStore.get(opts.peerId);
+      Object.values(StoreCodecs),
-      if (!peer)
+      opts?.peerId
-        throw `Failed to retrieve connection details for provided peer in peer store: ${opts.peerId.toString()}`;
+    );
-    } else {
-      peer = await this.randomPeer();
-      if (!peer)
-        throw "Failed to find known peer that registers waku store protocol";
-    }
 
-    let storeCodec = "";
+    if (!res) {
-    for (const codec of Object.values(StoreCodecs)) {
+      throw new Error("Failed to get a peer");
-      if (peer.protocols.includes(codec)) {
-        storeCodec = codec;
-        // Do not break as we want to keep the last value
     }
-    }
+    const { peer, protocol } = res;
-    log(`Use store codec ${storeCodec}`);
-    if (!storeCodec)
-      throw `Peer does not register waku store protocol: ${peer.id.toString()}`;
 
-    Object.assign(opts, { storeCodec });
     const connections = this.libp2p.connectionManager.getConnections(peer.id);
     const connection = selectConnection(connections);
 
@@ -199,7 +187,7 @@ export class WakuStore {
     const messages: WakuMessage[] = [];
     let cursor = undefined;
     while (true) {
-      const stream = await connection.newStream(storeCodec);
+      const stream = await connection.newStream(protocol);
       const queryOpts = Object.assign(opts, { cursor });
       const historyRpcQuery = HistoryRPC.createQuery(queryOpts);
       log("Querying store peer", connections[0].remoteAddr.toString());
@@ -317,13 +305,4 @@ export class WakuStore {
 
     return getPeersForProtocol(this.libp2p.peerStore, codecs);
   }
-
-  /**
-   * Returns a random peer that supports store protocol from the address
-   * book (`libp2p.peerStore`). Waku may or  may not be currently connected to
-   * this peer.
-   */
-  async randomPeer(): Promise<Peer | undefined> {
-    return selectRandomPeer(await this.peers());
-  }
 }