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 { DeferredPersonOverrideWriter } 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 personOverridesWorker: 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(),
+            personOverridesWorker?.stop(),
         ])
 
         if (piscina) {
@@ -428,6 +433,23 @@ export async function startPluginsServer(
             }
         }
 
+        if (capabilities.personOverrides) {
+            const postgres = hub?.postgres ?? new PostgresRouter(serverConfig)
+            const kafkaProducer = hub?.kafkaProducer ?? (await createKafkaProducerWrapper(serverConfig))
+            const overridesWriter = new DeferredPersonOverrideWriter(postgres)
+
+            personOverridesWorker = new PeriodicTask(async () => {
+                status.debug('👥', 'Processing pending overrides...')
+                const overridesCount = await overridesWriter.processPendingOverrides(kafkaProducer)
+                ;(overridesCount > 0 ? status.info : status.debug)(
+                    '👥',
+                    `Processed ${overridesCount} pending overrides.`
+                )
+            })
+
+            healthChecks['person-overrides'] = () => personOverridesWorker!.isRunning()
+        }
+
         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,47 @@
+import { status } from './status'
+import { sleep } from './utils'
+
+export class PeriodicTask {
+    private promise: Promise<void>
+    private running = true
+    private abortController: AbortController
+
+    constructor(task: () => Promise<void> | void, intervalMs = 1000, minimumWaitMs = 1000) {
+        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 waitTimeMs = Math.max(intervalMs - startTimeMs, minimumWaitMs)
+                    status.debug('🔄', `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

@@ -2,9 +2,10 @@ import { PluginEvent } from '@posthog/plugin-scaffold'
 import { DateTime } from 'luxon'
 import { Person } from 'types'
 
+import { defaultConfig } from '../../../config/config'
 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 +23,22 @@ export async function processPersonsStep(
         throw error
     }
 
+    let overridesWriter: PersonOverrideWriter | DeferredPersonOverrideWriter | undefined = undefined
+    if (runner.poEEmbraceJoin) {
+        if (defaultConfig.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

@@ -712,7 +712,7 @@ 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, private lockId: number = 567) {}
 
     /**
      * Enqueue an override for deferred processing.

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)
+            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)
+            expect(fn).toBeCalled()
+            await task.stop()
+            expect(task.isRunning()).toEqual(false)
+        })
+    })
+})