Fix up acks

plugin-server/functional_tests/api.ts

@@ -107,7 +107,6 @@ export const capture = async ({
             })
         ),
         key: teamId ? teamId.toString() : '',
-        waitForAck: true,
     })
 }
 

plugin-server/functional_tests/jobs-consumer.test.ts

@@ -43,7 +43,7 @@ describe('dlq handling', () => {
     test.concurrent(`handles empty messages`, async () => {
         const key = uuidv4()
 
-        await produce({ topic: 'jobs', message: null, key, waitForAck: true })
+        await produce({ topic: 'jobs', message: null, key })
 
         await waitForExpect(() => {
             const messages = dlq.filter((message) => message.key?.toString() === key)
@@ -54,7 +54,7 @@ describe('dlq handling', () => {
     test.concurrent(`handles invalid JSON`, async () => {
         const key = uuidv4()
 
-        await produce({ topic: 'jobs', message: Buffer.from('invalid json'), key, waitForAck: true })
+        await produce({ topic: 'jobs', message: Buffer.from('invalid json'), key })
 
         await waitForExpect(() => {
             const messages = dlq.filter((message) => message.key?.toString() === key)
@@ -72,7 +72,7 @@ describe('dlq handling', () => {
             labels: { topic: 'jobs', partition: '0', groupId: 'jobs-inserter' },
         })
 
-        await produce({ topic: 'jobs', message: Buffer.from(''), key: '', waitForAck: true })
+        await produce({ topic: 'jobs', message: Buffer.from(''), key: '' })
 
         await waitForExpect(async () => {
             const metricAfter = await getMetric({

plugin-server/functional_tests/kafka.ts

@@ -36,18 +36,8 @@ export async function createKafkaProducer() {
     return producer
 }
 
-export async function produce({
-    topic,
-    message,
-    key,
-    waitForAck,
-}: {
-    topic: string
-    message: Buffer | null
-    key: string
-    waitForAck: boolean
-}) {
+export async function produce({ topic, message, key }: { topic: string; message: Buffer | null; key: string }) {
     producer = producer ?? (await createKafkaProducer())
 
-    await new KafkaProducerWrapper(producer).produce({ topic, value: message, key: Buffer.from(key), waitForAck })
+    await new KafkaProducerWrapper(producer).produce({ topic, value: message, key: Buffer.from(key) })
 }

plugin-server/functional_tests/scheduled-tasks-runner.test.ts

@@ -43,7 +43,7 @@ describe('dlq handling', () => {
     test.concurrent(`handles empty messages`, async () => {
         const key = uuidv4()
 
-        await produce({ topic: 'scheduled_tasks', message: null, key, waitForAck: true })
+        await produce({ topic: 'scheduled_tasks', message: null, key })
 
         await waitForExpect(() => {
             const messages = dlq.filter((message) => message.key?.toString() === key)
@@ -54,7 +54,7 @@ describe('dlq handling', () => {
     test.concurrent(`handles invalid JSON`, async () => {
         const key = uuidv4()
 
-        await produce({ topic: 'scheduled_tasks', message: Buffer.from('invalid json'), key, waitForAck: true })
+        await produce({ topic: 'scheduled_tasks', message: Buffer.from('invalid json'), key })
 
         await waitForExpect(() => {
             const messages = dlq.filter((message) => message.key?.toString() === key)
@@ -69,7 +69,7 @@ describe('dlq handling', () => {
             topic: 'scheduled_tasks',
             message: Buffer.from(JSON.stringify({ taskType: 'invalidTaskType', pluginConfigId: 1 })),
             key,
-            waitForAck: true,
+    ,
         })
 
         await waitForExpect(() => {
@@ -85,7 +85,7 @@ describe('dlq handling', () => {
             topic: 'scheduled_tasks',
             message: Buffer.from(JSON.stringify({ taskType: 'runEveryMinute', pluginConfigId: 'asdf' })),
             key,
-            waitForAck: true,
+    ,
         })
 
         await waitForExpect(() => {
@@ -106,7 +106,7 @@ describe('dlq handling', () => {
 
         // NOTE: we don't actually care too much about the contents of the
         // message, just that it triggeres the consumer to try to process it.
-        await produce({ topic: 'scheduled_tasks', message: Buffer.from(''), key: '', waitForAck: true })
+        await produce({ topic: 'scheduled_tasks', message: Buffer.from(''), key: '' })
 
         await waitForExpect(async () => {
             const metricAfter = await getMetric({

plugin-server/functional_tests/session-recordings.test.ts

@@ -177,7 +177,6 @@ test.skip('consumer updates timestamp exported to prometheus', async () => {
         topic: KAFKA_SESSION_RECORDING_SNAPSHOT_ITEM_EVENTS,
         message: Buffer.from(''),
         key: '',
-        waitForAck: true,
     })
 
     await waitForExpect(async () => {
@@ -250,15 +249,13 @@ test.skip(`handles message with no token or with token and no associated team_id
         topic: KAFKA_SESSION_RECORDING_SNAPSHOT_ITEM_EVENTS,
         message: Buffer.from(JSON.stringify({ uuid: noTokenUuid, data: JSON.stringify({}) })),
         key: noTokenKey,
-        waitForAck: true,
     })
     await produce({
         topic: KAFKA_SESSION_RECORDING_SNAPSHOT_ITEM_EVENTS,
         message: Buffer.from(
             JSON.stringify({ uuid: noAssociatedTeamUuid, token: 'no associated team', data: JSON.stringify({}) })
         ),
         key: noAssociatedTeamKey,
-        waitForAck: true,
     })
 
     await capture(makeSessionMessage(teamId, 'should be ingested'))

plugin-server/src/cdp/cdp-consumers.ts

@@ -167,7 +167,6 @@ abstract class CdpConsumerBase {
                     },
                 ],
             })),
-            waitForAck: true,
         }).catch((reason) => {
             status.error('⚠️', `failed to produce message: ${reason}`)
         })

plugin-server/src/kafka/producer.ts

@@ -91,24 +91,19 @@ export class KafkaProducerWrapper {
         key,
         topic,
         headers,
-        waitForAck,
     }: {
         value: MessageValue
         key: MessageKey
         topic: string
         headers?: MessageHeader[]
-        waitForAck: boolean
     }): Promise<void> {
         try {
+            const produceTimer = ingestEventKafkaProduceLatency.labels({ topic }).startTimer()
+            const produceSpan = getSpan()?.startChild({ op: 'kafka_produce' })
             kafkaProducerMessagesQueuedCounter.labels({ topic_name: topic }).inc()
-
             status.debug('📤', 'Producing message', { topic: topic })
-            const produceSpan = getSpan()?.startChild({ op: 'kafka_produce' })
-            return await new Promise((resolve, reject) => {
-                const produceTimer = ingestEventKafkaProduceLatency
-                    .labels({ topic, waitForAck: waitForAck.toString() })
-                    .startTimer()
 
+            const result = await new Promise((resolve, reject) => {
                 this.producer.produce(
                     topic,
                     null,
@@ -117,28 +112,15 @@ export class KafkaProducerWrapper {
                     Date.now(),
                     headers ?? [],
                     (error: any, offset: NumberNullUndefined) => {
-                        produceSpan?.finish()
-                        kafkaProducerMessagesWrittenCounter.labels({ topic_name: topic }).inc()
-                        if (error) {
-                            status.error('⚠️', 'produce_error', { error: error, topic: topic })
-                        } else {
-                            status.debug('📤', 'Produced message', { topic: topic, offset: offset })
-                        }
-
-                        if (waitForAck) {
-                            produceTimer()
-                            // TODO: Modify this to return the offset in a way that works across the codebase
-                            return error ? reject(error) : resolve()
-                        }
+                        return error ? reject(error) : resolve(offset)
                     }
                 )
-
-                // If we're not waiting for an ack, we can resolve immediately
-                if (!waitForAck) {
-                    resolve(undefined)
-                    produceTimer()
-                }
             })
+
+            produceSpan?.finish()
+            kafkaProducerMessagesWrittenCounter.labels({ topic_name: topic }).inc()
+            status.debug('📤', 'Produced message', { topic: topic, offset: result })
+            produceTimer()
         } catch (error) {
             kafkaProducerMessagesFailedCounter.labels({ topic_name: topic }).inc()
             status.error('⚠️', 'kafka_produce_error', { error: error, topic: topic })
@@ -157,10 +139,8 @@ export class KafkaProducerWrapper {
 
     async queueMessages({
         kafkaMessages: kafkaMessage,
-        waitForAck,
     }: {
         kafkaMessages: ProducerRecord | ProducerRecord[]
-        waitForAck: boolean
     }): Promise<void> {
         const records = Array.isArray(kafkaMessage) ? kafkaMessage : [kafkaMessage]
 
@@ -173,7 +153,6 @@ export class KafkaProducerWrapper {
                             key: message.key ? Buffer.from(message.key) : null,
                             value: message.value ? Buffer.from(message.value) : null,
                             headers: convertKafkaJSHeadersToRdKafkaHeaders(message.headers),
-                            waitForAck: waitForAck,
                         })
                     )
                 )
@@ -255,6 +234,6 @@ export const kafkaProducerMessagesFailedCounter = new Counter({
 export const ingestEventKafkaProduceLatency = new Summary({
     name: 'ingest_event_kafka_produce_latency',
     help: 'Wait time for individual Kafka produces',
-    labelNames: ['topic', 'waitForAck'],
+    labelNames: ['topic'],
     percentiles: [0.5, 0.9, 0.95, 0.99],
 })

plugin-server/src/main/graphile-worker/schedule.ts

@@ -60,7 +60,6 @@ export async function runScheduledTasks(
                     topic: KAFKA_SCHEDULED_TASKS,
                     messages: [{ key: pluginConfigId.toString(), value: JSON.stringify({ taskType, pluginConfigId }) }],
                 },
-                waitForAck: true,
             })
             graphileScheduledTaskCounter.labels({ status: 'queued', task: taskType }).inc()
         }

plugin-server/src/main/ingestion-queues/batch-processing/each-batch-ingestion.ts

@@ -80,7 +80,6 @@ async function handleProcessingError(
                 value: message.value,
                 key: message.key ?? null, // avoid undefined, just to be safe
                 headers: headers,
-                waitForAck: true,
             })
         } catch (error) {
             // If we can't send to the DLQ and it's not retriable, just continue. We'll commit the
@@ -282,7 +281,6 @@ async function emitToOverflow(queue: IngestionConsumer, kafkaMessages: Message[]
                 // instead as that behavior is safer.
                 key: useRandomPartitioning ? null : message.key ?? null,
                 headers: message.headers,
-                waitForAck: true,
             })
         )
     )

plugin-server/src/main/ingestion-queues/jobs-consumer.ts

@@ -59,7 +59,6 @@ export const startJobsConsumer = async ({
                         topic: KAFKA_JOBS_DLQ,
                         messages: [{ value: message.value, key: message.key }],
                     },
-                    waitForAck: true,
                 })
                 resolveOffset(message.offset)
                 continue
@@ -79,7 +78,6 @@ export const startJobsConsumer = async ({
                         topic: KAFKA_JOBS_DLQ,
                         messages: [{ value: message.value, key: message.key }],
                     },
-                    waitForAck: true,
                 })
                 resolveOffset(message.offset)
                 continue

plugin-server/src/main/ingestion-queues/scheduled-tasks-consumer.ts

@@ -169,7 +169,6 @@ const getTasksFromBatch = async (batch: Batch, producer: KafkaProducerWrapper) =
                     topic: KAFKA_SCHEDULED_TASKS_DLQ,
                     messages: [{ value: message.value, key: message.key }],
                 },
-                waitForAck: true,
             })
             continue
         }
@@ -190,7 +189,6 @@ const getTasksFromBatch = async (batch: Batch, producer: KafkaProducerWrapper) =
                     topic: KAFKA_SCHEDULED_TASKS_DLQ,
                     messages: [{ value: message.value, key: message.key }],
                 },
-                waitForAck: true,
             })
             continue
         }
@@ -202,7 +200,6 @@ const getTasksFromBatch = async (batch: Batch, producer: KafkaProducerWrapper) =
                     topic: KAFKA_SCHEDULED_TASKS_DLQ,
                     messages: [{ value: message.value, key: message.key }],
                 },
-                waitForAck: true,
             })
             continue
         }

plugin-server/src/main/ingestion-queues/session-recording/services/console-logs-ingester.ts

@@ -163,7 +163,6 @@ export class ConsoleLogsIngester {
                             key: event.session_id,
                         })),
                     },
-                    waitForAck: true,
                 }),
             ]
         } catch (error) {

plugin-server/src/main/ingestion-queues/session-recording/services/replay-events-ingester.ts

@@ -180,7 +180,6 @@ export class ReplayEventsIngester {
                             },
                         ],
                     },
-                    waitForAck: true,
                 }),
             ]
         } catch (error) {

plugin-server/src/utils/db/db.ts

@@ -697,7 +697,7 @@ export class DB {
             })
         }
 
-        await this.kafkaProducer.queueMessages({ kafkaMessages, waitForAck: true })
+        await this.kafkaProducer.queueMessages({ kafkaMessages })
         return person
     }
 
@@ -881,7 +881,7 @@ export class DB {
     ): Promise<void> {
         const kafkaMessages = await this.addDistinctIdPooled(person, distinctId, version, tx)
         if (kafkaMessages.length) {
-            await this.kafkaProducer.queueMessages({ kafkaMessages: kafkaMessages, waitForAck: true })
+            await this.kafkaProducer.queueMessages({ kafkaMessages })
         }
     }
 
@@ -1435,7 +1435,6 @@ export class DB {
                     },
                 ],
             },
-            waitForAck: true,
         })
     }
 

plugin-server/src/utils/db/hub.ts

@@ -158,7 +158,6 @@ export async function createHub(
                     },
                 ],
             },
-            waitForAck: true,
         })
     }
 

plugin-server/src/worker/ingestion/app-metrics.ts

@@ -187,7 +187,6 @@ export class AppMetrics {
                 topic: KAFKA_APP_METRICS,
                 messages: kafkaMessages,
             },
-            waitForAck: true,
         })
         status.debug('🚽', `Finished flushing app metrics, took ${Date.now() - startTime}ms`)
     }

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

@@ -43,7 +43,6 @@ export async function extractHeatmapDataStep(
                             value: JSON.stringify(rawEvent),
                         })),
                     },
-                    waitForAck: true,
                 })
             )
         }

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

@@ -17,7 +17,6 @@ export function produceExceptionSymbolificationEventStep(
                     },
                 ],
             },
-            waitForAck: true,
         })
         .catch((error) => {
             status.warn('⚠️', 'Failed to produce exception event for symbolification', {

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

@@ -352,7 +352,7 @@ export class EventPipelineRunner {
                     teamId,
                     `plugin_server_ingest_event:${currentStepName}`
                 )
-                await this.hub.db.kafkaProducer!.queueMessages({ kafkaMessages: message, waitForAck: true })
+                await this.hub.db.kafkaProducer!.queueMessages({ kafkaMessages: message })
             } catch (dlqError) {
                 status.info('🔔', `Errored trying to add event to dead letter queue. Error: ${dlqError}`)
                 Sentry.captureException(dlqError, {

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

@@ -301,7 +301,7 @@ export class PersonState {
 
         if (Object.keys(update).length > 0) {
             const [updatedPerson, kafkaMessages] = await this.db.updatePersonDeprecated(person, update)
-            const kafkaAck = this.db.kafkaProducer.queueMessages({ kafkaMessages, waitForAck: true })
+            const kafkaAck = this.db.kafkaProducer.queueMessages({ kafkaMessages })
             return [updatedPerson, kafkaAck]
         }
 
@@ -768,7 +768,7 @@ export class PersonState {
             })
             .inc()
 
-        const kafkaAck = this.db.kafkaProducer.queueMessages({ kafkaMessages, waitForAck: true })
+        const kafkaAck = this.db.kafkaProducer.queueMessages({ kafkaMessages })
 
         return [mergedPerson, kafkaAck]
     }