Process events in batches

package.json

@@ -86,6 +86,6 @@
         }
     },
     "lint-staged": {
-        "*.{js,css,scss}": "prettier --write"
+        "*.{js,ts,css,scss}": "prettier --write"
     }
 }

src/plugins.ts

@@ -1,6 +1,6 @@
 import * as path from 'path'
 import * as fs from 'fs'
-import { createPluginConfigVM, prepareForRun } from './vm'
+import { createPluginConfigVM } from './vm'
 import { PluginsServer, PluginConfig, PluginJsonConfig, TeamId } from './types'
 import { PluginEvent, PluginAttachment } from 'posthog-plugins'
 import { clearError, processError } from './error'
@@ -180,26 +180,21 @@ async function loadPlugin(server: PluginsServer, pluginConfig: PluginConfig): Pr
 }
 
 export async function runPlugins(server: PluginsServer, event: PluginEvent): Promise<PluginEvent | null> {
-    const pluginsToRun = server.pluginConfigsPerTeam.get(event.team_id) || server.defaultConfigs
-
+    const pluginsToRun = getPluginsForTeam(server, event.team_id)
     let returnedEvent: PluginEvent | null = event
 
     for (const pluginConfig of pluginsToRun.reverse()) {
-        if (pluginConfig.vm) {
-            const processEvent = prepareForRun(server, event.team_id, pluginConfig, 'processEvent', event)
-
-            if (processEvent) {
-                const startTime = performance.now()
-                try {
-                    returnedEvent = (await processEvent(returnedEvent)) || null
-                    const ms = Math.round((performance.now() - startTime) * 1000) / 1000
-                    logTime(pluginConfig.plugin?.name || 'noname', ms)
-                } catch (error) {
-                    await processError(server, pluginConfig, error, returnedEvent)
-                    const ms = Math.round((performance.now() - startTime) * 1000) / 1000
-                    logTime(pluginConfig.plugin?.name || 'noname', ms, true)
-                }
+        if (pluginConfig.vm?.methods?.processEvent) {
+            let errored = false
+            const { processEvent } = pluginConfig.vm.methods
+            const startTime = performance.now()
+            try {
+                returnedEvent = (await processEvent(returnedEvent)) || null
+            } catch (error) {
+                errored = true
+                await processError(server, pluginConfig, error, returnedEvent)
             }
+            logTime(pluginConfig.plugin?.name || 'noname', performance.now() - startTime, errored)
 
             if (!returnedEvent) {
                 return null
@@ -209,3 +204,46 @@ export async function runPlugins(server: PluginsServer, event: PluginEvent): Pro
 
     return returnedEvent
 }
+
+export async function runPluginsOnBatch(server: PluginsServer, batch: PluginEvent[]): Promise<PluginEvent[]> {
+    const eventsByTeam = new Map<number, PluginEvent[]>()
+
+    for (const event of batch) {
+        if (eventsByTeam.has(event.team_id)) {
+            eventsByTeam.get(event.team_id)!.push(event)
+        } else {
+            eventsByTeam.set(event.team_id, [event])
+        }
+    }
+
+    let allReturnedEvents: PluginEvent[] = []
+
+    for (const [teamId, teamEvents] of eventsByTeam.entries()) {
+        const pluginsToRun = getPluginsForTeam(server, teamId)
+
+        let returnedEvents: PluginEvent[] = teamEvents
+
+        for (const pluginConfig of pluginsToRun.reverse()) {
+            const { processEventBatch } = pluginConfig.vm?.methods || {}
+            if (processEventBatch && returnedEvents.length > 0) {
+                const startTime = performance.now()
+                let errored = false
+                try {
+                    returnedEvents = (await processEventBatch(returnedEvents)) || []
+                } catch (error) {
+                    errored = true
+                    await processError(server, pluginConfig, error, returnedEvents[0])
+                }
+                logTime(pluginConfig.plugin?.name || 'noname', performance.now() - startTime, errored)
+            }
+        }
+
+        allReturnedEvents = allReturnedEvents.concat(returnedEvents)
+    }
+
+    return allReturnedEvents
+}
+
+function getPluginsForTeam(server: PluginsServer, teamId: number): PluginConfig[] {
+    return server.pluginConfigsPerTeam.get(teamId) || server.defaultConfigs
+}

src/stats.ts

@@ -3,7 +3,7 @@ const histories = new Map<string, Array<number>>()
 const historyIndex = new Map<string, number>()
 
 export function logTime(name: string, time: number, error?: boolean): void {
-    const ms = Math.round(time * 1000) / 1000
+    // const ms = Math.round(time * 1000) / 1000
     // TODO: add this back with better dev logging. Disabling since this trashes performance tests.
     // console.log(`Running plugin ${name}: ${error ? 'ERROR IN ' : ''}${ms}ms`)
     if (!histories.has(name)) {

src/types.ts

@@ -113,5 +113,6 @@ export interface PluginConfigVMReponse {
     vm: VM
     methods: {
         processEvent: (event: PluginEvent) => Promise<PluginEvent>
+        processEventBatch: (batch: PluginEvent[]) => Promise<PluginEvent[]>
     }
 }

src/vm.ts

@@ -65,9 +65,32 @@ export function createPluginConfigVM(
         // run the plugin setup script, if present
         __callWithMeta('setupPlugin');
 
+        // we have processEvent, but not processEventBatch
+        if (!__getExported('processEventBatch') && __getExported('processEvent')) {
+            exports.processEventBatch = async function processEventBatch (batch, meta) {
+                const processEvent = __getExported('processEvent');
+                let waitFor = false
+                const processedEvents = batch.map(event => {
+                    const e = processEvent(event, meta)
+                    if (e && typeof e.then !== 'undefined') {
+                        waitFor = true
+                    }
+                    return e
+                })
+                const response = waitFor ? (await Promise.all(processedEvents)) : processedEvents;
+                return response.filter(r => r)
+            }
+        // we have processEventBatch, but not processEvent
+        } else if (!__getExported('processEvent') && __getExported('processEventBatch')) {
+            exports.processEvent = async function processEvent (event, meta) {
+                return (await (__getExported('processEventBatch'))([event], meta))?.[0]
+            }
+        }
+
         // export various functions
         const __methods = {
-            processEvent: __bindMeta('processEvent')
+            processEvent: __bindMeta('processEvent'),
+            processEventBatch: __bindMeta('processEventBatch')
         };
         `
     )
@@ -77,33 +100,3 @@ export function createPluginConfigVM(
         methods: vm.run('__methods'),
     }
 }
-
-export function prepareForRun(
-    server: PluginsServer,
-    teamId: number,
-    pluginConfig: PluginConfig, // might have team_id=0
-    method: 'processEvent',
-    event?: PluginEvent
-): null | ((event: PluginEvent) => Promise<PluginEvent>) | (() => Promise<void>) {
-    if (!pluginConfig.vm?.methods[method]) {
-        return null
-    }
-
-    const { vm } = pluginConfig.vm
-
-    if (event?.properties?.token) {
-        // TODO: this should be nicer... and it's not optimised for batch processing
-        const posthog = createInternalPostHogInstance(
-            event.properties.token,
-            { apiHost: event.site_url, fetch },
-            {
-                performance: performance,
-            }
-        )
-        vm.freeze(posthog, 'posthog')
-    } else {
-        vm.freeze(null, 'posthog')
-    }
-
-    return pluginConfig.vm.methods[method]
-}

src/worker/worker.ts

@@ -1,5 +1,5 @@
+import { runPlugins, runPluginsOnBatch, setupPlugins } from '../plugins'
 import { cloneObject, setLogLevel } from '../utils'
-import { runPlugins, setupPlugins } from '../plugins'
 import { createServer } from '../server'
 import { PluginsServerConfig } from '../types'
 
@@ -29,5 +29,10 @@ export async function createWorker(config: PluginsServerConfig): Promise<TaskWor
             // must clone the object, as we may get from VM2 something like { ..., properties: Proxy {} }
             return cloneObject(processedEvent as Record<string, any>)
         }
+        if (task === 'processEventBatch') {
+            const processedEvents = await runPluginsOnBatch(server, args.batch)
+            // must clone the object, as we may get from VM2 something like { ..., properties: Proxy {} }
+            return cloneObject(processedEvents as any[])
+        }
     }
 }

tests/plugins.test.ts

@@ -52,7 +52,7 @@ test('setupPlugins and runPlugins', async () => {
         },
     })
     expect(pluginConfig.vm).toBeDefined()
-    expect(Object.keys(pluginConfig.vm!.methods)).toEqual(['processEvent'])
+    expect(Object.keys(pluginConfig.vm!.methods)).toEqual(['processEvent', 'processEventBatch'])
 
     expect(setError).toHaveBeenCalled()
     expect(setError.mock.calls[0][0]).toEqual(mockServer)