fix: cleanup routines on reliable channel and core protocols (#2733)

* fix: add stop methods to protocols to prevent event listener leaks

* fix: add abort signal support for graceful store query cancellation

* fix: call protocol stop methods in WakuNode.stop()

* fix: improve QueryOnConnect cleanup and abort signal handling

* fix: improve MissingMessageRetriever cleanup with abort signal

* fix: add stopAllRetries method to RetryManager for proper cleanup

* fix: implement comprehensive ReliableChannel stop() with proper cleanup

* fix: add active query tracking to QueryOnConnect and await its stop()

* fix: add stop() to IRelayAPI and IStore interfaces, implement in SDK wrappers

* align with usual naming (isStarted)

* remove unnecessary `await`

* test: `stop()` is now async

* chore: use more concise syntax

---------

Co-authored-by: Levente Kiss <[email protected]>

Author: fryorcraken

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

@@ -61,6 +61,7 @@ export class FilterCore {
   }
 
   public async stop(): Promise<void> {
+    this.streamManager.stop();
     try {
       await this.libp2p.unhandle(FilterCodecs.PUSH);
     } catch (e) {

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

@@ -33,6 +33,11 @@ export class LightPushCore {
     this.streamManager = new StreamManager(CODECS.v3, libp2p.components);
   }
 
+  public stop(): void {
+    this.streamManager.stop();
+    this.streamManagerV2.stop();
+  }
+
   public async send(
     encoder: IEncoder,
     message: IMessage,

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

@@ -35,6 +35,10 @@ export class StoreCore {
     this.streamManager = new StreamManager(StoreCodec, libp2p.components);
   }
 
+  public stop(): void {
+    this.streamManager.stop();
+  }
+
   public get maxTimeLimit(): number {
     return MAX_TIME_RANGE;
   }
@@ -68,6 +72,11 @@ export class StoreCore {
 
     let currentCursor = queryOpts.paginationCursor;
     while (true) {
+      if (queryOpts.abortSignal?.aborted) {
+        log.info("Store query aborted by signal");
+        break;
+      }
+
       const storeQueryRequest = StoreQueryRequest.create({
         ...queryOpts,
         paginationCursor: currentCursor
@@ -89,13 +98,22 @@ export class StoreCore {
         break;
       }
 
-      const res = await pipe(
-        [storeQueryRequest.encode()],
-        lp.encode,
-        stream,
-        lp.decode,
-        async (source) => await all(source)
-      );
+      let res;
+      try {
+        res = await pipe(
+          [storeQueryRequest.encode()],
+          lp.encode,
+          stream,
+          lp.decode,
+          async (source) => await all(source)
+        );
+      } catch (error) {
+        if (error instanceof Error && error.name === "AbortError") {
+          log.info(`Store query aborted for peer ${peerId.toString()}`);
+          break;
+        }
+        throw error;
+      }
 
       const bytes = new Uint8ArrayList();
       res.forEach((chunk) => {
@@ -122,6 +140,11 @@ export class StoreCore {
         `${storeQueryResponse.messages.length} messages retrieved from store`
       );
 
+      if (queryOpts.abortSignal?.aborted) {
+        log.info("Store query aborted by signal before processing messages");
+        break;
+      }
+
       const decodedMessages = storeQueryResponse.messages.map((protoMsg) => {
         if (!protoMsg.message) {
           return Promise.resolve(undefined);

packages/core/src/lib/stream_manager/stream_manager.ts

@@ -23,6 +23,15 @@ export class StreamManager {
     );
   }
 
+  public stop(): void {
+    this.libp2p.events.removeEventListener(
+      "peer:update",
+      this.handlePeerUpdateStreamPool
+    );
+    this.streamPool.clear();
+    this.ongoingCreation.clear();
+  }
+
   public async getStream(peerId: PeerId): Promise<Stream | undefined> {
     try {
       const peerIdStr = peerId.toString();

packages/interfaces/src/relay.ts

@@ -16,6 +16,7 @@ export interface IRelayAPI {
   readonly pubsubTopics: Set<PubsubTopic>;
   readonly gossipSub: GossipSub;
   start: () => Promise<void>;
+  stop: () => Promise<void>;
   waitForPeers: () => Promise<void>;
   getMeshPeers: (topic?: TopicStr) => PeerIdStr[];
 }

packages/interfaces/src/store.ts

@@ -88,11 +88,18 @@ export type QueryRequestParams = {
    * Only use if you know what you are doing.
    */
   peerId?: PeerId;
+
+  /**
+   * An optional AbortSignal to cancel the query.
+   * When the signal is aborted, the query will stop processing and return early.
+   */
+  abortSignal?: AbortSignal;
 };
 
 export type IStore = {
   readonly multicodec: string;
 
+  stop(): void;
   createCursor(message: IDecodedMessage): StoreCursor;
   queryGenerator: <T extends IDecodedMessage>(
     decoders: IDecoder<T>[],

packages/relay/src/relay.ts

@@ -67,6 +67,10 @@ export class Relay implements IRelay {
    * Observers under key `""` are always called.
    */
   private observers: Map<PubsubTopic, Map<ContentTopic, Set<unknown>>>;
+  private messageEventHandlers: Map<
+    PubsubTopic,
+    (event: CustomEvent<GossipsubMessage>) => void
+  > = new Map();
 
   public constructor(params: RelayConstructorParams) {
     if (!this.isRelayPubsub(params.libp2p.services.pubsub)) {
@@ -105,6 +109,19 @@ export class Relay implements IRelay {
     this.subscribeToAllTopics();
   }
 
+  public async stop(): Promise<void> {
+    for (const pubsubTopic of this.pubsubTopics) {
+      const handler = this.messageEventHandlers.get(pubsubTopic);
+      if (handler) {
+        this.gossipSub.removeEventListener("gossipsub:message", handler);
+      }
+      this.gossipSub.topicValidators.delete(pubsubTopic);
+      this.gossipSub.unsubscribe(pubsubTopic);
+    }
+    this.messageEventHandlers.clear();
+    this.observers.clear();
+  }
+
   /**
    * Wait for at least one peer with the given protocol to be connected and in the gossipsub
    * mesh for all pubsubTopics.
@@ -299,17 +316,17 @@ export class Relay implements IRelay {
    * @override
    */
   private gossipSubSubscribe(pubsubTopic: string): void {
-    this.gossipSub.addEventListener(
-      "gossipsub:message",
-      (event: CustomEvent<GossipsubMessage>) => {
-        if (event.detail.msg.topic !== pubsubTopic) return;
-
-        this.processIncomingMessage(
-          event.detail.msg.topic,
-          event.detail.msg.data
-        ).catch((e) => log.error("Failed to process incoming message", e));
-      }
-    );
+    const handler = (event: CustomEvent<GossipsubMessage>): void => {
+      if (event.detail.msg.topic !== pubsubTopic) return;
+
+      this.processIncomingMessage(
+        event.detail.msg.topic,
+        event.detail.msg.data
+      ).catch((e) => log.error("Failed to process incoming message", e));
+    };
+
+    this.messageEventHandlers.set(pubsubTopic, handler);
+    this.gossipSub.addEventListener("gossipsub:message", handler);
 
     this.gossipSub.topicValidators.set(pubsubTopic, messageValidator);
     this.gossipSub.subscribe(pubsubTopic);

packages/sdk/src/light_push/light_push.ts

@@ -65,6 +65,7 @@ export class LightPush implements ILightPush {
 
   public stop(): void {
     this.retryManager.stop();
+    this.protocol.stop();
   }
 
   public async send(

packages/sdk/src/query_on_connect/query_on_connect.spec.ts

@@ -158,14 +158,14 @@ describe("QueryOnConnect", () => {
       expect(wakuEventSpy.calledWith(WakuEvent.Health)).to.be.true;
     });
 
-    it("should remove event listeners when stopped", () => {
+    it("should remove event listeners when stopped", async () => {
       const peerRemoveSpy =
         mockPeerManagerEventEmitter.removeEventListener as sinon.SinonSpy;
       const wakuRemoveSpy =
         mockWakuEventEmitter.removeEventListener as sinon.SinonSpy;
 
       queryOnConnect.start();
-      queryOnConnect.stop();
+      await queryOnConnect.stop();
 
       expect(peerRemoveSpy.calledWith(PeerManagerEventNames.StoreConnect)).to.be
         .true;

packages/sdk/src/query_on_connect/query_on_connect.ts

@@ -52,6 +52,13 @@ export class QueryOnConnect<
   private lastTimeOffline: number;
   private readonly forceQueryThresholdMs: number;
 
+  private isStarted: boolean = false;
+  private abortController?: AbortController;
+  private activeQueryPromise?: Promise<void>;
+
+  private boundStoreConnectHandler?: (event: CustomEvent<PeerId>) => void;
+  private boundHealthHandler?: (event: CustomEvent<HealthStatus>) => void;
+
   public constructor(
     public decoders: IDecoder<T>[],
     public stopIfTrue: (msg: T) => boolean,
@@ -71,11 +78,37 @@ export class QueryOnConnect<
   }
 
   public start(): void {
+    if (this.isStarted) {
+      log.warn("QueryOnConnect already running");
+      return;
+    }
     log.info("starting query-on-connect service");
+    this.isStarted = true;
+    this.abortController = new AbortController();
     this.setupEventListeners();
   }
 
-  public stop(): void {
+  public async stop(): Promise<void> {
+    if (!this.isStarted) {
+      return;
+    }
+    log.info("stopping query-on-connect service");
+    this.isStarted = false;
+
+    if (this.abortController) {
+      this.abortController.abort();
+      this.abortController = undefined;
+    }
+
+    if (this.activeQueryPromise) {
+      log.info("Waiting for active query to complete...");
+      try {
+        await this.activeQueryPromise;
+      } catch (error) {
+        log.warn("Active query failed during stop:", error);
+      }
+    }
+
     this.unsetEventListeners();
   }
 
@@ -107,7 +140,10 @@ export class QueryOnConnect<
       this.lastTimeOffline > this.lastSuccessfulQuery ||
       timeSinceLastQuery > this.forceQueryThresholdMs
     ) {
-      await this.query(peerId);
+      this.activeQueryPromise = this.query(peerId).finally(() => {
+        this.activeQueryPromise = undefined;
+      });
+      await this.activeQueryPromise;
     } else {
       log.info(`no querying`);
     }
@@ -120,7 +156,8 @@ export class QueryOnConnect<
       for await (const page of this._queryGenerator(this.decoders, {
         timeStart,
         timeEnd,
-        peerId
+        peerId,
+        abortSignal: this.abortController?.signal
       })) {
         // Await for decoding
         const messages = (await Promise.all(page)).filter(
@@ -166,33 +203,41 @@ export class QueryOnConnect<
   }
 
   private setupEventListeners(): void {
+    this.boundStoreConnectHandler = (event: CustomEvent<PeerId>) => {
+      void this.maybeQuery(event.detail).catch((err) =>
+        log.error("query-on-connect error", err)
+      );
+    };
+
+    this.boundHealthHandler = this.updateLastOfflineDate.bind(this);
+
     this.peerManagerEventEmitter.addEventListener(
       PeerManagerEventNames.StoreConnect,
-      (event) =>
-        void this.maybeQuery(event.detail).catch((err) =>
-          log.error("query-on-connect error", err)
-        )
+      this.boundStoreConnectHandler
     );
 
     this.wakuEventEmitter.addEventListener(
       WakuEvent.Health,
-      this.updateLastOfflineDate.bind(this)
+      this.boundHealthHandler
     );
   }
 
   private unsetEventListeners(): void {
-    this.peerManagerEventEmitter.removeEventListener(
-      PeerManagerEventNames.StoreConnect,
-      (event) =>
-        void this.maybeQuery(event.detail).catch((err) =>
-          log.error("query-on-connect error", err)
-        )
-    );
+    if (this.boundStoreConnectHandler) {
+      this.peerManagerEventEmitter.removeEventListener(
+        PeerManagerEventNames.StoreConnect,
+        this.boundStoreConnectHandler
+      );
+      this.boundStoreConnectHandler = undefined;
+    }
 
-    this.wakuEventEmitter.removeEventListener(
-      WakuEvent.Health,
-      this.updateLastOfflineDate.bind(this)
-    );
+    if (this.boundHealthHandler) {
+      this.wakuEventEmitter.removeEventListener(
+        WakuEvent.Health,
+        this.boundHealthHandler
+      );
+      this.boundHealthHandler = undefined;
+    }
   }
 
   private updateLastOfflineDate(event: CustomEvent<HealthStatus>): void {