Suspend event deliveries until middlewares are ready (#552)

Author: zikaari

.changeset/ten-lions-look.md

@@ -0,0 +1,5 @@
+---
+'@segment/analytics-next': patch
+---
+
+Suspend event deliveries until middlewares are ready

packages/browser/package.json

@@ -41,7 +41,7 @@
   "size-limit": [
     {
       "path": "dist/umd/index.js",
-      "limit": "25.7 KB"
+      "limit": "25.85 KB"
     }
   ],
   "dependencies": {

packages/browser/src/core/analytics/index.ts

@@ -354,19 +354,22 @@ export class Analytics extends Emitter {
   }
 
   async addSourceMiddleware(fn: MiddlewareFunction): Promise<Analytics> {
-    const { sourceMiddlewarePlugin } = await import(
-      /* webpackChunkName: "middleware" */ '../../plugins/middleware'
-    )
+    await this.queue.criticalTasks.run(async () => {
+      const { sourceMiddlewarePlugin } = await import(
+        /* webpackChunkName: "middleware" */ '../../plugins/middleware'
+      )
 
-    const integrations: Record<string, boolean> = {}
-    this.queue.plugins.forEach((plugin) => {
-      if (plugin.type === 'destination') {
-        return (integrations[plugin.name] = true)
-      }
+      const integrations: Record<string, boolean> = {}
+      this.queue.plugins.forEach((plugin) => {
+        if (plugin.type === 'destination') {
+          return (integrations[plugin.name] = true)
+        }
+      })
+
+      const plugin = sourceMiddlewarePlugin(fn, integrations)
+      await this.register(plugin)
     })
 
-    const plugin = sourceMiddlewarePlugin(fn, integrations)
-    await this.register(plugin)
     return this
   }
 

packages/browser/src/core/inspector/__tests__/index.test.ts

@@ -19,7 +19,6 @@ describe('Inspector interface', () => {
     const deliveryPromise = analytics.track('Test event').catch(() => {})
 
     expect(window.__SEGMENT_INSPECTOR__?.triggered).toHaveBeenCalledTimes(1)
-    expect(window.__SEGMENT_INSPECTOR__?.enriched).toHaveBeenCalledTimes(1)
 
     expect(window.__SEGMENT_INSPECTOR__?.triggered).toHaveBeenCalledWith(
       expect.objectContaining({
@@ -33,6 +32,7 @@ describe('Inspector interface', () => {
 
     await deliveryPromise
 
+    expect(window.__SEGMENT_INSPECTOR__?.enriched).toHaveBeenCalledTimes(1)
     expect(window.__SEGMENT_INSPECTOR__?.delivered).toHaveBeenCalledTimes(1)
   })
 })

packages/browser/src/core/queue/__tests__/event-queue.test.ts

@@ -1,4 +1,5 @@
 /* eslint-disable @typescript-eslint/no-floating-promises */
+import { noop } from 'lodash'
 import { Analytics } from '../../analytics'
 import { pWhile } from '../../../lib/p-while'
 import * as timer from '../../../lib/priority-queue/backoff'
@@ -9,6 +10,7 @@ import {
 import { Context, ContextCancelation } from '../../context'
 import { Plugin } from '../../plugin'
 import { EventQueue } from '../event-queue'
+import { pTimeout } from '../../callback'
 
 async function flushAll(eq: EventQueue): Promise<Context[]> {
   const flushSpy = jest.spyOn(eq, 'flush')
@@ -164,6 +166,37 @@ describe('Flushing', () => {
     expect(eq.queue.getAttempts(fruitBasket)).toEqual(2)
   })
 
+  test('waits for critical tasks to finish before performing event deliveries', async () => {
+    jest.useRealTimers()
+
+    const eq = new EventQueue()
+
+    let finishCriticalTask: () => void = noop
+    const startTask = () =>
+      new Promise<void>((res) => (finishCriticalTask = res))
+
+    // some preceding events that've been scheduled
+    const p1 = eq.dispatch(fruitBasket)
+    const p2 = eq.dispatch(basketView)
+    // a critical task has been kicked off
+    eq.criticalTasks.run(startTask)
+    // a succeeding event
+    const p3 = eq.dispatch(shopper)
+
+    // even after a good amount of time, none of the events should be delivered
+    await expect(pTimeout(Promise.race([p1, p2, p3]), 1000)).rejects.toThrow()
+
+    // give the green light
+    finishCriticalTask()
+
+    // now that the task is complete, the delivery should resume
+    expect(await Promise.all([p1, p2, p3])).toMatchObject([
+      fruitBasket,
+      basketView,
+      shopper,
+    ])
+  })
+
   test('delivers events on retry', async () => {
     jest.useRealTimers()
 

packages/browser/src/core/queue/event-queue.ts

@@ -7,6 +7,7 @@ import { Context, ContextCancelation } from '../context'
 import { Emitter } from '@segment/analytics-core'
 import { Integrations } from '../events'
 import { Plugin } from '../plugin'
+import { createTaskGroup, TaskGroup } from '../task/task-group'
 import { attempt, ensure } from './delivery'
 import { inspectorHost } from '../inspector'
 
@@ -18,6 +19,14 @@ type PluginsByType = {
 }
 
 export class EventQueue extends Emitter {
+  /**
+   * All event deliveries get suspended until all the tasks in this task group are complete.
+   * For example: a middleware that augments the event object should be loaded safely as a
+   * critical task, this way, event queue will wait for it to be ready before sending events.
+   *
+   * This applies to all the events already in the queue, and the upcoming ones
+   */
+  criticalTasks: TaskGroup = createTaskGroup()
   queue: PriorityQueue<Context>
   plugins: Plugin[] = []
   failedInitializations: string[] = []
@@ -152,6 +161,8 @@ export class EventQueue extends Emitter {
   }
 
   private async deliver(ctx: Context): Promise<Context> {
+    await this.criticalTasks.done()
+
     const start = Date.now()
     try {
       ctx = await this.flushOne(ctx)

packages/browser/src/core/task/task-group.ts

@@ -0,0 +1,31 @@
+import { isThenable } from '../../lib/is-thenable'
+
+export type TaskGroup = {
+  done: () => Promise<void>
+  run: <Operation extends (...args: any[]) => any>(
+    op: Operation
+  ) => ReturnType<Operation>
+}
+
+export const createTaskGroup = (): TaskGroup => {
+  let taskCompletionPromise: Promise<void>
+  let resolvePromise: () => void
+  let count = 0
+
+  return {
+    done: () => taskCompletionPromise,
+    run: (op) => {
+      const returnValue = op()
+
+      if (isThenable(returnValue)) {
+        if (++count === 1) {
+          taskCompletionPromise = new Promise((res) => (resolvePromise = res))
+        }
+
+        returnValue.finally(() => --count === 0 && resolvePromise())
+      }
+
+      return returnValue
+    },
+  }
+}