feat(plugin-server): Add capability to use deferred overrides writer and worker

plugin-server/src/capabilities.ts

@@ -19,6 +19,7 @@ export function getPluginServerCapabilities(config: PluginsServerConfig): Plugin
                 processAsyncOnEventHandlers: true,
                 processAsyncWebhooksHandlers: true,
                 sessionRecordingBlobIngestion: true,
+                personOverrides: config.POE_DEFERRED_WRITES_ENABLED,
                 transpileFrontendApps: true,
                 preflightSchedules: true,
                 ...sharedCapabilities,
@@ -75,5 +76,10 @@ export function getPluginServerCapabilities(config: PluginsServerConfig): Plugin
                 transpileFrontendApps: true, // TODO: move this away from pod startup, into a graphile job
                 ...sharedCapabilities,
             }
+        case PluginServerMode.person_overrides:
+            return {
+                personOverrides: true,
+                ...sharedCapabilities,
+            }
     }
 }

plugin-server/src/config/config.ts

@@ -128,6 +128,7 @@ export function getDefaultConfig(): PluginsServerConfig {
         EXTERNAL_REQUEST_TIMEOUT_MS: 10 * 1000, // 10 seconds
         DROP_EVENTS_BY_TOKEN_DISTINCT_ID: '',
         DROP_EVENTS_BY_TOKEN: '',
+        POE_DEFERRED_WRITES_ENABLED: false,
         POE_EMBRACE_JOIN_FOR_TEAMS: '',
         RELOAD_PLUGIN_JITTER_MAX_MS: 60000,
 

plugin-server/src/main/pluginsServer.ts

@@ -15,11 +15,13 @@ import { Hub, PluginServerCapabilities, PluginsServerConfig } from '../types'
 import { createHub, createKafkaClient, createKafkaProducerWrapper } from '../utils/db/hub'
 import { PostgresRouter } from '../utils/db/postgres'
 import { cancelAllScheduledJobs } from '../utils/node-schedule'
+import { PeriodicTask } from '../utils/periodic-task'
 import { PubSub } from '../utils/pubsub'
 import { status } from '../utils/status'
 import { delay } from '../utils/utils'
 import { AppMetrics } from '../worker/ingestion/app-metrics'
 import { OrganizationManager } from '../worker/ingestion/organization-manager'
+import { DeferredPersonOverrideWorker } from '../worker/ingestion/person-state'
 import { TeamManager } from '../worker/ingestion/team-manager'
 import Piscina, { makePiscina as defaultMakePiscina } from '../worker/piscina'
 import { GraphileWorker } from './graphile-worker/graphile-worker'
@@ -108,6 +110,8 @@ export async function startPluginsServer(
     let jobsConsumer: Consumer | undefined
     let schedulerTasksConsumer: Consumer | undefined
 
+    let personOverridesPeriodicTask: PeriodicTask | undefined
+
     let httpServer: Server | undefined // healthcheck server
 
     let graphileWorker: GraphileWorker | undefined
@@ -146,6 +150,7 @@ export async function startPluginsServer(
             jobsConsumer?.disconnect(),
             stopSessionRecordingBlobConsumer?.(),
             schedulerTasksConsumer?.disconnect(),
+            personOverridesPeriodicTask?.stop(),
         ])
 
         if (piscina) {
@@ -428,6 +433,18 @@ export async function startPluginsServer(
             }
         }
 
+        if (capabilities.personOverrides) {
+            const postgres = hub?.postgres ?? new PostgresRouter(serverConfig)
+            const kafkaProducer = hub?.kafkaProducer ?? (await createKafkaProducerWrapper(serverConfig))
+
+            personOverridesPeriodicTask = new DeferredPersonOverrideWorker(postgres, kafkaProducer).runTask(5000)
+            personOverridesPeriodicTask.promise.catch(async () => {
+                status.error('⚠️', 'Person override worker task crashed! Requesting shutdown...')
+                await closeJobs()
+                process.exit(1)
+            })
+        }
+
         if (capabilities.http) {
             httpServer = createHttpServer(serverConfig.HTTP_SERVER_PORT, healthChecks, analyticsEventsIngestionConsumer)
         }

plugin-server/src/types.ts

@@ -76,6 +76,7 @@ export enum PluginServerMode {
     scheduler = 'scheduler',
     analytics_ingestion = 'analytics-ingestion',
     recordings_blob_ingestion = 'recordings-blob-ingestion',
+    person_overrides = 'person-overrides',
 }
 
 export const stringToPluginServerMode = Object.fromEntries(
@@ -198,6 +199,7 @@ export interface PluginsServerConfig {
     DROP_EVENTS_BY_TOKEN_DISTINCT_ID: string
     DROP_EVENTS_BY_TOKEN: string
     POE_EMBRACE_JOIN_FOR_TEAMS: string
+    POE_DEFERRED_WRITES_ENABLED: boolean
     RELOAD_PLUGIN_JITTER_MAX_MS: number
     SKIP_UPDATE_EVENT_AND_PROPERTIES_STEP: boolean
 
@@ -286,6 +288,7 @@ export interface PluginServerCapabilities {
     processAsyncOnEventHandlers?: boolean
     processAsyncWebhooksHandlers?: boolean
     sessionRecordingBlobIngestion?: boolean
+    personOverrides?: boolean
     transpileFrontendApps?: boolean // TODO: move this away from pod startup, into a graphile job
     preflightSchedules?: boolean // Used for instance health checks on hobby deploy, not useful on cloud
     http?: boolean

plugin-server/src/utils/periodic-task.ts

@@ -0,0 +1,52 @@
+import { status } from './status'
+import { sleep } from './utils'
+
+export class PeriodicTask {
+    public readonly promise: Promise<void>
+    private running = true
+    private abortController: AbortController
+
+    constructor(task: () => Promise<void> | void, intervalMs: number, minimumWaitMs = 0) {
+        this.abortController = new AbortController()
+
+        const abortRequested = new Promise((resolve) => {
+            this.abortController.signal.addEventListener('abort', resolve, { once: true })
+        })
+
+        this.promise = new Promise<void>(async (resolve, reject) => {
+            try {
+                status.debug('🔄', 'Periodic task starting...', { task })
+                while (!this.abortController.signal.aborted) {
+                    const startTimeMs = Date.now()
+                    await task()
+                    const durationMs = Date.now() - startTimeMs
+                    const waitTimeMs = Math.max(intervalMs - durationMs, minimumWaitMs)
+                    status.debug(
+                        '🔄',
+                        `Task completed in ${durationMs / 1000}s, next evaluation in ${waitTimeMs / 1000}s`,
+                        { task }
+                    )
+                    await Promise.race([sleep(waitTimeMs), abortRequested])
+                }
+                status.info('✅', 'Periodic task stopped by request.', { task })
+                resolve()
+            } catch (error) {
+                status.warn('⚠️', 'Error in periodic task!', { task, error })
+                reject(error)
+            } finally {
+                this.running = false
+            }
+        })
+    }
+
+    public isRunning(): boolean {
+        return this.running
+    }
+
+    public async stop(): Promise<void> {
+        this.abortController.abort()
+        try {
+            await this.promise
+        } catch {}
+    }
+}

plugin-server/src/worker/ingestion/event-pipeline/processPersonsStep.ts

@@ -4,7 +4,7 @@ import { Person } from 'types'
 
 import { normalizeEvent } from '../../../utils/event'
 import { status } from '../../../utils/status'
-import { PersonOverrideWriter, PersonState } from '../person-state'
+import { DeferredPersonOverrideWriter, PersonOverrideWriter, PersonState } from '../person-state'
 import { parseEventTimestamp } from '../timestamps'
 import { EventPipelineRunner } from './runner'
 
@@ -22,13 +22,22 @@ export async function processPersonsStep(
         throw error
     }
 
+    let overridesWriter: PersonOverrideWriter | DeferredPersonOverrideWriter | undefined = undefined
+    if (runner.poEEmbraceJoin) {
+        if (runner.hub.POE_DEFERRED_WRITES_ENABLED) {
+            overridesWriter = new DeferredPersonOverrideWriter(runner.hub.db.postgres)
+        } else {
+            overridesWriter = new PersonOverrideWriter(runner.hub.db.postgres)
+        }
+    }
+
     const person = await new PersonState(
         event,
         event.team_id,
         String(event.distinct_id),
         timestamp,
         runner.hub.db,
-        runner.poEEmbraceJoin ? new PersonOverrideWriter(runner.hub.db.postgres) : undefined
+        overridesWriter
     ).update()
 
     return [event, person]

plugin-server/src/worker/ingestion/person-state.ts

@@ -10,6 +10,7 @@ import { Person, PropertyUpdateOperation, TimestampFormat } from '../../types'
 import { DB } from '../../utils/db/db'
 import { PostgresRouter, PostgresUse, TransactionClient } from '../../utils/db/postgres'
 import { timeoutGuard } from '../../utils/db/utils'
+import { PeriodicTask } from '../../utils/periodic-task'
 import { promiseRetry } from '../../utils/retries'
 import { status } from '../../utils/status'
 import { castTimestampOrNow, UUIDT } from '../../utils/utils'
@@ -708,11 +709,7 @@ export class PersonOverrideWriter {
 }
 
 export class DeferredPersonOverrideWriter {
-    /**
-     * @param lockId the lock identifier/key used to ensure that only one
-     *               process is updating the overrides at a time
-     */
-    constructor(private postgres: PostgresRouter, private lockId: number) {}
+    constructor(private postgres: PostgresRouter) {}
 
     /**
      * Enqueue an override for deferred processing.
@@ -741,6 +738,18 @@ export class DeferredPersonOverrideWriter {
 
         return []
     }
+}
+
+export class DeferredPersonOverrideWorker {
+    // The advisory lock identifier/key used to ensure that only one process is
+    // updating the overrides at a time.
+    public readonly lockId = 567
+
+    private writer: PersonOverrideWriter
+
+    constructor(private postgres: PostgresRouter, private kafkaProducer: KafkaProducerWrapper) {
+        this.writer = new PersonOverrideWriter(this.postgres)
+    }
 
     /**
      * Process all (or up to the given limit) pending overrides.
@@ -751,9 +760,7 @@ export class DeferredPersonOverrideWriter {
      *
      * @returns the number of overrides processed
      */
-    public async processPendingOverrides(kafkaProducer: KafkaProducerWrapper, limit?: number): Promise<number> {
-        const writer = new PersonOverrideWriter(this.postgres)
-
+    public async processPendingOverrides(limit?: number): Promise<number> {
         return await this.postgres.transaction(PostgresUse.COMMON_WRITE, 'processPendingOverrides', async (tx) => {
             const {
                 rows: [{ acquired }],
@@ -778,7 +785,7 @@ export class DeferredPersonOverrideWriter {
 
             const messages: ProducerRecord[] = []
             for (const { id, ...mergeOperation } of rows) {
-                messages.push(...(await writer.addPersonOverride(tx, mergeOperation)))
+                messages.push(...(await this.writer.addPersonOverride(tx, mergeOperation)))
                 await this.postgres.query(
                     tx,
                     SQL`DELETE FROM posthog_pendingpersonoverride WHERE id = ${id}`,
@@ -797,11 +804,19 @@ export class DeferredPersonOverrideWriter {
             // Postgres for some reason -- the same row state should be
             // generated each call, and the receiving ReplacingMergeTree will
             // ensure we keep only the latest version after all writes settle.)
-            await kafkaProducer.queueMessages(messages, true)
+            await this.kafkaProducer.queueMessages(messages, true)
 
             return rows.length
         })
     }
+
+    public runTask(intervalMs: number): PeriodicTask {
+        return new PeriodicTask(async () => {
+            status.debug('👥', 'Processing pending overrides...')
+            const overridesCount = await this.processPendingOverrides()
+            ;(overridesCount > 0 ? status.info : status.debug)('👥', `Processed ${overridesCount} pending overrides.`)
+        }, intervalMs)
+    }
 }
 
 function SQL(sqlParts: TemplateStringsArray, ...args: any[]): { text: string; values: any[] } {

plugin-server/tests/utils/periodic-task.test.ts

@@ -0,0 +1,24 @@
+import { PeriodicTask } from '../../src/utils/periodic-task'
+
+describe('PeriodicTask', () => {
+    describe('updates completion status', () => {
+        it('on success', async () => {
+            const fn = jest.fn()
+            const task = new PeriodicTask(fn, 1000)
+            expect(fn).toBeCalled()
+            expect(task.isRunning()).toEqual(true)
+            await task.stop()
+            expect(task.isRunning()).toEqual(false)
+        })
+
+        it('on failure', async () => {
+            const fn = jest.fn(() => {
+                throw new Error()
+            })
+            const task = new PeriodicTask(fn, 1000)
+            expect(fn).toBeCalled()
+            await task.stop()
+            expect(task.isRunning()).toEqual(false)
+        })
+    })
+})