fix: remove onCancel when worker is terminated (#9033)

Author: sheremet-va

docs/api/advanced/vitest.md

@@ -458,11 +458,15 @@ Register a handler that will be called when the server is restarted due to a con
 ## onCancel
 
 ```ts
-function onCancel(fn: (reason: CancelReason) => Awaitable<void>): void
+function onCancel(fn: (reason: CancelReason) => Awaitable<void>): () => void
 ```
 
 Register a handler that will be called when the test run is cancelled with [`vitest.cancelCurrentRun`](#cancelcurrentrun).
 
+::: warning EXPERIMENTAL
+Since 4.0.10, `onCancel` returns a teardown function that will remove the listener.
+:::
+
 ## onClose
 
 ```ts

packages/vitest/src/node/core.ts

@@ -139,7 +139,7 @@ export class Vitest {
   private _onRestartListeners: OnServerRestartHandler[] = []
   private _onClose: (() => Awaitable<void>)[] = []
   private _onSetServer: OnServerRestartHandler[] = []
-  private _onCancelListeners: ((reason: CancelReason) => Awaitable<void>)[] = []
+  private _onCancelListeners = new Set<(reason: CancelReason) => Awaitable<void>>()
   private _onUserTestsRerun: OnTestsRerunHandler[] = []
   private _onFilterWatchedSpecification: ((spec: TestSpecification) => boolean)[] = []
 
@@ -706,7 +706,7 @@ export class Vitest {
 
     // previous run
     await this.runningPromise
-    this._onCancelListeners = []
+    this._onCancelListeners.clear()
     this.isCancelling = false
 
     // schedule the new run
@@ -807,7 +807,7 @@ export class Vitest {
 
     // previous run
     await this.runningPromise
-    this._onCancelListeners = []
+    this._onCancelListeners.clear()
     this.isCancelling = false
 
     // schedule the new run
@@ -859,7 +859,7 @@ export class Vitest {
    */
   async cancelCurrentRun(reason: CancelReason): Promise<void> {
     this.isCancelling = true
-    await Promise.all(this._onCancelListeners.splice(0).map(listener => listener(reason)))
+    await Promise.all([...this._onCancelListeners].map(listener => listener(reason)))
     await this.runningPromise
   }
 
@@ -1257,8 +1257,11 @@ export class Vitest {
   /**
    * Register a handler that will be called when the test run is cancelled with `vitest.cancelCurrentRun`.
    */
-  onCancel(fn: (reason: CancelReason) => Awaitable<void>): void {
-    this._onCancelListeners.push(fn)
+  onCancel(fn: (reason: CancelReason) => Awaitable<void>): () => void {
+    this._onCancelListeners.add(fn)
+    return () => {
+      this._onCancelListeners.delete(fn)
+    }
   }
 
   /**

packages/vitest/src/node/pools/poolRunner.ts

@@ -37,6 +37,7 @@ export class PoolRunner {
     rpc: [unknown]
   }> = new EventEmitter()
 
+  private _offCancel: () => void
   private _rpc: BirpcReturn<RunnerRPC, RuntimeRPC>
 
   public get isTerminated(): boolean {
@@ -63,7 +64,7 @@ export class PoolRunner {
       },
     )
 
-    this.project.vitest.onCancel(reason => this._rpc.onCancel(reason))
+    this._offCancel = this.project.vitest.onCancel(reason => this._rpc.onCancel(reason))
   }
 
   postMessage(message: WorkerRequest): void {
@@ -181,6 +182,7 @@ export class PoolRunner {
       )
 
       this._eventEmitter.removeAllListeners()
+      this._offCancel()
       this._rpc.$close(new Error('[vitest-pool-runner]: Pending methods while closing rpc'))
 
       // Stop the worker process (this sets _fork/_thread to undefined)