feat: optimistically create streams for light protocols

packages/core/src/lib/base_protocol.ts

@@ -1,32 +1,49 @@
 import type { Libp2p } from "@libp2p/interface";
-import type { Stream } from "@libp2p/interface/connection";
 import type { PeerId } from "@libp2p/interface/peer-id";
 import { Peer, PeerStore } from "@libp2p/interface/peer-store";
 import type { IBaseProtocol, Libp2pComponents } from "@waku/interfaces";
-import {
-  getPeersForProtocol,
-  selectConnection,
-  selectPeerForProtocol
-} from "@waku/utils/libp2p";
+import { getPeersForProtocol, selectPeerForProtocol } from "@waku/utils/libp2p";
+import debug from "debug";
+
+import { StreamManager } from "./stream_manager.js";
 
 /**
  * A class with predefined helpers, to be used as a base to implement Waku
  * Protocols.
  */
-export class BaseProtocol implements IBaseProtocol {
+export class BaseProtocol extends StreamManager implements IBaseProtocol {
   public readonly addLibp2pEventListener: Libp2p["addEventListener"];
   public readonly removeLibp2pEventListener: Libp2p["removeEventListener"];
 
   constructor(
     public multicodec: string,
-    private components: Libp2pComponents
+    private components: Libp2pComponents,
+    log: debug.Debugger
   ) {
+    super(
+      multicodec,
+      components.connectionManager.getConnections.bind(
+        components.connectionManager
+      ),
+      log
+    );
+
     this.addLibp2pEventListener = components.events.addEventListener.bind(
       components.events
     );
     this.removeLibp2pEventListener = components.events.removeEventListener.bind(
       components.events
     );
+
+    this.addLibp2pEventListener(
+      "peer:update",
+      super.handlePeerUpdateStreamPool.bind(this)
+    );
+    // TODO: might be better to check with `connection:close` event
+    this.addLibp2pEventListener(
+      "peer:disconnect",
+      super.handlePeerDisconnectStreamPool.bind(this)
+    );
   }
 
   public get peerStore(): PeerStore {
@@ -50,15 +67,4 @@ export class BaseProtocol implements IBaseProtocol {
     );
     return peer;
   }
-  protected async newStream(peer: Peer): Promise<Stream> {
-    const connections = this.components.connectionManager.getConnections(
-      peer.id
-    );
-    const connection = selectConnection(connections);
-    if (!connection) {
-      throw new Error("Failed to get a connection to the peer");
-    }
-
-    return connection.newStream(this.multicodec);
-  }
 }

packages/core/src/lib/filter/index.ts

@@ -246,7 +246,7 @@ class Filter extends BaseProtocol implements IReceiver {
   }
 
   constructor(libp2p: Libp2p, options?: ProtocolCreateOptions) {
-    super(FilterCodecs.SUBSCRIBE, libp2p.components);
+    super(FilterCodecs.SUBSCRIBE, libp2p.components, log);
 
     libp2p.handle(FilterCodecs.PUSH, this.onRequest.bind(this)).catch((e) => {
       log("Failed to register ", FilterCodecs.PUSH, e);
@@ -271,7 +271,7 @@ class Filter extends BaseProtocol implements IReceiver {
       this.setActiveSubscription(
         _pubSubTopic,
         peer.id.toString(),
-        new Subscription(_pubSubTopic, peer, this.newStream.bind(this, peer))
+        new Subscription(_pubSubTopic, peer, this.getStream.bind(this, peer))
       );
 
     return subscription;

packages/core/src/lib/light_push/index.ts

@@ -34,7 +34,7 @@ class LightPush extends BaseProtocol implements ILightPush {
   options: ProtocolCreateOptions;
 
   constructor(libp2p: Libp2p, options?: ProtocolCreateOptions) {
-    super(LightPushCodec, libp2p.components);
+    super(LightPushCodec, libp2p.components, log);
     this.options = options || {};
   }
 
@@ -99,7 +99,7 @@ class LightPush extends BaseProtocol implements ILightPush {
     }
 
     const peer = await this.getPeer(opts?.peerId);
-    const stream = await this.newStream(peer);
+    const stream = await this.getStream(peer);
 
     try {
       const res = await pipe(

packages/core/src/lib/store/index.ts

@@ -82,7 +82,7 @@ class Store extends BaseProtocol implements IStore {
   options: ProtocolCreateOptions;
 
   constructor(libp2p: Libp2p, options?: ProtocolCreateOptions) {
-    super(StoreCodec, libp2p.components);
+    super(StoreCodec, libp2p.components, log);
     this.options = options ?? {};
   }
 
@@ -254,7 +254,7 @@ class Store extends BaseProtocol implements IStore {
     const peer = await this.getPeer(options?.peerId);
 
     for await (const messages of paginate<T>(
-      this.newStream.bind(this, peer),
+      this.getStream.bind(this, peer),
       queryOpts,
       decodersAsMap,
       options?.cursor

packages/core/src/lib/stream_manager.ts

@@ -0,0 +1,126 @@
+import type { PeerUpdate } from "@libp2p/interface";
+import type { Stream } from "@libp2p/interface/connection";
+import type { PeerId } from "@libp2p/interface/peer-id";
+import { Peer } from "@libp2p/interface/peer-store";
+import { Libp2p } from "@waku/interfaces";
+import { selectConnection } from "@waku/utils/libp2p";
+import debug from "debug";
+
+export class StreamManager {
+  private streamsPool: Map<string, Stream[]> = new Map();
+  private ongoingStreamCreations: Map<string, Promise<Stream>> = new Map();
+  private readonly MAX_STREAMS_PER_PEER = 3;
+  private readonly MAX_RETRIES = 3;
+
+  constructor(
+    public multicodec: string,
+    public getConnections: Libp2p["getConnections"],
+    private log: debug.Debugger
+  ) {}
+
+  private async newStream(
+    peer: Peer,
+    retries = this.MAX_RETRIES
+  ): Promise<Stream> {
+    try {
+      const connections = this.getConnections(peer.id);
+      const connection = selectConnection(connections);
+      if (!connection) {
+        throw new Error("Failed to get a connection to the peer");
+      }
+      return connection.newStream(this.multicodec);
+    } catch (error) {
+      if (retries > 0) {
+        this.log(`Retrying stream creation. Retries left: ${retries}`);
+        return this.newStream(peer, retries - 1);
+      }
+      throw error;
+    }
+  }
+
+  protected async getStream(peer: Peer): Promise<Stream> {
+    const peerIdStr = peer.id.toString();
+
+    const ongoingCreation = this.ongoingStreamCreations.get(peerIdStr);
+    if (ongoingCreation) {
+      // Wait for the ongoing stream creation to complete
+      await ongoingCreation;
+    }
+
+    const peerStreams = this.streamsPool.get(peerIdStr) || [];
+    if (peerStreams.length > 0) {
+      const stream = peerStreams.pop();
+      if (!stream) {
+        throw new Error("Failed to get a stream from the pool");
+      }
+      this.replenishStreamPool(peer);
+      return stream;
+    }
+
+    this.log(`No stream available for peer ${peerIdStr}. Opening a new one.`);
+    return this.createAndSaveStream(peer);
+  }
+
+  private async createAndSaveStream(peer: Peer): Promise<Stream> {
+    const peerIdStr = peer.id.toString();
+
+    const streamCreationPromise = (async () => {
+      const stream = await this.newStream(peer);
+      const peerStreams = this.streamsPool.get(peerIdStr) || [];
+      peerStreams.push(stream);
+      this.streamsPool.set(peerIdStr, peerStreams);
+      return stream;
+    })();
+
+    this.ongoingStreamCreations.set(peerIdStr, streamCreationPromise);
+
+    try {
+      return await streamCreationPromise;
+    } finally {
+      this.ongoingStreamCreations.delete(peerIdStr);
+    }
+  }
+
+  private replenishStreamPool(peer: Peer): void {
+    const peerIdStr = peer.id.toString();
+    const ongoingCreationsCount = this.ongoingStreamCreations.has(peerIdStr)
+      ? 1
+      : 0;
+    const availableStreamsCount = (this.streamsPool.get(peerIdStr) || [])
+      .length;
+
+    if (
+      ongoingCreationsCount + availableStreamsCount <
+      this.MAX_STREAMS_PER_PEER
+    ) {
+      this.createAndSaveStream(peer)
+        .then(() => {
+          this.log(`Replenished stream pool for peer ${peer.id.toString()}`);
+        })
+        .catch((err) => {
+          this.log(
+            `Error replenishing stream pool for peer ${peer.id.toString()}: ${err}`
+          );
+        });
+    }
+  }
+
+  protected handlePeerUpdateStreamPool(evt: CustomEvent<PeerUpdate>): void {
+    const peer = evt.detail.peer;
+    if (peer.protocols.includes(this.multicodec)) {
+      const peerIdStr = peer.id.toString();
+      if (!this.streamsPool.has(peerIdStr)) {
+        this.streamsPool.set(peerIdStr, []);
+      }
+      this.log(`Optimistically opening a stream to ${peer.id.toString()}`);
+      this.replenishStreamPool(peer);
+    }
+  }
+
+  protected handlePeerDisconnectStreamPool(evt: CustomEvent<PeerId>): void {
+    const peerId = evt.detail;
+    this.streamsPool.delete(peerId.toString());
+    // Cancel ongoing stream creation if any
+    this.ongoingStreamCreations.delete(peerId.toString());
+  }
+}

packages/peer-exchange/src/waku_peer_exchange.ts

@@ -29,7 +29,7 @@ export class WakuPeerExchange extends BaseProtocol implements IPeerExchange {
    * @param components - libp2p components
    */
   constructor(components: Libp2pComponents) {
-    super(PeerExchangeCodec, components);
+    super(PeerExchangeCodec, components, log);
     this.multicodec = PeerExchangeCodec;
   }
 
@@ -47,7 +47,7 @@ export class WakuPeerExchange extends BaseProtocol implements IPeerExchange {
 
     const peer = await this.getPeer(params.peerId);
 
-    const stream = await this.newStream(peer);
+    const stream = await this.getStream(peer);
 
     const res = await pipe(
       [rpcQuery.encode()],