Await heartbeat promises on shutdown

When AppSignal is "gracefully stopped" by calling `Appsignal.stop()`,
the heartbeats might not have finished sending. AppSignal should await
the delivery of those heartbeats before shutting down.

To do this, the pending heartbeats are collected into a set, which
automatically cleans itself up when all promises are settled. A promise
for the settling of all remaining promises in the set is awaited when
`Appsignal.stop()` is called.

The method that awaits these promises is named `Heartbeat.shutdown()`,
and not `Heartbeat.stop()`, to avoid any potential confusion where it
could be mistaken for the counterpart of `Heartbeat.start()`.

This is a sort of breaking change, in that `Appsignal.stop()` is now
an asynchronous method, meaning that it returns a promise, and as
such it must be awaited.

However, note that `Appsignal.stop()` should have already been
awaiting the promise returned by the OpenTelemetry SDK's `.shutdown()`
method. As such, this change makes sense beyond the heartbeats.

Author: unflxw

.changesets/appsignal-stop-must-be-awaited.md

@@ -0,0 +1,25 @@
+---
+bump: "patch"
+type: "change"
+---
+
+`Appsignal.stop()` now returns a promise. For your application to wait until
+AppSignal has been gracefully stopped, this promise must be awaited:
+
+```javascript
+import { Appsignal } from "@appsignal/nodejs"
+
+await Appsignal.stop()
+process.exit(0)
+```
+
+In older Node.js versions where top-level await is not available, terminate
+the application when the promise is settled:
+
+```javascript
+import { Appsignal } from "@appsignal/nodejs"
+
+Appsignal.stop().finally(() => {
+  process.exit(0)
+})
+```

src/__tests__/client.test.ts

@@ -23,8 +23,8 @@ describe("Client", () => {
     client = new Client({ ...DEFAULT_OPTS })
   })
 
-  afterEach(() => {
-    client.stop()
+  afterEach(async () => {
+    await client.stop()
   })
 
   it("starts the client", () => {
@@ -33,18 +33,18 @@ describe("Client", () => {
     expect(startSpy).toHaveBeenCalled()
   })
 
-  it("stops the client", () => {
+  it("stops the client", async () => {
     const extensionStopSpy = jest.spyOn(Extension.prototype, "stop")
-    client.stop()
+    await client.stop()
     expect(extensionStopSpy).toHaveBeenCalled()
   })
 
-  it("stops the probes when the client is active", () => {
+  it("stops the probes when the client is active", async () => {
     client = new Client({ ...DEFAULT_OPTS, active: true })
     const probes = client.metrics().probes()
     expect(probes.isRunning).toEqual(true)
 
-    client.stop()
+    await client.stop()
     expect(probes.isRunning).toEqual(false)
   })
 

src/__tests__/span_processor.test.ts

@@ -35,9 +35,9 @@ describe("Span processor", () => {
     tracer = tracerProvider.getTracer("unknown-instrumentation")
   })
 
-  afterEach(() => {
-    tracerProvider.shutdown()
-    client.stop()
+  afterEach(async () => {
+    await tracerProvider.shutdown()
+    await client.stop()
   })
 
   it("processes unknown OpenTelemetry spans", () => {

src/__tests__/winston_transport.test.ts

@@ -20,8 +20,8 @@ describe("BaseLogger", () => {
     })
   })
 
-  afterAll(() => {
-    client.stop()
+  afterAll(async () => {
+    await client.stop()
   })
 
   beforeEach(() => {

src/client.ts

@@ -43,6 +43,7 @@ import {
   LayerType as RestifyLayerType
 } from "@opentelemetry/instrumentation-restify"
 import { SpanProcessor, TestModeSpanProcessor } from "./span_processor"
+import { Heartbeat } from "./heartbeat"
 
 const DefaultInstrumentations = {
   "@opentelemetry/instrumentation-express": ExpressInstrumentation,
@@ -206,16 +207,17 @@ export class Client {
    * Call this before the end of your program to make sure the
    * agent is stopped as well.
    */
-  public stop(calledBy?: string): void {
+  public async stop(calledBy?: string): Promise<void> {
     if (calledBy) {
       console.log(`Stopping AppSignal (${calledBy})`)
     } else {
       console.log("Stopping AppSignal")
     }
 
-    this.#sdk?.shutdown()
+    await this.#sdk?.shutdown()
     this.metrics().probes().stop()
     this.extension.stop()
+    await Heartbeat.shutdown()
   }
 
   /**

src/heartbeat.ts

@@ -11,7 +11,21 @@ export type Event = {
   timestamp: number
 }
 
+class PendingPromiseSet<T> extends Set<Promise<T>> {
+  add(promise: Promise<T>) {
+    super.add(promise)
+    promise.finally(() => this.delete(promise))
+    return this
+  }
+
+  async allSettled() {
+    await Promise.allSettled(this)
+  }
+}
+
 export class Heartbeat {
+  private static heartbeatPromises = new PendingPromiseSet<any>()
+
   name: string
   id: string
 
@@ -20,6 +34,10 @@ export class Heartbeat {
     this.id = crypto.randomBytes(8).toString("hex")
   }
 
+  public static async shutdown() {
+    await Heartbeat.heartbeatPromises.allSettled()
+  }
+
   public start() {
     this.transmit(this.event("start"))
   }
@@ -45,16 +63,26 @@ export class Heartbeat {
       return
     }
 
-    new Transmitter(
+    const promise = new Transmitter(
       `${Client.config.data.loggingEndpoint}/heartbeats/json`,
       JSON.stringify(event)
-    ).transmit().then(({status}: {status: number}) => {
-      if (status !== 200) {
-        Client.internalLogger.warn(`Failed to transmit heartbeat: status code ${status}`)
-      }
-    }).catch(({error}: {error: Error}) => {
-      Client.internalLogger.warn(`Failed to transmit heartbeat: ${error.message}`)
-    })
+    ).transmit()
+
+    const handledPromise = promise
+      .then(({ status }: { status: number }) => {
+        if (status !== 200) {
+          Client.internalLogger.warn(
+            `Failed to transmit heartbeat: status code ${status}`
+          )
+        }
+      })
+      .catch(({ error }: { error: Error }) => {
+        Client.internalLogger.warn(
+          `Failed to transmit heartbeat: ${error.message}`
+        )
+      })
+
+    Heartbeat.heartbeatPromises.add(handledPromise)
   }
 }