Optimize Kafka queue offset handling

src/ingestion/kafka-queue.ts

@@ -1,13 +1,48 @@
 import * as Sentry from '@sentry/node'
 import { Kafka, Consumer, Message, EachBatchPayload } from 'kafkajs'
 import { PluginsServer, Queue, RawEventMessage } from 'types'
-import { KAFKA_EVENTS_INGESTION_HANDOFF } from './topics'
+import { KAFKA_EVENTS_WAL } from './topics'
 import { PluginEvent } from '@posthog/plugin-scaffold'
 import { status } from '../status'
-import { killGracefully } from '../utils'
 
 export type BatchCallback = (messages: Message[]) => Promise<void>
 
+/**
+ * We use this to avoid a situation where the event was discarded and because of that its Kafka offset is not resolved,
+ * potentially causing unnecessary message retries.
+ *
+ * Works in the following way: for each event UUID the algorithm sets the Kafka offset to resolve
+ * based on the last event after it from the Kafka topic that WAS discarded. Example, we've got events:
+ * ```JS
+ * [{ uuid: 'a', offset: 1 }, { uuid: 's', offset: 2 }, { uuid: 'd', offset: 3 }, { uuid: 'f', offset: 4 }]
+ * ```
+ * Now some plugin discards the last two! Returned map for use in resolveOffset() will in result look like this:
+ * ```JS
+ * { 'a': 1, 's': 4, 'd': 4, 'f': 4 }
+ * ```
+ * Because 'd' and 'f' were discarded by a plugin, when we save 's' to the database, we'll know that at the same time
+ * we in fact also covered 'd' and 'f' and should resolve the last offset - belonging to 'f' – not an intermediary one.
+ * As for 'a', we simply resolve its offset outright, because its next event ('s') is not discarded.
+ */
+function aliasEventUuidForDiscardedKafkaOffsets(
+    rawEventMessages: RawEventMessage[],
+    processedEvents: PluginEvent[]
+): Map<string, string> {
+    rawEventMessages = [...rawEventMessages] // Shallow copy to avoid side effects
+    const eventUuidToKafkaOffset = new Map<string, string>()
+    const processedUuids: Set<string> = new Set(processedEvents.map((event) => event.uuid!))
+    rawEventMessages.reverse()
+    // This initial value below is in fact the last message, due to the array being reversed
+    let currentNotDiscardedOffset: string = rawEventMessages.shift()!.kafka_offset
+    for (const rawEventMessage of rawEventMessages) {
+        if (processedUuids.has(rawEventMessage.uuid)) {
+            currentNotDiscardedOffset = rawEventMessage.kafka_offset
+        }
+        eventUuidToKafkaOffset.set(rawEventMessage.uuid, currentNotDiscardedOffset)
+    }
+    return eventUuidToKafkaOffset
+}
+
 export class KafkaQueue implements Queue {
     private pluginsServer: PluginsServer
     private kafka: Kafka
@@ -39,27 +74,23 @@ export class KafkaQueue implements Queue {
         isStale,
     }: EachBatchPayload): Promise<void> {
         const batchProcessingTimer = new Date()
-        const rawEvents: RawEventMessage[] = batch.messages.map((message) => ({
+        const rawEventMessages: RawEventMessage[] = batch.messages.map((message) => ({
             ...JSON.parse(message.value!.toString()),
             kafka_offset: message.offset,
         }))
-        const parsedEvents = rawEvents.map((rawEvent) => ({
+        const parsedEventMessages = rawEventMessages.map((rawEvent) => ({
             ...rawEvent,
             data: JSON.parse(rawEvent.data),
         }))
-        const pluginEvents: PluginEvent[] = rawEvents.map((rawEvent) => {
-            const { data: dataStr, ...restOfRawEvent } = rawEvent
-            const event = { ...restOfRawEvent, ...JSON.parse(dataStr) }
-            return {
-                ...event,
-                kafka_offset: restOfRawEvent.kafka_offset,
-                site_url: event.site_url || null,
-                ip: event.ip || null,
-            }
-        })
+        const pluginEvents: PluginEvent[] = parsedEventMessages.map((parsedEvent) => ({
+            ...parsedEvent,
+            event: parsedEvent.data.event,
+            properties: parsedEvent.data.properties,
+        }))
         const processedEvents: PluginEvent[] = (
             await this.processEventBatch(pluginEvents)
         ).filter((event: PluginEvent[] | false | null | undefined) => Boolean(event))
+        const eventUuidToKafkaOffset = aliasEventUuidForDiscardedKafkaOffsets(rawEventMessages, processedEvents)
         for (const event of processedEvents) {
             if (!isRunning()) {
                 status.info('😮', 'Consumer not running anymore, canceling batch processing!')
@@ -71,7 +102,7 @@ export class KafkaQueue implements Queue {
             }
             const singleIngestionTimer = new Date()
             await this.saveEvent(event)
-            resolveOffset(event.kafka_offset!)
+            resolveOffset(eventUuidToKafkaOffset.get(event.uuid!)!)
             await heartbeat()
             await commitOffsetsIfNecessary()
             this.pluginsServer.statsd?.timing('kafka_queue.single_ingestion', singleIngestionTimer)
@@ -85,7 +116,7 @@ export class KafkaQueue implements Queue {
             this.consumer.on(this.consumer.events.CRASH, ({ payload: { error } }) => reject(error))
             status.info('⏬', `Connecting Kafka consumer to ${this.pluginsServer.KAFKA_HOSTS}...`)
             this.wasConsumerRan = true
-            await this.consumer.subscribe({ topic: KAFKA_EVENTS_INGESTION_HANDOFF })
+            await this.consumer.subscribe({ topic: KAFKA_EVENTS_WAL })
             // KafkaJS batching: https://kafka.js.org/docs/consuming#a-name-each-batch-a-eachbatch
             await this.consumer.run({
                 // TODO: eachBatchAutoResolve: false, // don't autoresolve whole batch in case we exit it early
@@ -104,7 +135,7 @@ export class KafkaQueue implements Queue {
             return
         }
         status.info('⏳', 'Pausing Kafka consumer...')
-        await this.consumer.pause([{ topic: KAFKA_EVENTS_INGESTION_HANDOFF }])
+        await this.consumer.pause([{ topic: KAFKA_EVENTS_WAL }])
         status.info('⏸', 'Kafka consumer paused!')
     }
 
@@ -113,12 +144,12 @@ export class KafkaQueue implements Queue {
             return
         }
         status.info('⏳', 'Resuming Kafka consumer...')
-        await this.consumer.resume([{ topic: KAFKA_EVENTS_INGESTION_HANDOFF }])
+        await this.consumer.resume([{ topic: KAFKA_EVENTS_WAL }])
         status.info('▶️', 'Kafka consumer resumed!')
     }
 
     isPaused(): boolean {
-        return this.consumer.paused().some(({ topic }) => topic === KAFKA_EVENTS_INGESTION_HANDOFF)
+        return this.consumer.paused().some(({ topic }) => topic === KAFKA_EVENTS_WAL)
     }
 
     async stop(): Promise<void> {
@@ -136,7 +167,7 @@ export class KafkaQueue implements Queue {
 
     private static buildConsumer(kafka: Kafka): Consumer {
         const consumer = kafka.consumer({
-            groupId: 'clickhouse-ingestion',
+            groupId: 'plugin-server',
             readUncommitted: false,
         })
         const { GROUP_JOIN, CRASH, CONNECT, DISCONNECT } = consumer.events
@@ -146,7 +177,6 @@ export class KafkaQueue implements Queue {
         consumer.on(CRASH, ({ payload: { error, groupId } }) => {
             status.error('⚠️', `Kafka consumer group ${groupId} crashed:\n`, error)
             Sentry.captureException(error)
-            killGracefully()
         })
         consumer.on(CONNECT, () => {
             status.info('✅', 'Kafka consumer connected!')