diff --git a/jest.setup.fetch-mock.js b/jest.setup.fetch-mock.js index 9f684662..c9b2f01b 100644 --- a/jest.setup.fetch-mock.js +++ b/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'), }) ) ) diff --git a/src/celery/broker.ts b/src/celery/broker.ts index da3559de..cf391376 100644 --- a/src/celery/broker.ts +++ b/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)) } /** diff --git a/src/config.ts b/src/config.ts new file mode 100644 index 00000000..ee66b571 --- /dev/null +++ b/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 { + 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 = process.env +): PluginsServerConfig { + const defaultConfig = getDefaultConfig() + + const newConfig: Record = { ...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 +} diff --git a/src/index.ts b/src/index.ts index dda7ee66..b9fc22a7 100644 --- a/src/index.ts +++ b/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 = 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) diff --git a/src/plugins.ts b/src/plugins.ts index 94585670..ca634167 100644 --- a/src/plugins.ts +++ b/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) { diff --git a/src/server.ts b/src/server.ts index 7faecd5a..b8bc172c 100644 --- a/src/server.ts +++ b/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 = { ...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 = {} ): Promise<[PluginsServer, () => Promise]> { @@ -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, diff --git a/src/types.ts b/src/types.ts index 227b6dcc..be331d86 100644 --- a/src/types.ts +++ b/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,6 +24,7 @@ export interface PluginsServerConfig { DISABLE_WEB: boolean WEB_PORT: number WEB_HOSTNAME: string + LOG_LEVEL: LogLevel __jestMock?: { getPluginRows: Plugin[] @@ -21,6 +32,7 @@ export interface PluginsServerConfig { getPluginAttachmentRows: PluginAttachmentDB[] } } +export type PluginsServerConfigKey = Exclude export interface PluginsServer extends PluginsServerConfig { // active connections to postgres and redis diff --git a/src/utils.ts b/src/utils.ts index e8825644..e35cf5e1 100644 --- a/src/utils.ts +++ b/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 + } +} diff --git a/src/vm.ts b/src/vm.ts index bbc0951c..5db05b2c 100644 --- a/src/vm.ts +++ b/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( diff --git a/src/worker/config.ts b/src/worker/config.ts index 35d82d37..aa1a9e52 100644 --- a/src/worker/config.ts +++ b/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 } diff --git a/src/worker/piscina.js b/src/worker/piscina.js index 2d19e3db..56614345 100644 --- a/src/worker/piscina.js +++ b/src/worker/piscina.js @@ -15,8 +15,6 @@ if (isMainThread) { }, } } else { - console.log('🧵 Starting Piscina Worker Thread') - if (areWeTestingWithJest()) { require('ts-node').register() } diff --git a/src/worker/worker.ts b/src/worker/worker.ts index 155bc222..08bf2113 100644 --- a/src/worker/worker.ts +++ b/src/worker/worker.ts @@ -1,3 +1,4 @@ +import { setLogLevel } from '../utils' import { runPlugins, setupPlugins } from '../plugins' import { createServer } from '../server' import { PluginsServerConfig } from '../types' @@ -5,6 +6,10 @@ import { PluginsServerConfig } from '../types' type TaskWorker = ({ task, args }: { task: string; args: any }) => Promise export async function createWorker(config: PluginsServerConfig): Promise { + setLogLevel(config.LOG_LEVEL) + + console.info('🧵 Starting Piscina Worker Thread') + const [server, closeServer] = await createServer(config) await setupPlugins(server) diff --git a/tests/config.test.ts b/tests/config.test.ts new file mode 100644 index 00000000..84d62432 --- /dev/null +++ b/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) +}) diff --git a/tests/piscina.test.ts b/tests/piscina.test.ts deleted file mode 100644 index 074d10c5..00000000 --- a/tests/piscina.test.ts +++ /dev/null @@ -1,95 +0,0 @@ -import { defaultConfig } from '../src/server' -import { makePiscina } from '../src/worker/piscina' -import { PluginEvent } from 'posthog-plugins/src/types' -import { performance } from 'perf_hooks' -import { mockJestWithIndex } from './helpers/plugins' -import * as os from 'os' - -jest.mock('../src/sql') -jest.setTimeout(300000) // 300 sec timeout - -function processOneEvent(processEvent: (event: PluginEvent) => Promise): Promise { - const defaultEvent = { - distinct_id: 'my_id', - ip: '127.0.0.1', - site_url: 'http://localhost', - team_id: 2, - now: new Date().toISOString(), - event: 'default event', - properties: { key: 'value' }, - } - - return processEvent(defaultEvent) -} - -async function processCountEvents(count: number, piscina: ReturnType) { - const startTime = performance.now() - const promises = Array(count) - const processEvent = (event: PluginEvent) => piscina.runTask({ task: 'processEvent', args: { event } }) - for (let i = 0; i < count; i++) { - promises[i] = processOneEvent(processEvent) - } - // this will get heavy for tests > 10k events, should chunk them somehow... - await Promise.all(promises) - - const ms = Math.round((performance.now() - startTime) * 1000) / 1000 - - const log = { - eventsPerSecond: 1000 / (ms / count), - events: count, - concurrency: piscina.threads.length, - totalMs: ms, - averageEventMs: ms / count, - } - - return log -} - -function setupPiscina(workers: number) { - return makePiscina({ - ...defaultConfig, - WORKER_CONCURRENCY: workers, - __jestMock: mockJestWithIndex(` - function processEvent (event, meta) { - let j = 0; for(let i = 0; i < 200000; i++) { j = i }; - event.properties = { "somewhere": "over the rainbow" }; - return event - } - `), - }) -} - -test('piscina 2-24 workers', async () => { - const cpuCount = os.cpus().length - - const workers = [1, 2, 4, 8, 12, 16, 24, 32, 48, 64].filter((cores) => - cpuCount === 2 ? cores <= cpuCount : cores < cpuCount - ) - const events = 10000 - const rounds = 5 - - const results: Record = {} - for (const cores of workers) { - const piscina = setupPiscina(cores) - - // warmup - await processCountEvents(cpuCount * 4, piscina) - - // start - let throughput = 0 - for (let i = 0; i < rounds; i++) { - const { eventsPerSecond } = await processCountEvents(events, piscina) - throughput += eventsPerSecond - } - results[cores] = Math.round(throughput / rounds) - await piscina.destroy() - } - - console.log({ cpuCount }) - console.log(JSON.stringify(results, null, 2)) - - // expect that adding more cores (up to cpuCount) increases throughput - for (let i = 1; i < workers.length; i++) { - expect(results[workers[i - 1]]).toBeLessThan(results[workers[i]]) - } -}) diff --git a/tests/plugins.test.ts b/tests/plugins.test.ts index 0d7dd617..0f58f3cc 100644 --- a/tests/plugins.test.ts +++ b/tests/plugins.test.ts @@ -1,6 +1,6 @@ import { runPlugins, setupPlugins } from '../src/plugins' import { createServer } from '../src/server' -import { PluginsServer } from '../src/types' +import { LogLevel, PluginsServer } from '../src/types' import { PluginEvent } from 'posthog-plugins/src/types' import { mockPluginTempFolder, @@ -14,7 +14,7 @@ jest.mock('../src/sql') let mockServer: PluginsServer beforeEach(async () => { - ;[mockServer] = await createServer() + ;[mockServer] = await createServer({ LOG_LEVEL: LogLevel.Log }) }) test('setupPlugins and runPlugins', async () => { diff --git a/tests/queue.test.ts b/tests/queue.test.ts index a91c49fe..9e3f75e6 100644 --- a/tests/queue.test.ts +++ b/tests/queue.test.ts @@ -1,5 +1,5 @@ import { startQueue } from '../src/worker/queue' -import { createServer, defaultConfig } from '../src/server' +import { createServer } from '../src/server' import { PluginsServer } from '../src/types' import Client from '../src/celery/client' import { runPlugins } from '../src/plugins' @@ -14,7 +14,7 @@ beforeEach(async () => { // silence logs console.info = jest.fn() - mockServer = (await createServer(defaultConfig))[0] + mockServer = (await createServer())[0] }) test('worker and task passing via redis', async () => { @@ -53,7 +53,7 @@ test('worker and task passing via redis', async () => { expect(item['headers']['task']).toBe('posthog.tasks.process_event.process_event_with_plugins') expect(item['properties']['body_encoding']).toBe('base64') - const body = new Buffer(item['body'], 'base64').toString() + const body = Buffer.from(item['body'], 'base64').toString() const [args2, kwargs2] = JSON.parse(body) expect(args2).toEqual(args) @@ -74,7 +74,7 @@ test('worker and task passing via redis', async () => { expect(processedItem['headers']['task']).toBe('posthog.tasks.process_event.process_event') expect(processedItem['properties']['body_encoding']).toBe('base64') - const processedBody = new Buffer(processedItem['body'], 'base64').toString() + const processedBody = Buffer.from(processedItem['body'], 'base64').toString() const [args3, kwargs3] = JSON.parse(processedBody) expect(args3).toEqual([]) diff --git a/tests/utils.test.ts b/tests/utils.test.ts index 871c8098..a89b2c8a 100644 --- a/tests/utils.test.ts +++ b/tests/utils.test.ts @@ -1,4 +1,5 @@ -import { getFileFromTGZ, getFileFromZip, getFileFromArchive, bufferToStream } from '../src/utils' +import { getFileFromTGZ, getFileFromZip, getFileFromArchive, bufferToStream, setLogLevel } from '../src/utils' +import { LogLevel } from '../src/types' // .zip in Base64: github repo posthog/helloworldplugin const zip = @@ -55,3 +56,91 @@ test('bufferToStream', async () => { const stream = bufferToStream(buffer) expect(stream.read()).toEqual(buffer) }) + +test('setLogLevel', async () => { + function resetMocks() { + console.debug = jest.fn() + console.info = jest.fn() + console.log = jest.fn() + console.warn = jest.fn() + console.error = jest.fn() + } + + resetMocks() + setLogLevel(LogLevel.Debug) + console.debug('debug') + console.info('debug') + console.log('debug') + console.warn('debug') + console.error('debug') + expect(console.debug).toHaveBeenCalledWith('debug') + expect(console.info).toHaveBeenCalledWith('debug') + expect(console.log).toHaveBeenCalledWith('debug') + expect(console.warn).toHaveBeenCalledWith('debug') + expect(console.error).toHaveBeenCalledWith('debug') + + resetMocks() + setLogLevel(LogLevel.Info) + console.debug('info') + console.info('info') + console.log('info') + console.warn('info') + console.error('info') + expect((console.debug as any)._original).toBeDefined() + expect(console.info).toHaveBeenCalledWith('info') + expect(console.log).toHaveBeenCalledWith('info') + expect(console.warn).toHaveBeenCalledWith('info') + expect(console.error).toHaveBeenCalledWith('info') + + resetMocks() + setLogLevel(LogLevel.Log) + console.debug('log') + console.info('log') + console.log('log') + console.warn('log') + console.error('log') + expect((console.debug as any)._original).toBeDefined() + expect((console.info as any)._original).toBeDefined() + expect(console.log).toHaveBeenCalledWith('log') + expect(console.warn).toHaveBeenCalledWith('log') + expect(console.error).toHaveBeenCalledWith('log') + + resetMocks() + setLogLevel(LogLevel.Warn) + console.debug('warn') + console.info('warn') + console.log('warn') + console.warn('warn') + console.error('warn') + expect((console.debug as any)._original).toBeDefined() + expect((console.info as any)._original).toBeDefined() + expect((console.log as any)._original).toBeDefined() + expect(console.warn).toHaveBeenCalledWith('warn') + expect(console.error).toHaveBeenCalledWith('warn') + + resetMocks() + setLogLevel(LogLevel.Error) + console.debug('error') + console.info('error') + console.log('error') + console.warn('error') + console.error('error') + expect((console.debug as any)._original).toBeDefined() + expect((console.info as any)._original).toBeDefined() + expect((console.log as any)._original).toBeDefined() + expect((console.warn as any)._original).toBeDefined() + expect(console.error).toHaveBeenCalledWith('error') + + resetMocks() + setLogLevel(LogLevel.None) + console.debug('none') + console.info('none') + console.log('none') + console.warn('none') + console.error('none') + expect((console.debug as any)._original).toBeDefined() + expect((console.info as any)._original).toBeDefined() + expect((console.log as any)._original).toBeDefined() + expect((console.warn as any)._original).toBeDefined() + expect((console.error as any)._original).toBeDefined() +}) diff --git a/tests/vm.test.ts b/tests/vm.test.ts index 168b84b3..ea009912 100644 --- a/tests/vm.test.ts +++ b/tests/vm.test.ts @@ -1,7 +1,7 @@ import { createPluginConfigVM, prepareForRun } from '../src/vm' import { PluginConfig, PluginsServer, Plugin } from '../src/types' import { PluginEvent } from 'posthog-plugins' -import { createServer, defaultConfig } from '../src/server' +import { createServer } from '../src/server' import * as fetch from 'node-fetch' const defaultEvent = { @@ -40,7 +40,7 @@ const mockConfig: PluginConfig = { } beforeEach(async () => { - mockServer = (await createServer(defaultConfig))[0] + mockServer = (await createServer())[0] }) afterEach(async () => { diff --git a/tests/worker.test.ts b/tests/worker.test.ts new file mode 100644 index 00000000..ee9c0612 --- /dev/null +++ b/tests/worker.test.ts @@ -0,0 +1,130 @@ +import { makePiscina } from '../src/worker/piscina' +import { defaultConfig } from '../src/config' +import { PluginEvent } from 'posthog-plugins/src/types' +import { performance } from 'perf_hooks' +import { mockJestWithIndex } from './helpers/plugins' +import * as os from 'os' +import { LogLevel } from '../src/types' + +jest.mock('../src/sql') +jest.setTimeout(300000) // 300 sec timeout + +function processOneEvent(processEvent: (event: PluginEvent) => Promise): Promise { + const defaultEvent = { + distinct_id: 'my_id', + ip: '127.0.0.1', + site_url: 'http://localhost', + team_id: 2, + now: new Date().toISOString(), + event: 'default event', + properties: { key: 'value' }, + } + + return processEvent(defaultEvent) +} + +async function processCountEvents(count: number, piscina: ReturnType) { + const maxPromises = 1000 + const startTime = performance.now() + const promises = Array(maxPromises) + const processEvent = (event: PluginEvent) => piscina.runTask({ task: 'processEvent', args: { event } }) + + const groups = Math.ceil(count / maxPromises) + for (let j = 0; j < groups; j++) { + const groupCount = j === groups - 1 ? count % maxPromises : maxPromises + for (let i = 0; i < groupCount; i++) { + promises[i] = processOneEvent(processEvent) + } + await Promise.all(promises) + } + + const ms = Math.round((performance.now() - startTime) * 1000) / 1000 + + const log = { + eventsPerSecond: 1000 / (ms / count), + events: count, + concurrency: piscina.threads.length, + totalMs: ms, + averageEventMs: ms / count, + } + + return log +} + +function setupPiscina(workers: number, code: string, tasksPerWorker: number) { + return makePiscina({ + ...defaultConfig, + WORKER_CONCURRENCY: workers, + TASKS_PER_WORKER: tasksPerWorker, + LOG_LEVEL: LogLevel.Log, + __jestMock: mockJestWithIndex(code), + }) +} + +test('piscina worker test', async () => { + const coreCount = os.cpus().length + + const workers = [1, 2, 4, 8, 12, 16].filter((cores) => cores <= coreCount) + const rounds = 5 + + const tests: { testName: string; events: number; testCode: string }[] = [ + { + testName: 'simple', + events: 10000, + testCode: ` + function processEvent (event, meta) { + event.properties = { "somewhere": "over the rainbow" }; + return event + } + `, + }, + { + testName: 'for200k', + events: 10000, + testCode: ` + function processEvent (event, meta) { + let j = 0; for(let i = 0; i < 200000; i++) { j = i }; + event.properties = { "somewhere": "over the rainbow" }; + return event + } + `, + }, + { + testName: 'timeout100ms', + events: 2000, + testCode: ` + async function processEvent (event, meta) { + await new Promise(resolve => __jestSetTimeout(() => resolve(), 100)) + event.properties = { "somewhere": "over the rainbow" }; + return event + } + `, + }, + ] + + const results: Array> = [] + for (const { testName, events, testCode } of tests) { + const result: Record = { + testName, + coreCount, + } + for (const cores of workers) { + const piscina = setupPiscina(cores, testCode, 100) + + // warmup + await processCountEvents(cores * 4, piscina) + + // start + let throughput = 0 + for (let i = 0; i < rounds; i++) { + const { eventsPerSecond } = await processCountEvents(events, piscina) + throughput += eventsPerSecond + } + result[`${cores} cores`] = Math.round(throughput / rounds) + await piscina.destroy() + } + results.push(result) + console.log(JSON.stringify({ result }, null, 2)) + } + console.table(results) +})