refactor(plugin-server): Add deferred person override writer

latest_migrations.manifest

@@ -5,7 +5,7 @@ contenttypes: 0002_remove_content_type_name
 ee: 0015_add_verified_properties
 otp_static: 0002_throttling
 otp_totp: 0002_auto_20190420_0723
-posthog: 0367_job_inputs
+posthog: 0368_pendingpersonoverride
 sessions: 0001_initial
 social_django: 0010_uid_db_index
 two_factor: 0007_auto_20201201_1019

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 { PersonState } from '../person-state'
+import { PersonOverrideWriter, PersonState } from '../person-state'
 import { parseEventTimestamp } from '../timestamps'
 import { EventPipelineRunner } from './runner'
 
@@ -29,7 +29,7 @@ export async function processPersonsStep(
         timestamp,
         runner.hub.db,
         runner.hub.statsd,
-        runner.poEEmbraceJoin
+        runner.poEEmbraceJoin ? new PersonOverrideWriter(runner.hub.db.postgres) : undefined
     ).update()
 
     return [event, person]

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

@@ -4,6 +4,7 @@ import { StatsD } from 'hot-shots'
 import { ProducerRecord } from 'kafkajs'
 import { DateTime } from 'luxon'
 import { Counter } from 'prom-client'
+import { KafkaProducerWrapper } from 'utils/db/kafka-producer-wrapper'
 
 import { KAFKA_PERSON_OVERRIDE } from '../../config/kafka-topics'
 import { Person, PropertyUpdateOperation, TimestampFormat } from '../../types'
@@ -92,7 +93,6 @@ export class PersonState {
     private db: DB
     private statsd: StatsD | undefined
     public updateIsIdentified: boolean // TODO: remove this from the class and being hidden
-    private poEEmbraceJoin: boolean
 
     constructor(
         event: PluginEvent,
@@ -101,7 +101,7 @@ export class PersonState {
         timestamp: DateTime,
         db: DB,
         statsd: StatsD | undefined = undefined,
-        poEEmbraceJoin = false,
+        private personOverrideWriter?: PersonOverrideWriter | DeferredPersonOverrideWriter,
         uuid: UUIDT | undefined = undefined,
         maxMergeAttempts: number = MAX_FAILED_PERSON_MERGE_ATTEMPTS
     ) {
@@ -119,9 +119,6 @@ export class PersonState {
         // If set to true, we'll update `is_identified` at the end of `updateProperties`
         // :KLUDGE: This is an indirect communication channel between `handleIdentifyOrAlias` and `updateProperties`
         this.updateIsIdentified = false
-
-        // For persons on events embrace the join gradual roll-out, remove after fully rolled out
-        this.poEEmbraceJoin = poEEmbraceJoin
     }
 
     async update(): Promise<Person> {
@@ -451,7 +448,7 @@ export class PersonState {
         const properties: Properties = { ...otherPerson.properties, ...mergeInto.properties }
         this.applyEventPropertyUpdates(properties)
 
-        if (this.poEEmbraceJoin) {
+        if (this.personOverrideWriter) {
             // Optimize merging persons to keep using the person id that has longer history,
             // which means we'll have less events to update during the squash later
             if (otherPerson.created_at < mergeInto.created_at) {
@@ -486,7 +483,7 @@ export class PersonState {
                 call: this.event.event, // $identify, $create_alias or $merge_dangerously
                 oldPersonIdentified: String(otherPerson.is_identified),
                 newPersonIdentified: String(mergeInto.is_identified),
-                poEEmbraceJoin: String(this.poEEmbraceJoin),
+                poEEmbraceJoin: String(!!this.personOverrideWriter),
             })
             .inc()
 
@@ -518,12 +515,10 @@ export class PersonState {
                 const deletePersonMessages = await this.db.deletePerson(otherPerson, tx)
 
                 let personOverrideMessages: ProducerRecord[] = []
-                if (this.poEEmbraceJoin) {
-                    personOverrideMessages = await new PersonOverrideWriter(this.db.postgres).addPersonOverride(
+                if (this.personOverrideWriter) {
+                    personOverrideMessages = await this.personOverrideWriter.addPersonOverride(
                         tx,
-                        this.teamId,
-                        otherPerson,
-                        mergeInto
+                        getMergeOperation(this.teamId, otherPerson, mergeInto)
                     )
                 }
 
@@ -544,37 +539,54 @@ export class PersonState {
                 call: this.event.event, // $identify, $create_alias or $merge_dangerously
                 oldPersonIdentified: String(otherPerson.is_identified),
                 newPersonIdentified: String(mergeInto.is_identified),
-                poEEmbraceJoin: String(this.poEEmbraceJoin),
+                poEEmbraceJoin: String(!!this.personOverrideWriter),
             })
             .inc()
         return result
     }
 }
 
-class PersonOverrideWriter {
-    constructor(private postgres: PostgresRouter) {}
+type MergeOperation = {
+    team_id: number
+    old_person_id: string
+    override_person_id: string
+    oldest_event: DateTime
+}
 
-    public async addPersonOverride(
-        tx: TransactionClient,
-        teamId: number,
-        oldPerson: Person,
-        overridePerson: Person
-    ): Promise<ProducerRecord[]> {
-        if (teamId != oldPerson.team_id || teamId != overridePerson.team_id) {
-            throw new Error('cannot merge persons across different teams')
-        }
+function getMergeOperation(teamId: number, oldPerson: Person, overridePerson: Person): MergeOperation {
+    if (teamId != oldPerson.team_id || teamId != overridePerson.team_id) {
+        throw new Error('cannot merge persons across different teams')
+    }
+    return {
+        team_id: teamId,
+        old_person_id: oldPerson.uuid,
+        override_person_id: overridePerson.uuid,
+        oldest_event: overridePerson.created_at,
+    }
+}
 
+export class PersonOverrideWriter {
+    constructor(private postgres: PostgresRouter) {}
+
+    public async addPersonOverride(tx: TransactionClient, mergeOperation: MergeOperation): Promise<ProducerRecord[]> {
         const mergedAt = DateTime.now()
-        const oldestEvent = overridePerson.created_at
         /**
             We'll need to do 4 updates:
 
          1. Add the persons involved to the helper table (2 of them)
          2. Add an override from oldPerson to override person
          3. Update any entries that have oldPerson as the override person to now also point to the new override person. Note that we don't update `oldest_event`, because it's a heuristic (used to optimise squashing) tied to the old_person and nothing changed about the old_person who's events need to get squashed.
          */
-        const oldPersonId = await this.addPersonOverrideMapping(tx, oldPerson)
-        const overridePersonId = await this.addPersonOverrideMapping(tx, overridePerson)
+        const oldPersonMappingId = await this.addPersonOverrideMapping(
+            tx,
+            mergeOperation.team_id,
+            mergeOperation.old_person_id
+        )
+        const overridePersonMappingId = await this.addPersonOverrideMapping(
+            tx,
+            mergeOperation.team_id,
+            mergeOperation.override_person_id
+        )
 
         await this.postgres.query(
             tx,
@@ -586,10 +598,10 @@ class PersonOverrideWriter {
                     oldest_event,
                     version
                 ) VALUES (
-                    ${teamId},
-                    ${oldPersonId},
-                    ${overridePersonId},
-                    ${oldestEvent},
+                    ${mergeOperation.team_id},
+                    ${oldPersonMappingId},
+                    ${overridePersonMappingId},
+                    ${mergeOperation.oldest_event},
                     0
                 )
             `,
@@ -606,9 +618,9 @@ class PersonOverrideWriter {
                     UPDATE
                         posthog_personoverride
                     SET
-                        override_person_id = ${overridePersonId}, version = COALESCE(version, 0)::numeric + 1
+                        override_person_id = ${overridePersonMappingId}, version = COALESCE(version, 0)::numeric + 1
                     WHERE
-                        team_id = ${teamId} AND override_person_id = ${oldPersonId}
+                        team_id = ${mergeOperation.team_id} AND override_person_id = ${oldPersonMappingId}
                     RETURNING
                         old_person_id,
                         version,
@@ -637,21 +649,21 @@ class PersonOverrideWriter {
                 messages: [
                     {
                         value: JSON.stringify({
-                            team_id: teamId,
+                            team_id: mergeOperation.team_id,
+                            old_person_id: mergeOperation.old_person_id,
+                            override_person_id: mergeOperation.override_person_id,
+                            oldest_event: castTimestampOrNow(mergeOperation.oldest_event, TimestampFormat.ClickHouse),
                             merged_at: castTimestampOrNow(mergedAt, TimestampFormat.ClickHouse),
-                            override_person_id: overridePerson.uuid,
-                            old_person_id: oldPerson.uuid,
-                            oldest_event: castTimestampOrNow(oldestEvent, TimestampFormat.ClickHouse),
                             version: 0,
                         }),
                     },
                     ...transitiveUpdates.map(({ old_person_id, version, oldest_event }) => ({
                         value: JSON.stringify({
-                            team_id: teamId,
-                            merged_at: castTimestampOrNow(mergedAt, TimestampFormat.ClickHouse),
-                            override_person_id: overridePerson.uuid,
+                            team_id: mergeOperation.team_id,
                             old_person_id: old_person_id,
+                            override_person_id: mergeOperation.override_person_id,
                             oldest_event: castTimestampOrNow(oldest_event, TimestampFormat.ClickHouse),
+                            merged_at: castTimestampOrNow(mergedAt, TimestampFormat.ClickHouse),
                             version: version,
                         }),
                     })),
@@ -662,7 +674,7 @@ class PersonOverrideWriter {
         return personOverrideMessages
     }
 
-    private async addPersonOverrideMapping(tx: TransactionClient, person: Person): Promise<number> {
+    private async addPersonOverrideMapping(tx: TransactionClient, teamId: number, personId: string): Promise<number> {
         /**
             Update the helper table that serves as a mapping between a serial ID and a Person UUID.
 
@@ -684,8 +696,8 @@ class PersonOverrideWriter {
                         uuid
                     )
                     VALUES (
-                        ${person.team_id},
-                        '${person.uuid}'
+                        ${teamId},
+                        '${personId}'
                     )
                     ON CONFLICT("team_id", "uuid") DO NOTHING
                     RETURNING id
@@ -694,7 +706,7 @@ class PersonOverrideWriter {
                 UNION ALL
                 SELECT id
                 FROM posthog_personoverridemapping
-                WHERE uuid = '${person.uuid}'
+                WHERE team_id = ${teamId} AND uuid = '${personId}'
             `,
             undefined,
             'personOverrideMapping'
@@ -704,6 +716,91 @@ 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) {}
+
+    /**
+     * Enqueue an override for deferred processing.
+     */
+    public async addPersonOverride(tx: TransactionClient, mergeOperation: MergeOperation): Promise<ProducerRecord[]> {
+        await this.postgres.query(
+            tx,
+            SQL`
+            INSERT INTO posthog_pendingpersonoverride (
+                team_id,
+                old_person_id,
+                override_person_id,
+                oldest_event
+            ) VALUES (
+                ${mergeOperation.team_id},
+                ${mergeOperation.old_person_id},
+                ${mergeOperation.override_person_id},
+                ${mergeOperation.oldest_event}
+            )`,
+            undefined,
+            'pendingPersonOverride'
+        )
+
+        return []
+    }
+
+    /**
+     * Process all pending overrides. An advisory lock is acquired prior to
+     * processing to ensure that this function has exclusive access to the
+     * pending overrides during the update process.
+     *
+     * @returns the number of overrides processed
+     */
+    public async processPendingOverrides(kafkaProducer: KafkaProducerWrapper): Promise<number> {
+        const writer = new PersonOverrideWriter(this.postgres)
+
+        return await this.postgres.transaction(PostgresUse.COMMON_WRITE, 'processPendingOverrides', async (tx) => {
+            const {
+                rows: [{ acquired }],
+            } = await this.postgres.query(
+                tx,
+                SQL`SELECT pg_try_advisory_xact_lock(${this.lockId}) as acquired`,
+                undefined,
+                'processPendingOverrides'
+            )
+            if (!acquired) {
+                throw new Error('could not acquire lock')
+            }
+
+            // n.b.: Ordering by id ensures we are processing in (roughly) FIFO order
+            const { rows } = await this.postgres.query(
+                tx,
+                `SELECT * FROM posthog_pendingpersonoverride ORDER BY id`,
+                undefined,
+                'processPendingOverrides'
+            )
+
+            const messages: ProducerRecord[] = []
+            for (const { id, ...mergeOperation } of rows) {
+                messages.push(...(await writer.addPersonOverride(tx, mergeOperation)))
+                await this.postgres.query(
+                    tx,
+                    SQL`DELETE FROM posthog_pendingpersonoverride WHERE id = ${id}`,
+                    undefined,
+                    'processPendingOverrides'
+                )
+            }
+
+            // n.b.: We publish the messages here (and wait for acks) to ensure
+            // that all of our override updates are sent to Kafka before
+            // prior to committing the transaction. If we're unable to publish,
+            // we should discard updates and try again later when it's available
+            await kafkaProducer.queueMessages(messages, true)
+
+            return rows.length
+        })
+    }
+}
+
 export function ageInMonthsLowCardinality(timestamp: DateTime): number {
     const ageInMonths = Math.max(-Math.floor(timestamp.diffNow('months').months), 0)
     // for getting low cardinality for statsd metrics tags, which can cause issues in e.g. InfluxDB:

plugin-server/tests/helpers/promises.ts

@@ -13,3 +13,22 @@ export function createPromise<T = void>(): MockPromise<T> {
 
     return result as MockPromise<T>
 }
+
+export class WaitEvent {
+    private promise: Promise<void>
+    private resolve: () => void
+
+    constructor() {
+        this.promise = new Promise((resolve) => {
+            this.resolve = resolve
+        })
+    }
+
+    public set() {
+        this.resolve()
+    }
+
+    public async wait() {
+        return this.promise
+    }
+}