Benchmarks & various changes

jest.setup.fetch-mock.js

@@ -6,7 +6,8 @@ jest.mock('node-fetch', () => {
         (url) =>
             new Promise((resolve) =>
                 resolve({
-                    json: () => new Promise((resolve) => resolve(responsesToUrls[url])),
+                    json: () => new Promise((resolve) => resolve(responsesToUrls[url]) || { fetch: 'mock' }),
+                    text: () => new Promise((resolve) => resolve(JSON.stringify(responsesToUrls[url])) || 'fetchmock'),
                 })
             )
     )

src/celery/broker.ts

@@ -131,7 +131,7 @@ export default class RedisBroker {
                 Promise.resolve()
             })
             .then(() => this.receive(index, resolve, queue, callback))
-            .catch((err) => console.log(err))
+            .catch((err) => console.error(err))
     }
 
     /**

src/config.ts

@@ -0,0 +1,59 @@
+import { LogLevel, PluginsServerConfig, PluginsServerConfigKey } from './types'
+
+export const defaultConfig = overrideWithEnv(getDefaultConfig())
+export const configHelp = getConfigHelp()
+
+export function getDefaultConfig(): PluginsServerConfig {
+    return {
+        CELERY_DEFAULT_QUEUE: 'celery',
+        DATABASE_URL: 'postgres://localhost:5432/posthog',
+        PLUGINS_CELERY_QUEUE: 'posthog-plugins',
+        REDIS_URL: 'redis://localhost/',
+        BASE_DIR: '.',
+        PLUGINS_RELOAD_PUBSUB_CHANNEL: 'reload-plugins',
+        DISABLE_WEB: false,
+        WEB_PORT: 3008,
+        WEB_HOSTNAME: '0.0.0.0',
+        WORKER_CONCURRENCY: 0, // use all cores
+        TASKS_PER_WORKER: 100,
+        LOG_LEVEL: LogLevel.Info,
+    }
+}
+
+export function getConfigHelp(): Record<PluginsServerConfigKey, string> {
+    return {
+        CELERY_DEFAULT_QUEUE: 'celery outgoing queue',
+        DATABASE_URL: 'url for postgres',
+        PLUGINS_CELERY_QUEUE: 'celery incoming queue',
+        REDIS_URL: 'url for redis',
+        BASE_DIR: 'base path for resolving local plugins',
+        PLUGINS_RELOAD_PUBSUB_CHANNEL: 'redis channel for reload events',
+        DISABLE_WEB: 'do not start the web service',
+        WEB_PORT: 'port for web server',
+        WEB_HOSTNAME: 'hostname for web server',
+        WORKER_CONCURRENCY: 'number of concurrent worker threads',
+        TASKS_PER_WORKER: 'number of parallel tasks per worker thread',
+        LOG_LEVEL: 'minimum log level',
+    }
+}
+
+export function overrideWithEnv(
+    config: PluginsServerConfig,
+    env: Record<string, string | undefined> = process.env
+): PluginsServerConfig {
+    const defaultConfig = getDefaultConfig()
+
+    const newConfig: Record<PluginsServerConfigKey, any> = { ...config }
+    for (const key of Object.keys(config) as PluginsServerConfigKey[]) {
+        if (typeof env[key] !== 'undefined') {
+            if (typeof defaultConfig[key] === 'number') {
+                newConfig[key] = env[key]?.indexOf('.') ? parseFloat(env[key]!) : parseInt(env[key]!)
+            } else if (typeof defaultConfig[key] === 'boolean') {
+                newConfig[key] = env[key] === 'true' || env[key] === 'True' || env[key] === '1'
+            } else {
+                newConfig[key] = env[key]
+            }
+        }
+    }
+    return newConfig
+}

src/index.ts

@@ -1,7 +1,9 @@
 import * as yargs from 'yargs'
-import { PluginsServerConfig } from './types'
+import { PluginsServerConfig, PluginsServerConfigKey } from './types'
 import { startPluginsServer } from './server'
 import { makePiscina } from './worker/piscina'
+import { defaultConfig, configHelp } from './config'
+import { setLogLevel } from './utils'
 
 type Argv = {
     config: string
@@ -11,31 +13,32 @@ type Argv = {
     concurrency: number
 }
 
-yargs
+let app: any = yargs
+    .wrap(yargs.terminalWidth())
     .scriptName('posthog-plugins')
     .option('config', { alias: 'c', describe: 'Config options JSON.', type: 'string' })
-    .option('disable-web', { describe: 'Whether web server should be disabled.', type: 'boolean' })
-    .option('web-port', { alias: 'p', describe: 'Web server port.', type: 'number' })
-    .option('web-hostname', { alias: 'h', describe: 'Web server hostname.', type: 'string' })
-    .option('concurrency', { describe: 'Concurrenct Worker Threads', type: 'number' })
-    .help()
-    .command({
-        command: ['start', '$0'],
-        describe: 'start the server',
-        handler: ({ config, disableWeb, webPort, webHostname, concurrency }: Argv) => {
-            const parsedConfig: PluginsServerConfig = config ? JSON.parse(config) : {}
-            if (typeof webHostname !== 'undefined') {
-                parsedConfig['WEB_HOSTNAME'] = webHostname
-            }
-            if (typeof webPort !== 'undefined') {
-                parsedConfig['WEB_PORT'] = webPort
-            }
-            if (typeof disableWeb !== 'undefined') {
-                parsedConfig['DISABLE_WEB'] = disableWeb
-            }
-            if (typeof concurrency !== 'undefined') {
-                parsedConfig['WORKER_CONCURRENCY'] = concurrency
-            }
-            startPluginsServer(parsedConfig, makePiscina)
-        },
-    }).argv
+
+for (const [key, value] of Object.entries(defaultConfig)) {
+    app = app.option(key.toLowerCase().split('_').join('-'), {
+        describe: `${configHelp[key as PluginsServerConfigKey] || key} [${value}]`,
+        type: typeof value,
+    })
+}
+
+const { config, ...otherArgs }: Argv = app.help().argv
+
+const parsedConfig: Record<string, any> = config ? JSON.parse(config) : {}
+for (const [key, value] of Object.entries(otherArgs)) {
+    if (typeof value !== 'undefined') {
+        // convert camelCase argument keys to under_score
+        const newKey = key
+            .replace(/(?:^|\.?)([A-Z])/g, (x, y) => '_' + y.toUpperCase())
+            .replace(/^_/, '')
+            .toUpperCase()
+        if (newKey in defaultConfig) {
+            parsedConfig[newKey] = value
+        }
+    }
+}
+setLogLevel(parsedConfig.LOG_LEVEL || defaultConfig.LOG_LEVEL)
+startPluginsServer(parsedConfig as PluginsServerConfig, makePiscina)

src/plugins.ts

@@ -161,7 +161,7 @@ async function loadPlugin(server: PluginsServer, pluginConfig: PluginConfig): Pr
             if (indexJs) {
                 try {
                     pluginConfig.vm = createPluginConfigVM(server, pluginConfig, indexJs, libJs || '')
-                    console.log(`Loaded plugin "${plugin.name}"!`)
+                    console.info(`Loaded plugin "${plugin.name}"!`)
                     await clearError(server, pluginConfig)
                     return true
                 } catch (error) {

src/server.ts

@@ -8,31 +8,9 @@ import { startFastifyInstance, stopFastifyInstance } from './web/server'
 import { Worker } from 'celery/worker'
 import { version } from '../package.json'
 import { PluginEvent } from 'posthog-plugins'
+import { defaultConfig } from './config'
 import Piscina from 'piscina'
 
-function overrideWithEnv(config: PluginsServerConfig): PluginsServerConfig {
-    const newConfig: Record<string, any> = { ...config }
-    for (const [key, value] of Object.entries(config)) {
-        if (process.env[key]) {
-            newConfig[key] = process.env[key]
-        }
-    }
-    return newConfig as PluginsServerConfig
-}
-
-export const defaultConfig: PluginsServerConfig = overrideWithEnv({
-    CELERY_DEFAULT_QUEUE: 'celery',
-    DATABASE_URL: 'postgres://localhost:5432/posthog',
-    PLUGINS_CELERY_QUEUE: 'posthog-plugins',
-    REDIS_URL: 'redis://localhost/',
-    BASE_DIR: '.',
-    PLUGINS_RELOAD_PUBSUB_CHANNEL: 'reload-plugins',
-    DISABLE_WEB: false,
-    WEB_PORT: 3008,
-    WEB_HOSTNAME: '0.0.0.0',
-    WORKER_CONCURRENCY: 0, // use all cores
-})
-
 export async function createServer(
     config: Partial<PluginsServerConfig> = {}
 ): Promise<[PluginsServer, () => Promise<void>]> {
@@ -47,6 +25,11 @@ export async function createServer(
 
     const redis = new Redis(serverConfig.REDIS_URL)
 
+    redis.on('error', (error) => {
+        console.error('🔴 Redis error!', error)
+        process.kill(process.pid, 'SIGTERM')
+    })
+
     const server: PluginsServer = {
         ...serverConfig,
         db,

src/types.ts

@@ -3,8 +3,18 @@ import { Redis } from 'ioredis'
 import { PluginEvent, PluginAttachment, PluginConfigSchema } from 'posthog-plugins'
 import { VM, VMScript } from 'vm2'
 
+export enum LogLevel {
+    Debug = 'debug',
+    Info = 'info',
+    Log = 'log',
+    Warn = 'warn',
+    Error = 'error',
+    None = 'none',
+}
+
 export interface PluginsServerConfig {
     WORKER_CONCURRENCY: number
+    TASKS_PER_WORKER: number
     CELERY_DEFAULT_QUEUE: string
     DATABASE_URL: string
     PLUGINS_CELERY_QUEUE: string
@@ -14,13 +24,15 @@ export interface PluginsServerConfig {
     DISABLE_WEB: boolean
     WEB_PORT: number
     WEB_HOSTNAME: string
+    LOG_LEVEL: LogLevel
 
     __jestMock?: {
         getPluginRows: Plugin[]
         getPluginConfigRows: PluginConfig[]
         getPluginAttachmentRows: PluginAttachmentDB[]
     }
 }
+export type PluginsServerConfigKey = Exclude<keyof PluginsServerConfig, '__jestMock'>
 
 export interface PluginsServer extends PluginsServerConfig {
     // active connections to postgres and redis

src/utils.ts

@@ -2,6 +2,7 @@ import { Readable } from 'stream'
 import * as tar from 'tar-stream'
 import AdmZip from 'adm-zip'
 import * as zlib from 'zlib'
+import { LogLevel } from './types'
 
 /**
  * @param binary Buffer
@@ -94,3 +95,15 @@ export function getFileFromZip(archive: Buffer, file: string): string | null {
 
     return null
 }
+
+export function setLogLevel(logLevel: LogLevel): void {
+    for (const loopLevel of ['debug', 'info', 'log', 'warn', 'error', 'none']) {
+        if (loopLevel === logLevel) {
+            break
+        }
+        const originalFunction = (console as any)[loopLevel]._original || (console as any)[loopLevel]
+        // eslint-disable-next-line @typescript-eslint/no-empty-function
+        ;(console as any)[loopLevel] = () => {}
+        ;(console as any)[loopLevel]._original = originalFunction
+    }
+}

src/vm.ts

@@ -7,6 +7,10 @@ import { createCache } from './extensions/cache'
 import { createInternalPostHogInstance } from 'posthog-js-lite'
 import { performance } from 'perf_hooks'
 
+function areWeTestingWithJest() {
+    return process.env.JEST_WORKER_ID !== undefined
+}
+
 export function createPluginConfigVM(
     server: PluginsServer,
     pluginConfig: PluginConfig, // NB! might have team_id = 0
@@ -18,6 +22,9 @@ export function createPluginConfigVM(
     })
     vm.freeze(createConsole(), 'console')
     vm.freeze(fetch, 'fetch')
+    if (areWeTestingWithJest()) {
+        vm.freeze(setTimeout, '__jestSetTimeout')
+    }
     vm.freeze(
         {
             cache: createCache(

src/worker/config.ts

@@ -31,5 +31,9 @@ export function createConfig(serverConfig: PluginsServerConfig, filename: string
         config.maxThreads = serverConfig.WORKER_CONCURRENCY
     }
 
+    if (serverConfig.TASKS_PER_WORKER > 1) {
+        config.concurrentTasksPerWorker = serverConfig.TASKS_PER_WORKER
+    }
+
     return config
 }

src/worker/piscina.js

@@ -15,8 +15,6 @@ if (isMainThread) {
         },
     }
 } else {
-    console.log('🧵 Starting Piscina Worker Thread')
-
     if (areWeTestingWithJest()) {
         require('ts-node').register()
     }

src/worker/worker.ts

@@ -1,10 +1,15 @@
+import { setLogLevel } from '../utils'
 import { runPlugins, setupPlugins } from '../plugins'
 import { createServer } from '../server'
 import { PluginsServerConfig } from '../types'
 
 type TaskWorker = ({ task, args }: { task: string; args: any }) => Promise<any>
 
 export async function createWorker(config: PluginsServerConfig): Promise<TaskWorker> {
+    setLogLevel(config.LOG_LEVEL)
+
+    console.info('🧵 Starting Piscina Worker Thread')
+
     const [server, closeServer] = await createServer(config)
     await setupPlugins(server)
 

tests/config.test.ts

@@ -0,0 +1,29 @@
+import { getDefaultConfig, overrideWithEnv } from '../src/config'
+
+test('overrideWithEnv 1', async () => {
+    const defaultConfig = getDefaultConfig()
+    const env = {
+        DISABLE_WEB: 'false',
+        WEB_PORT: '3008',
+        WEB_HOSTNAME: '0.0.0.0',
+        BASE_DIR: undefined,
+    }
+    const config = overrideWithEnv(getDefaultConfig(), env)
+
+    expect(config.DISABLE_WEB).toEqual(false)
+    expect(config.WEB_PORT).toEqual(3008)
+    expect(config.WEB_HOSTNAME).toEqual('0.0.0.0')
+    expect(config.BASE_DIR).toEqual(defaultConfig.BASE_DIR)
+})
+
+test('overrideWithEnv 2', async () => {
+    const defaultConfig = getDefaultConfig()
+    const env = {
+        DISABLE_WEB: '1',
+        WEB_PORT: '3008.12',
+    }
+    const config = overrideWithEnv(getDefaultConfig(), env)
+
+    expect(config.DISABLE_WEB).toEqual(true)
+    expect(config.WEB_PORT).toEqual(3008.12)
+})