fix(core): improve event listener cleanup

Closes #4186.

Before this fix, we weren't removing event listeners registered by the
plugin event brokers (mostly for the `_exit` and `_restart` control
events).

This is now done in a simple way via the new `onKey`  and `clearKey`
methods on the event bus, which facilitates removing all listeners
matching a given key.

We use `garden.sessionId` as the key now, which is not quite ideal
when running concurrent commands in `garden dev`, but we'll be assigning
a command-unique `sessionId` in an upcoming PR (which gives us precisely
the semantics we want here).

core/src/cli/command-line.ts

@@ -719,6 +719,7 @@ ${chalk.white.underline("Keys:")}
         this.flashError(getCmdFailMsg(name))
       })
       .finally(() => {
+        this.garden.events.clearKey(this.garden.sessionId)
         delete this.runningCommands[id]
         this.renderStatus()
       })

core/src/events.ts

@@ -30,12 +30,17 @@ export type GardenEventListener<T extends EventName> = (payload: Events[T]) => v
  * See below for the event interfaces.
  */
 export class EventBus extends EventEmitter2 {
-  constructor() {
+  private keyIndex: {
+    [key: string]: { [eventName: string]: ((payload: any) => void)[] }
+  }
+
+  constructor(name?: string) {
     super({
       wildcard: false,
       newListener: false,
       maxListeners: 5000, // we may need to adjust this
     })
+    this.keyIndex = {}
   }
 
   emit<T extends EventName>(name: T, payload: Events[T]) {
@@ -46,6 +51,53 @@ export class EventBus extends EventEmitter2 {
     return super.on(name, listener)
   }
 
+  /**
+   * Registers the listener under the provided key for easy cleanup via `offKey`. This is useful e.g. for the
+   * plugin event broker, which is instantiated in several places and where there isn't a single obvious place to
+   * remove listeners from all instances generated in a single command run.
+   */
+  onKey<T extends EventName>(name: T, listener: (payload: Events[T]) => void, key: string) {
+    if (!this.keyIndex[key]) {
+      this.keyIndex[key] = {}
+    }
+    if (!this.keyIndex[key][name]) {
+      this.keyIndex[key][name] = []
+    }
+    this.keyIndex[key][name].push(listener)
+    return super.on(name, listener)
+  }
+
+  /**
+   * Removes all event listeners for the event `name` that were registered under `key` (via `onKey`).
+   */
+  offKey<T extends EventName>(name: T, key: string) {
+    if (!this.keyIndex[key]) {
+      return
+    }
+    if (!this.keyIndex[key][name]) {
+      return
+    }
+    for (const listener of this.keyIndex[key][name]) {
+      this.removeListener(name, listener)
+    }
+    delete this.keyIndex[key][name]
+  }
+
+  /**
+   * Removes all event listeners that were registered under `key` (via `onKey`).
+   */
+  clearKey(key: string) {
+    if (!this.keyIndex[key]) {
+      return
+    }
+    for (const name of Object.keys(this.keyIndex[key])) {
+      for (const listener of this.keyIndex[key][name]) {
+        this.removeListener(name, listener)
+      }
+    }
+    delete this.keyIndex[key]
+  }
+
   onAny(listener: <T extends EventName>(name: T, payload: Events[T]) => void) {
     return super.onAny(<any>listener)
   }
@@ -260,8 +312,8 @@ export type EventName = keyof Events
 
 // Note: Does not include logger events.
 export const pipedEventNames: EventName[] = [
-  "_exit",
-  "_restart",
+  // "_exit",
+  // "_restart",
   "_test",
   "_workflowRunRegistered",
   "sessionCompleted",

core/src/plugin-context.ts

@@ -131,21 +131,22 @@ export class PluginEventBroker extends EventEmitter<PluginEvents, PluginEventTyp
   private done: boolean
   private failed: boolean
   private error: Error | undefined
+  private garden: Garden
+  private abortHandler: () => void
 
   constructor(garden: Garden) {
     super()
 
     this.aborted = false
     this.done = false
     this.failed = false
+    this.garden = garden
+    this.abortHandler = () => this.emit("abort")
 
+    // console.trace()
     // Always respond to exit and restart events
-    garden.events.on("_exit", () => {
-      this.emit("abort")
-    })
-    garden.events.on("_restart", () => {
-      this.emit("abort")
-    })
+    this.garden.events.onKey("_exit", this.abortHandler, garden.sessionId)
+    this.garden.events.onKey("_restart", this.abortHandler, garden.sessionId)
 
     this.on("abort", () => {
       this.aborted = true

core/src/server/server.ts

@@ -270,6 +270,7 @@ export class GardenServer extends EventEmitter {
         args,
         opts,
       })
+      this.garden.events.clearKey(this.garden.sessionId)
 
       if (result.errors?.length) {
         throw result.errors[0]

core/test/unit/src/events.ts

@@ -6,6 +6,7 @@
  * file, You can obtain one at http://mozilla.org/MPL/2.0/.
  */
 
+import { range } from "lodash"
 import { EventBus } from "../../../src/events"
 import { expect } from "chai"
 
@@ -45,4 +46,61 @@ describe("EventBus", () => {
       events.emit("_test", "bar")
     })
   })
+
+  describe("onKey", () => {
+    it("should add a listener under the specified key", (done) => {
+      const key = "gandalf"
+      events.onKey("_test", (payload) => {
+        expect(payload).to.equal("foo")
+        expect(events["keyIndex"][key]["_test"].length).to.eql(1)
+        done()
+      }, key)
+      events.emit("_test", "foo")
+    })
+  })
+
+  describe("offKey", () => {
+    it("should remove all listeners with the specified key and name", () => {
+      const key = "gandalf"
+      const otherKey = "blob"
+      for (const _i of range(3)) {
+        events.onKey("_test", () => {}, key)
+        events.onKey("_restart", () => {}, key)
+
+        events.onKey("_test", () => {}, otherKey)
+        events.onKey("_restart", () => {}, otherKey)
+      }
+      expect(events.listenerCount()).to.eql(12)
+      events.offKey("_test", key)
+      expect(events.listenerCount()).to.eql(9)
+      expect(events["keyIndex"][key]["_test"]).to.be.undefined
+
+      // We expect the index for other key + name combinations to be the same.
+      expect(events["keyIndex"][key]["_restart"].length).to.eql(3)
+      expect(events["keyIndex"][otherKey]["_test"].length).to.eql(3)
+      expect(events["keyIndex"][otherKey]["_restart"].length).to.eql(3)
+    })
+  })
+
+  describe("clearKey", () => {
+    it("should remove all listeners with the specified key", () => {
+      const key = "gandalf"
+      const otherKey = "blob"
+      for (const _i of range(3)) {
+        events.onKey("_test", () => {}, key)
+        events.onKey("_restart", () => {}, key)
+
+        events.onKey("_test", () => {}, otherKey)
+        events.onKey("_restart", () => {}, otherKey)
+      }
+      expect(events.listenerCount()).to.eql(12)
+      events.clearKey(key)
+      expect(events.listenerCount()).to.eql(6)
+      expect(events["keyIndex"][key]).to.be.undefined
+
+      // We expect the index for the other key to be the same.
+      expect(events["keyIndex"][otherKey]["_test"].length).to.eql(3)
+      expect(events["keyIndex"][otherKey]["_restart"].length).to.eql(3)
+    })
+  })
 })