feat: pre-emptive stream creation for protocols

packages/core/src/index.ts

@@ -30,3 +30,4 @@ export { waitForRemotePeer } from "./lib/wait_for_remote_peer.js";
 export { ConnectionManager } from "./lib/connection_manager.js";
 
 export { KeepAliveManager } from "./lib/keep_alive_manager.js";
+export { StreamManager } from "./lib/stream_manager.js";

packages/core/src/lib/base_protocol.ts

@@ -1,13 +1,11 @@
 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
@@ -16,17 +14,28 @@ import {
 export class BaseProtocol implements IBaseProtocol {
   public readonly addLibp2pEventListener: Libp2p["addEventListener"];
   public readonly removeLibp2pEventListener: Libp2p["removeEventListener"];
+  protected streamManager: StreamManager;
 
   constructor(
     public multicodec: string,
-    private components: Libp2pComponents
+    private components: Libp2pComponents,
+    log: debug.Debugger
   ) {
     this.addLibp2pEventListener = components.events.addEventListener.bind(
       components.events
     );
     this.removeLibp2pEventListener = components.events.removeEventListener.bind(
       components.events
     );
+
+    this.streamManager = new StreamManager(
+      multicodec,
+      components.connectionManager.getConnections.bind(
+        components.connectionManager
+      ),
+      this.addLibp2pEventListener,
+      log
+    );
   }
 
   public get peerStore(): PeerStore {
@@ -50,15 +59,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,11 @@ 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.streamManager.getStream.bind(this, peer)
+        )
       );
 
     return subscription;

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

@@ -44,7 +44,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 || {};
   }
 
@@ -103,7 +103,7 @@ class LightPush extends BaseProtocol implements ILightPush {
 
     let error: undefined | SendError = undefined;
     const peer = await this.getPeer(opts?.peerId);
-    const stream = await this.newStream(peer);
+    const stream = await this.streamManager.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.streamManager.getStream.bind(this, peer),
       queryOpts,
       decodersAsMap,
       options?.cursor

packages/core/src/lib/stream_manager.ts

@@ -0,0 +1,68 @@
+import type { PeerUpdate } from "@libp2p/interface";
+import type { Stream } from "@libp2p/interface/connection";
+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 streamPool: Map<string, Promise<Stream>>;
+
+  constructor(
+    public multicodec: string,
+    public getConnections: Libp2p["getConnections"],
+    public addEventListener: Libp2p["addEventListener"],
+    private log: debug.Debugger
+  ) {
+    this.addEventListener(
+      "peer:update",
+      this.handlePeerUpdateStreamPool.bind(this)
+    );
+    this.getStream = this.getStream.bind(this);
+    this.streamPool = new Map();
+  }
+
+  public async getStream(peer: Peer): Promise<Stream> {
+    const peerIdStr = peer.id.toString();
+    const streamPromise = this.streamPool.get(peerIdStr);
+
+    if (!streamPromise) {
+      return this.newStream(peer); // fallback by creating a new stream on the spot
+    }
+
+    // We have the stream, let's remove it from the map
+    this.streamPool.delete(peerIdStr);
+
+    this.prepareNewStream(peer);
+
+    const stream = await streamPromise;
+
+    if (stream.status === "closed") {
+      return this.newStream(peer); // fallback by creating a new stream on the spot
+    }
+
+    return stream;
+  }
+
+  private async newStream(peer: Peer): Promise<Stream> {
+    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);
+  }
+
+  private prepareNewStream(peer: Peer): void {
+    const streamPromise = this.newStream(peer);
+    this.streamPool.set(peer.id.toString(), streamPromise);
+  }
+
+  private handlePeerUpdateStreamPool = (evt: CustomEvent<PeerUpdate>): void => {
+    const peer = evt.detail.peer;
+    if (peer.protocols.includes(this.multicodec)) {
+      this.log(`Optimistically opening a stream to ${peer.id.toString()}`);
+      this.prepareNewStream(peer);
+    }
+  };
+}

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.streamManager.getStream(peer);
 
     const res = await pipe(
       [rpcQuery.encode()],