Fix in-flight being unable to track event requests made by cloned clients

packages/core/lib/clone-client.js

@@ -1,5 +1,7 @@
 const assign = require('./es-utils/assign')
 
+const onCloneCallbacks = []
+
 module.exports = (client) => {
   const clone = new client.Client({}, {}, [], client._notifier)
 
@@ -25,5 +27,9 @@ module.exports = (client) => {
   clone._delivery = client._delivery
   clone._sessionDelegate = client._sessionDelegate
 
+  onCloneCallbacks.forEach(callback => { callback(clone) })
+
   return clone
 }
+
+module.exports.registerCallback = callback => { onCloneCallbacks.push(callback) }

packages/core/test/lib/clone-client.test.ts

@@ -287,4 +287,91 @@ describe('@bugsnag/core/lib/clone-client', () => {
       expect(cloned._sessionDelegate).toBe(original._sessionDelegate)
     })
   })
+
+  describe('registerCallback', () => {
+    it('allows registering callbacks that are called when a client is cloned', () => {
+      const callback1 = jest.fn()
+      const callback2 = jest.fn()
+      const callback3 = jest.fn()
+
+      clone.registerCallback(callback1)
+      clone.registerCallback(callback2)
+      clone.registerCallback(callback3)
+
+      const original = new Client({ apiKey })
+
+      expect(callback1).not.toHaveBeenCalled()
+      expect(callback2).not.toHaveBeenCalled()
+      expect(callback3).not.toHaveBeenCalled()
+
+      clone(original)
+
+      expect(callback1).toHaveBeenCalledTimes(1)
+      expect(callback2).toHaveBeenCalledTimes(1)
+      expect(callback3).toHaveBeenCalledTimes(1)
+
+      clone(original)
+
+      expect(callback1).toHaveBeenCalledTimes(2)
+      expect(callback2).toHaveBeenCalledTimes(2)
+      expect(callback3).toHaveBeenCalledTimes(2)
+    })
+
+    it('passes the clone to the callbacks', () => {
+      const callback1 = jest.fn()
+      const callback2 = jest.fn()
+      const callback3 = jest.fn()
+
+      clone.registerCallback(callback1)
+      clone.registerCallback(callback2)
+      clone.registerCallback(callback3)
+
+      const original = new Client({ apiKey })
+
+      expect(callback1).not.toHaveBeenCalled()
+      expect(callback2).not.toHaveBeenCalled()
+      expect(callback3).not.toHaveBeenCalled()
+
+      const cloned = clone(original)
+
+      expect(callback1).toHaveBeenCalledWith(cloned)
+      expect(callback2).toHaveBeenCalledWith(cloned)
+      expect(callback3).toHaveBeenCalledWith(cloned)
+
+      const cloned2 = clone(original)
+
+      expect(callback1).toHaveBeenCalledWith(cloned2)
+      expect(callback2).toHaveBeenCalledWith(cloned2)
+      expect(callback3).toHaveBeenCalledWith(cloned2)
+
+      // the callbacks should not be called with the original client
+      expect(callback1).not.toHaveBeenCalledWith(original)
+      expect(callback2).not.toHaveBeenCalledWith(original)
+      expect(callback3).not.toHaveBeenCalledWith(original)
+    })
+
+    it('calls callbacks in the order they are registered', () => {
+      const order: string[] = []
+
+      const callback1 = () => { order.push('callback1') }
+      const callback2 = () => { order.push('callback2') }
+      const callback3 = () => { order.push('callback3') }
+
+      clone.registerCallback(callback1)
+      clone.registerCallback(callback2)
+      clone.registerCallback(callback3)
+
+      const original = new Client({ apiKey })
+
+      expect(order).toEqual([])
+
+      clone(original)
+
+      expect(order).toEqual(['callback1', 'callback2', 'callback3'])
+
+      clone(original)
+
+      expect(order).toEqual(['callback1', 'callback2', 'callback3', 'callback1', 'callback2', 'callback3'])
+    })
+  })
 })

packages/in-flight/src/in-flight.js

@@ -1,47 +1,19 @@
 const cuid = require('@bugsnag/cuid')
+const clone = require('@bugsnag/core/lib/clone-client')
 
 const FLUSH_POLL_INTERVAL_MS = 50
 const inFlightRequests = new Map()
 
 const noop = () => {}
 
+// when a client is cloned, make sure to patch the clone's notify method too
+// we don't need to patch delivery when a client is cloned because the
+// original client's delivery method will be copied over to the clone
+clone.registerCallback(patchNotify)
+
 module.exports = {
   trackInFlight (client) {
-    const originalNotify = client._notify
-
-    client._notify = function (event, onError, callback = noop) {
-      const id = cuid()
-      inFlightRequests.set(id, true)
-
-      const _callback = function () {
-        inFlightRequests.delete(id)
-        callback.apply(null, arguments)
-      }
-
-      client._depth += 1
-
-      try {
-        originalNotify.call(client, event, onError, _callback)
-      } finally {
-        client._depth -= 1
-      }
-    }
-
-    const patchDelivery = (delivery) => {
-      const originalSendSession = delivery.sendSession
-
-      delivery.sendSession = function (session, callback = noop) {
-        const id = cuid()
-        inFlightRequests.set(id, true)
-
-        const _callback = function () {
-          inFlightRequests.delete(id)
-          callback.apply(null, arguments)
-        }
-
-        originalSendSession.call(delivery, session, _callback)
-      }
-    }
+    patchNotify(client)
 
     let delivery = client._delivery
     patchDelivery(delivery)
@@ -85,3 +57,52 @@ module.exports = {
     })
   }
 }
+
+// patch a client's _notify method to track in-flight requests
+// we patch _notify directly to track requests as early as possible and use the
+// "post report" delivery callback to know when a request finishes
+function patchNotify (client) {
+  const originalNotify = client._notify
+
+  client._notify = function (event, onError, callback = noop) {
+    const id = cuid()
+    inFlightRequests.set(id, true)
+
+    const _callback = function () {
+      inFlightRequests.delete(id)
+      callback.apply(null, arguments)
+    }
+
+    client._depth += 1
+
+    try {
+      originalNotify.call(client, event, onError, _callback)
+    } finally {
+      client._depth -= 1
+    }
+  }
+}
+
+// patch a delivery delegate's sendSession method to track in-flight requests
+// we do this on the delivery delegate because the client object doesn't
+// actually deliver sessions, a session delegate does
+// we can't patch the session delegate either because it will deliver sessions
+// in a way that makes sense on the platform, e.g. on node sessions are batched
+// into 1 request made every x seconds
+// therefore the only thing that knows when a session request is started and
+// when it finishes is the delivery delegate itself
+function patchDelivery (delivery) {
+  const originalSendSession = delivery.sendSession
+
+  delivery.sendSession = function (session, callback = noop) {
+    const id = cuid()
+    inFlightRequests.set(id, true)
+
+    const _callback = function () {
+      inFlightRequests.delete(id)
+      callback.apply(null, arguments)
+    }
+
+    originalSendSession.call(delivery, session, _callback)
+  }
+}

packages/in-flight/test/in-flight.test.ts

@@ -1,3 +1,4 @@
+import clone from '@bugsnag/core/lib/clone-client'
 import Client, { EventDeliveryPayload, SessionDeliveryPayload } from '@bugsnag/core/client'
 
 // The in-flight package has module level state which can leak between tests
@@ -43,6 +44,44 @@ describe('@bugsnag/in-flight', () => {
     expect(sendSession).not.toHaveBeenCalled()
   })
 
+  it('can track in-flight events after a client is cloned', () => {
+    const client = new Client({ apiKey: 'AN_API_KEY' })
+
+    // eslint thinks this is never reassigned, but it clearly is
+    let cloned: Client // eslint-disable-line prefer-const
+
+    const payloads: EventDeliveryPayload[] = []
+    const sendSession = jest.fn()
+
+    client._setDelivery(() => ({
+      sendEvent: (payload, cb) => {
+        expect(cloned._depth).toBe(2)
+        payloads.push(payload)
+        cb()
+      },
+      sendSession
+    }))
+
+    bugsnagInFlight.trackInFlight(client)
+
+    expect(payloads.length).toBe(0)
+
+    const onError = jest.fn()
+    const callback = jest.fn()
+
+    cloned = clone(client)
+
+    expect(cloned._depth).toBe(1)
+
+    cloned.notify(new Error('xyz'), onError, callback)
+
+    expect(cloned._depth).toBe(1)
+    expect(onError).toHaveBeenCalledTimes(1)
+    expect(callback).toHaveBeenCalledTimes(1)
+    expect(payloads.length).toBe(1)
+    expect(sendSession).not.toHaveBeenCalled()
+  })
+
   it('tracks in-flight sessions', () => {
     const client = new Client({ apiKey: 'AN_API_KEY' })
     const payloads: SessionDeliveryPayload[] = []
@@ -84,6 +123,49 @@ describe('@bugsnag/in-flight', () => {
     expect(client._sessionDelegate.resumeSession).not.toHaveBeenCalled()
   })
 
+  it('tracks in-flight sessions after a client has been cloned', () => {
+    const client = new Client({ apiKey: 'AN_API_KEY' })
+    const payloads: SessionDeliveryPayload[] = []
+    const sendEvent = jest.fn()
+    const callback = jest.fn()
+
+    client._sessionDelegate = {
+      startSession: jest.fn(function (client, session) {
+        client._delivery.sendSession(session, callback)
+
+        return client
+      }),
+      pauseSession: jest.fn(),
+      resumeSession: jest.fn()
+    }
+
+    client._setDelivery(() => ({
+      sendEvent,
+      sendSession: (payload, cb) => {
+        payloads.push(payload)
+        cb()
+      }
+    }))
+
+    bugsnagInFlight.trackInFlight(client)
+
+    expect(payloads.length).toBe(0)
+    expect(callback).not.toHaveBeenCalled()
+    expect(client._sessionDelegate.startSession).not.toHaveBeenCalled()
+    expect(client._sessionDelegate.pauseSession).not.toHaveBeenCalled()
+    expect(client._sessionDelegate.resumeSession).not.toHaveBeenCalled()
+
+    const cloned = clone(client)
+
+    cloned.startSession()
+
+    expect(payloads.length).toBe(1)
+    expect(callback).toHaveBeenCalledTimes(1)
+    expect(cloned._sessionDelegate.startSession).toHaveBeenCalledTimes(1)
+    expect(cloned._sessionDelegate.pauseSession).not.toHaveBeenCalled()
+    expect(cloned._sessionDelegate.resumeSession).not.toHaveBeenCalled()
+  })
+
   it('tracks all in-flight requests', () => {
     const client = new Client({ apiKey: 'AN_API_KEY' })
     const eventPayloads: EventDeliveryPayload[] = []