From 0241415c46bab546b70a061e504958de202e1c53 Mon Sep 17 00:00:00 2001 From: Matheus Veras Batista Date: Wed, 26 Aug 2026 19:36:26 -0300 Subject: [PATCH 1/9] feat(cohorts): sweep membership rows a completed reconcile did not assert A person who stops matching an edited cohort gets no live transition, so their in_cohort = true row survives forever. A reconcile run replays full membership and certifies itself with per-partition markers; the sweep deletes what the run did not re-assert. Rows are stamped with the producer version, and the oldest version a run asserted bounds the sweep threshold. Rows, marks, and consumer progress commit in one transaction. Deletes wait until the snapshot minimum is known and progress passes the marker-completion watermarks; progress also refreshes from committed offsets so quiet partitions cannot stall the gate. Gated by COHORT_MEMBERSHIP_SWEEP_ENABLED, off by default. Version persistence stays ungated: a mixed fleet writing NULLs would get current members swept. --- nodejs/src/cdp/config.ts | 15 + .../cdp-cohort-membership.consumer.test.ts | 203 +++++++ .../cdp-cohort-membership.consumer.ts | 378 ++++++++++-- nodejs/src/cdp/dev/cohort-sweep-e2e.test.ts | 131 +++++ .../cohort-membership/sweeper.service.test.ts | 252 ++++++++ .../cohort-membership/sweeper.service.ts | 545 ++++++++++++++++++ nodejs/src/common/config/kafka-topics.ts | 3 + .../src/common/kafka/consumer/consumer-v1.ts | 8 + .../src/common/kafka/consumer/consumer-v2.ts | 42 ++ nodejs/src/common/kafka/consumer/index.ts | 8 +- nodejs/tests/helpers/sql.ts | 8 +- 11 files changed, 1552 insertions(+), 41 deletions(-) create mode 100644 nodejs/src/cdp/dev/cohort-sweep-e2e.test.ts create mode 100644 nodejs/src/cdp/services/cohort-membership/sweeper.service.test.ts create mode 100644 nodejs/src/cdp/services/cohort-membership/sweeper.service.ts diff --git a/nodejs/src/cdp/config.ts b/nodejs/src/cdp/config.ts index 6a6842d38ef5..363ac8d1a747 100644 --- a/nodejs/src/cdp/config.ts +++ b/nodejs/src/cdp/config.ts @@ -186,6 +186,13 @@ export type CdpConfig = ClickhouseConfig & { CYCLOTRON_NODE_RESCHEDULE_MAX_CHUNKS_PER_CALL: number CYCLOTRON_NODE_RESCHEDULE_CHUNK_SLEEP_MS: number + // Mark-and-sweep deletion of cohort_membership rows a completed reconcile run did not re-assert + COHORT_MEMBERSHIP_SWEEP_ENABLED: boolean + COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS: number + COHORT_MEMBERSHIP_SWEEP_BATCH_SIZE: number + COHORT_MEMBERSHIP_SWEEP_CLAIM_TIMEOUT_MS: number + COHORT_MEMBERSHIP_SWEEP_ABANDON_AFTER_DAYS: number + // Email reputation evaluator (daily Temporal-scheduled bounce/complaint snapshots for workflows email) } @@ -345,5 +352,13 @@ export function getDefaultCdpConfig(): CdpConfig { CYCLOTRON_NODE_RESCHEDULE_CHUNK_SIZE: 5000, CYCLOTRON_NODE_RESCHEDULE_MAX_CHUNKS_PER_CALL: 20, CYCLOTRON_NODE_RESCHEDULE_CHUNK_SLEEP_MS: 100, + + // Stays off until the processor writes membership to the live topic and the whole fleet + // persists row versions. Sweeping before either would delete rows nothing re-asserts. + COHORT_MEMBERSHIP_SWEEP_ENABLED: false, + COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS: 60000, + COHORT_MEMBERSHIP_SWEEP_BATCH_SIZE: 1000, + COHORT_MEMBERSHIP_SWEEP_CLAIM_TIMEOUT_MS: 300000, + COHORT_MEMBERSHIP_SWEEP_ABANDON_AFTER_DAYS: 3, } } diff --git a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts index 9c3b48592d0b..18f283b42eb0 100644 --- a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts +++ b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts @@ -255,6 +255,189 @@ describe('CdpCohortMembershipConsumer', () => { }) }) + // The producer's format, which is fixed-width so that string order is chronological. + const OLDER_VERSION = '2026-05-26 12:00:00.000000' + const VERSION = '2026-05-26 12:34:56.789123' + const NEWER_VERSION = '2026-05-26 13:00:00.000000' + + const readVersion = async (cohortId: number): Promise => { + const result = await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + `SELECT in_cohort, to_char(version, 'YYYY-MM-DD HH24:MI:SS.US') AS version + FROM cohort_membership WHERE team_id = 1 AND person_id = $1 AND cohort_id = $2`, + [personId1, cohortId], + 'testQuery' + ) + return result.rows[0].version + } + + it.each([ + ['a live transition', {}], + ['a seed row', { origin: 'seed' as const, run_id: new UUIDT().toString() }], + ['a reconcile row', { origin: 'reconcile' as const, run_id: new UUIDT().toString() }], + ])('should persist the message version for %s', async (_label, tags) => { + const message = createKafkaMessage( + createCohortMembershipEvent({ + person_id: personId1, + cohort_id: 456, + team_id: 1, + last_updated: VERSION, + ...tags, + }), + { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: 0 } + ) + + await consumer['handleBatch']([message]) + + expect(await readVersion(456)).toEqual(VERSION) + }) + + it.each([ + ['an older version is rejected', VERSION, OLDER_VERSION, true, VERSION], + ['an equal version is applied', VERSION, VERSION, false, VERSION], + ['a newer version is applied', VERSION, NEWER_VERSION, false, NEWER_VERSION], + ['any version beats a versionless row', undefined, OLDER_VERSION, false, OLDER_VERSION], + ])( + 'should apply last-writer-wins on replay: %s', + async (_label, storedVersion, incomingVersion, expectedInCohort, expectedVersion) => { + await consumer['handleBatch']([ + createKafkaMessage( + createCohortMembershipEvent({ + person_id: personId1, + cohort_id: 456, + team_id: 1, + status: 'entered', + last_updated: storedVersion, + }), + { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: 0 } + ), + ]) + + await consumer['handleBatch']([ + createKafkaMessage( + createCohortMembershipEvent({ + person_id: personId1, + cohort_id: 456, + team_id: 1, + status: 'left', + last_updated: incomingVersion, + }), + { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: 1 } + ), + ]) + + const result = await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + `SELECT in_cohort, to_char(version, 'YYYY-MM-DD HH24:MI:SS.US') AS version + FROM cohort_membership WHERE team_id = 1 AND person_id = $1 AND cohort_id = 456`, + [personId1], + 'testQuery' + ) + + expect(result.rows[0].in_cohort).toBe(expectedInCohort) + expect(result.rows[0].version).toEqual(expectedVersion) + } + ) + + it('should track the oldest version each reconcile run asserted, across batches', async () => { + const runId = new UUIDT().toString() + const reconcileRow = (personId: string, lastUpdated: string) => + createKafkaMessage( + createCohortMembershipEvent({ + person_id: personId, + cohort_id: 456, + team_id: 1, + origin: 'reconcile', + run_id: runId, + last_updated: lastUpdated, + }), + { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: 0 } + ) + + await consumer['handleBatch']([ + reconcileRow(personId1, NEWER_VERSION), + reconcileRow(personId2, VERSION), + // A live transition rides the same topic and must not count as an assertion. + createKafkaMessage( + createCohortMembershipEvent({ + person_id: personId3, + cohort_id: 456, + team_id: 1, + last_updated: OLDER_VERSION, + }), + { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: 0 } + ), + ]) + + const readSweep = async () => { + const result = await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + `SELECT team_id, marker_bits, status, snapshot_rows, + to_char(min_snapshot_version, 'YYYY-MM-DD HH24:MI:SS.US') AS min_snapshot_version + FROM cohort_membership_sweeps WHERE run_id = $1 AND cohort_id = 456`, + [runId], + 'testQuery' + ) + return result.rows + } + + expect(await readSweep()).toEqual([ + { + team_id: '1', + marker_bits: '0', + status: 'collecting', + snapshot_rows: '2', + min_snapshot_version: VERSION, + }, + ]) + + await consumer['handleBatch']([reconcileRow(personId3, OLDER_VERSION)]) + + expect(await readSweep()).toEqual([ + { + team_id: '1', + marker_bits: '0', + status: 'collecting', + snapshot_rows: '3', + min_snapshot_version: OLDER_VERSION, + }, + ]) + }) + + it('should advance consumer progress to the next offset without ever regressing it', async () => { + const readProgress = async () => { + const result = await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + 'SELECT partition, next_offset FROM cohort_membership_consumer_progress ORDER BY partition', + undefined, + 'testQuery' + ) + return result.rows + } + + const rowOnPartition = (partition: number, offset: number) => + createKafkaMessage(createCohortMembershipEvent({ person_id: new UUIDT().toString(), team_id: 1 }), { + topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, + partition, + offset, + }) + + await consumer['handleBatch']([rowOnPartition(3, 10), rowOnPartition(3, 11), rowOnPartition(7, 5)]) + + expect(await readProgress()).toEqual([ + { partition: 3, next_offset: '12' }, + { partition: 7, next_offset: '6' }, + ]) + + // A rebalance can replay from an older committed offset; the gate must not walk back. + await consumer['handleBatch']([rowOnPartition(3, 0)]) + + expect(await readProgress()).toEqual([ + { partition: 3, next_offset: '12' }, + { partition: 7, next_offset: '6' }, + ]) + }) + it('should reject entire batch when invalid messages are present', async () => { const validEvent = { person_id: personId1, @@ -298,6 +481,26 @@ describe('CdpCohortMembershipConsumer', () => { expect(result.rows).toHaveLength(0) }) + it('should not record consumer progress when the batch write fails', async () => { + // A non-integer cohort_id passes schema validation but Postgres rejects it, so the + // membership write fails while the progress write on its own would have succeeded. + const message = createKafkaMessage( + createCohortMembershipEvent({ person_id: personId1, team_id: 1, cohort_id: 1.5 }), + { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, partition: 9, offset: 42 } + ) + + await expect(consumer['handleBatch']([message])).rejects.toThrow() + + const progress = await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + 'SELECT * FROM cohort_membership_consumer_progress', + undefined, + 'testQuery' + ) + + expect(progress.rows).toHaveLength(0) + }) + it('should not produce side effects when database insertion fails', async () => { const testEvents = createCohortMembershipEvents([ { diff --git a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts index 12095a0300b1..11776e5fa954 100644 --- a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts +++ b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts @@ -1,16 +1,31 @@ import { Message } from 'node-rdkafka' import { z } from 'zod' -import { KAFKA_COHORT_MEMBERSHIP_CHANGED } from '~/common/config/kafka-topics' -import { KafkaConsumerInterface, createKafkaConsumer } from '~/common/kafka/consumer' +import { KAFKA_COHORT_MEMBERSHIP_CHANGED, KAFKA_COHORT_RECONCILE_MARKERS } from '~/common/config/kafka-topics' +import { KafkaConsumerInterface, RdKafkaConsumerConfig, createKafkaConsumer } from '~/common/kafka/consumer' import { instrumentFn } from '~/common/tracing/tracing-utils' -import { PostgresUse } from '~/common/utils/db/postgres' +import { PostgresUse, TransactionClient } from '~/common/utils/db/postgres' import { parseJSON } from '~/common/utils/json-parse' import { logger } from '~/common/utils/logger' import { HealthCheckResult } from '../../types' +import type { CdpConfig } from '../config' +import { CohortMembershipSweeper, MembershipWatermarks } from '../services/cohort-membership/sweeper.service' import { CdpConsumerBase, CdpConsumerBaseConfig, CdpConsumerBaseDeps } from './cdp-base.consumer' +export type CdpCohortMembershipConsumerConfig = CdpConsumerBaseConfig & + Pick< + CdpConfig, + | 'COHORT_MEMBERSHIP_SWEEP_ENABLED' + | 'COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS' + | 'COHORT_MEMBERSHIP_SWEEP_BATCH_SIZE' + | 'COHORT_MEMBERSHIP_SWEEP_CLAIM_TIMEOUT_MS' + | 'COHORT_MEMBERSHIP_SWEEP_ABANDON_AFTER_DAYS' + > + +// Markers from runs that predate the sweeper are useless — the next reconcile heals anyway. +const startAtLatest = { ['auto.offset.reset' as keyof RdKafkaConsumerConfig]: 'latest' as never } + // `origin` is a loose string, not an enum: nothing branches on it, and rejecting a value the // processor added later would throw the whole batch and crash-loop on redelivery. const CohortMembershipChangeSchema = z.object({ @@ -25,20 +40,173 @@ const CohortMembershipChangeSchema = z.object({ export type CohortMembershipChange = z.infer -export class CdpCohortMembershipConsumer extends CdpConsumerBase { +/** How far the consumer has applied a partition, in the same units as a Kafka high watermark. */ +type PartitionProgress = { + partition: number + nextOffset: number +} + +/** What one reconcile run asserted in a single batch, folded into the run's running minimum. */ +type SnapshotMark = { + runId: string + cohortId: number + teamId: number + minVersion: string + rows: number +} + +/** + * A high watermark is the next offset to be produced, so consumer progress is recorded the same + * way — as the next offset to consume. Comparing the two is then exact rather than off by one. + */ +function partitionProgressFromMessages(messages: Message[]): PartitionProgress[] { + const maxOffsets = new Map() + + for (const message of messages) { + const current = maxOffsets.get(message.partition) + if (current === undefined || message.offset > current) { + maxOffsets.set(message.partition, message.offset) + } + } + + return Array.from(maxOffsets, ([partition, offset]) => ({ partition, nextOffset: offset + 1 })) +} + +/** + * The oldest version each reconcile run asserted in this batch. A sweep may only delete rows below + * every version its run re-asserted, so the run's minimum is folded across batches in Postgres. + * + * Versions are compared as strings: the producer's format is fixed-width, so lexicographic order + * is chronological order. A row carrying no version bounds nothing and is left out. + */ +function collectSnapshotMarks(changes: CohortMembershipChange[]): SnapshotMark[] { + const marks = new Map() + + for (const change of changes) { + if (change.origin !== 'reconcile' || !change.run_id) { + continue + } + + if (!change.last_updated) { + logger.warn('Reconcile membership change without a version, excluded from run marking', { + run_id: change.run_id, + cohort_id: change.cohort_id, + team_id: change.team_id, + }) + continue + } + + const key = `${change.run_id}:${change.cohort_id}` + const existing = marks.get(key) + + if (!existing) { + marks.set(key, { + runId: change.run_id, + cohortId: change.cohort_id, + teamId: change.team_id, + minVersion: change.last_updated, + rows: 1, + }) + continue + } + + existing.rows += 1 + if (change.last_updated < existing.minVersion) { + existing.minVersion = change.last_updated + } + } + + return Array.from(marks.values()) +} + +export class CdpCohortMembershipConsumer extends CdpConsumerBase { protected name = 'CdpCohortMembershipConsumer' private kafkaConsumer: KafkaConsumerInterface + private markerKafkaConsumer: KafkaConsumerInterface | null = null + private sweeper: CohortMembershipSweeper | null = null - constructor(config: CdpConsumerBaseConfig, deps: CdpConsumerBaseDeps) { + constructor(config: CdpCohortMembershipConsumerConfig, deps: CdpConsumerBaseDeps) { super(config, deps) this.kafkaConsumer = createKafkaConsumer({ groupId: 'cdp-cohort-membership-consumer', topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, }) + + if (config.COHORT_MEMBERSHIP_SWEEP_ENABLED) { + this.markerKafkaConsumer = createKafkaConsumer( + { groupId: 'cdp-cohort-membership-sweeper', topic: KAFKA_COHORT_RECONCILE_MARKERS }, + startAtLatest + ) + this.sweeper = new CohortMembershipSweeper(config, deps.postgres, { + captureMembershipWatermarks: () => this.captureMembershipWatermarks(), + refreshConsumerProgress: () => this.refreshConsumerProgressFromCommittedOffsets(), + }) + } } - private async persistCohortMembershipChanges(changes: CohortMembershipChange[]): Promise { - if (changes.length === 0) { + /** + * A partition only gets a progress row when it delivers a batch, so one that holds retained + * messages but receives no new traffic would block the gate forever. Committed offsets cover + * every assigned partition, and the consumer commits only after a batch is applied, so they + * are a safe lower bound on what reached Postgres. Each pod refreshes its own assignment; + * together the fleet covers the topic. + */ + private async refreshConsumerProgressFromCommittedOffsets(): Promise { + const committed = await this.kafkaConsumer.committedOffsets() + + const progress = committed + .filter((entry) => entry.topic === KAFKA_COHORT_MEMBERSHIP_CHANGED && entry.offset >= 0) + .map((entry) => ({ partition: entry.partition, nextOffset: entry.offset })) + + if (progress.length === 0) { + return + } + + await this.upsertConsumerProgress(PostgresUse.BEHAVIORAL_COHORTS_RW, progress) + } + + /** + * Read off the membership consumer's own client rather than the marker one: the two topics + * share brokers today, but the processor keeps a separate marker producer so they could split. + */ + private async captureMembershipWatermarks(): Promise { + // Asked of the broker rather than assumed. The processor's partition count governs the + // marker set, not this topic's layout; hard-coding it here would silently stop covering + // partitions above the assumed count, and the gate cannot wait on what it never captured. + const metadata = await this.kafkaConsumer.getPartitionsForTopic(KAFKA_COHORT_MEMBERSHIP_CHANGED) + + if (metadata.length === 0) { + throw new Error(`No partition metadata for ${KAFKA_COHORT_MEMBERSHIP_CHANGED}`) + } + + const results = await Promise.all( + metadata.map(async ({ id: partition }) => { + const [, high] = await this.kafkaConsumer.queryWatermarkOffsets( + KAFKA_COHORT_MEMBERSHIP_CHANGED, + partition + ) + return { partition, high } + }) + ) + + const watermarks: MembershipWatermarks = {} + for (const { partition, high } of results) { + watermarks[partition] = high + } + + return watermarks + } + + /** + * Membership rows, the snapshot marks they carry, and the consumer's progress all land in one + * transaction: progress must never claim offsets whose rows rolled back, or a sweep could + * delete rows it believes were re-asserted. + */ + private async persistCohortMembershipChanges( + changes: CohortMembershipChange[], + progress: PartitionProgress[] = [] + ): Promise { + if (changes.length === 0 && progress.length === 0) { return } @@ -49,35 +217,22 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { deduped.set(`${change.team_id}:${change.cohort_id}:${change.person_id}`, change) } - // Build the VALUES clause for batch upsert - const values: any[] = [] - const placeholders: string[] = [] - let paramIndex = 1 + const marks = collectSnapshotMarks(changes) - for (const change of deduped.values()) { - const inCohort = change.status === 'entered' - placeholders.push( - `($${paramIndex}, $${paramIndex + 1}, $${paramIndex + 2}, $${paramIndex + 3}, CURRENT_TIMESTAMP)` - ) - values.push(change.team_id, change.cohort_id, change.person_id, inCohort) - paramIndex += 4 - } - - // Single batch UPSERT query - handles both entered and left events - const query = ` - INSERT INTO cohort_membership (team_id, cohort_id, person_id, in_cohort, last_updated) - VALUES ${placeholders.join(', ')} - ON CONFLICT (team_id, cohort_id, person_id) - DO UPDATE SET - in_cohort = EXCLUDED.in_cohort, - last_updated = CURRENT_TIMESTAMP - ` - - await this.deps.postgres.query( + await this.deps.postgres.transaction( PostgresUse.BEHAVIORAL_COHORTS_RW, - query, - values, - 'batchUpsertCohortMembership' + 'cohortMembershipBatch', + async (tx) => { + if (deduped.size > 0) { + await this.upsertMembershipRows(tx, Array.from(deduped.values())) + } + if (marks.length > 0) { + await this.upsertSnapshotMarks(tx, marks) + } + if (progress.length > 0) { + await this.upsertConsumerProgress(tx, progress) + } + } ) } catch (error) { logger.error('Failed to process batch cohort membership changes', { @@ -88,6 +243,125 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { } } + /** + * The version guard is what makes an at-least-once replay safe: offsets commit after the write, + * so a crash replays a batch, and a stale reconcile row must not overwrite a newer live change + * that was already applied. It is the same last-writer-wins rule ClickHouse applies to this + * data. A row with no version counts as the oldest. + */ + private async upsertMembershipRows(tx: TransactionClient, changes: CohortMembershipChange[]): Promise { + const values: any[] = [] + const placeholders: string[] = [] + let paramIndex = 1 + + for (const change of changes) { + placeholders.push( + `($${paramIndex}, $${paramIndex + 1}, $${paramIndex + 2}, $${paramIndex + 3}, CURRENT_TIMESTAMP, $${ + paramIndex + 4 + }::timestamp)` + ) + values.push( + change.team_id, + change.cohort_id, + change.person_id, + change.status === 'entered', + change.last_updated ?? null + ) + paramIndex += 5 + } + + await this.deps.postgres.query( + tx, + ` + INSERT INTO cohort_membership (team_id, cohort_id, person_id, in_cohort, last_updated, version) + VALUES ${placeholders.join(', ')} + ON CONFLICT (team_id, cohort_id, person_id) + DO UPDATE SET + in_cohort = EXCLUDED.in_cohort, + last_updated = CURRENT_TIMESTAMP, + version = EXCLUDED.version + WHERE cohort_membership.version IS NULL OR EXCLUDED.version >= cohort_membership.version + `, + values, + 'batchUpsertCohortMembership' + ) + } + + /** + * Folds this batch's contribution into each run's ledger row. The row is created here when the + * snapshot arrives before the run's completion markers, which is the normal ordering. + */ + private async upsertSnapshotMarks(tx: TransactionClient, marks: SnapshotMark[]): Promise { + const values: any[] = [] + const placeholders: string[] = [] + let paramIndex = 1 + + for (const mark of marks) { + placeholders.push( + `($${paramIndex}, $${paramIndex + 1}, $${paramIndex + 2}, $${paramIndex + 3}::timestamp, $${ + paramIndex + 4 + })` + ) + values.push(mark.runId, mark.cohortId, mark.teamId, mark.minVersion, mark.rows) + paramIndex += 5 + } + + await this.deps.postgres.query( + tx, + ` + INSERT INTO cohort_membership_sweeps + (run_id, cohort_id, team_id, min_snapshot_version, snapshot_rows) + VALUES ${placeholders.join(', ')} + ON CONFLICT (run_id, cohort_id) + DO UPDATE SET + min_snapshot_version = LEAST( + cohort_membership_sweeps.min_snapshot_version, + EXCLUDED.min_snapshot_version + ), + snapshot_rows = cohort_membership_sweeps.snapshot_rows + EXCLUDED.snapshot_rows, + updated_at = CURRENT_TIMESTAMP + `, + values, + 'markCohortMembershipSnapshot' + ) + } + + /** + * GREATEST keeps the gate monotone, so a rebalance that replays from an older committed offset + * cannot walk a partition's progress backwards and unblock nothing that was already blocked. + */ + private async upsertConsumerProgress( + target: PostgresUse | TransactionClient, + progress: PartitionProgress[] + ): Promise { + const values: any[] = [] + const placeholders: string[] = [] + let paramIndex = 1 + + for (const entry of progress) { + placeholders.push(`($${paramIndex}, $${paramIndex + 1})`) + values.push(entry.partition, entry.nextOffset) + paramIndex += 2 + } + + await this.deps.postgres.query( + target, + ` + INSERT INTO cohort_membership_consumer_progress (partition, next_offset) + VALUES ${placeholders.join(', ')} + ON CONFLICT (partition) + DO UPDATE SET + next_offset = GREATEST( + cohort_membership_consumer_progress.next_offset, + EXCLUDED.next_offset + ), + updated_at = CURRENT_TIMESTAMP + `, + values, + 'upsertCohortMembershipProgress' + ) + } + private _parseAndValidateBatch(messages: Message[]): CohortMembershipChange[] { const cohortMembershipChanges: CohortMembershipChange[] = [] @@ -126,6 +400,11 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { return cohortMembershipChanges } + private async handleBatch(messages: Message[]): Promise { + const cohortMembershipChanges = this._parseAndValidateBatch(messages) + await this.persistCohortMembershipChanges(cohortMembershipChanges, partitionProgressFromMessages(messages)) + } + public override async start(): Promise { await super.start() @@ -137,15 +416,34 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { }) return instrumentFn('cdpCohortMembershipConsumer.handleEachBatch', async () => { - const cohortMembershipChanges = this._parseAndValidateBatch(messages) - await this.persistCohortMembershipChanges(cohortMembershipChanges) + await this.handleBatch(messages) }) }) + + // Started after the membership consumer so its client is connected for watermark reads. + if (this.markerKafkaConsumer && this.sweeper) { + await this.markerKafkaConsumer.connect(async (messages) => { + return instrumentFn('cdpCohortMembershipConsumer.handleMarkerBatch', async () => { + await this.handleMarkerBatch(messages) + }) + }) + this.sweeper.start() + } + } + + private async handleMarkerBatch(messages: Message[]): Promise { + const sweeper = this.sweeper + if (!sweeper) { + return + } + + await sweeper.applyMarkers(sweeper.parseMarkers(messages)) } public override async stop(): Promise { logger.info('💤', `Stopping ${this.name}...`) - await this.kafkaConsumer.disconnect() + await this.sweeper?.stop() + await Promise.all([this.kafkaConsumer.disconnect(), this.markerKafkaConsumer?.disconnect()]) // IMPORTANT: super always comes last await super.stop() @@ -153,6 +451,10 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { } public isHealthy(): HealthCheckResult { - return this.kafkaConsumer.isHealthy() + const results = [this.kafkaConsumer.isHealthy(), this.markerKafkaConsumer?.isHealthy()].filter( + (result): result is HealthCheckResult => result !== undefined + ) + + return results.find((result) => result.status !== 'ok') ?? results[0] } } diff --git a/nodejs/src/cdp/dev/cohort-sweep-e2e.test.ts b/nodejs/src/cdp/dev/cohort-sweep-e2e.test.ts new file mode 100644 index 000000000000..16f0ca36f926 --- /dev/null +++ b/nodejs/src/cdp/dev/cohort-sweep-e2e.test.ts @@ -0,0 +1,131 @@ +import { KAFKA_COHORT_MEMBERSHIP_CHANGED, KAFKA_COHORT_RECONCILE_MARKERS } from '~/common/config/kafka-topics' +import { KafkaProducerWrapper } from '~/common/kafka/producer' +import { closeHub, createHub } from '~/common/utils/db/hub' +import { PostgresUse } from '~/common/utils/db/postgres' +import { UUIDT } from '~/common/utils/utils' +import { createCdpConsumerDeps } from '~/tests/helpers/cdp' +import { waitForExpect } from '~/tests/helpers/expectations' +import { resetBehavioralCohortsDatabase } from '~/tests/helpers/sql' +import { Hub } from '~/types' + +import { CdpCohortMembershipConsumer } from '../consumers/cdp-cohort-membership.consumer' +import { COHORT_PARTITION_COUNT } from '../services/cohort-membership/sweeper.service' + +// Drives the whole chain against the running dev stack: real Kafka topics, the real membership +// consumer, the real marker consumer, and the real sweep loop. Lives under dev/ so CI skips it. +jest.setTimeout(180_000) + +describe('cohort membership mark and sweep, end to end', () => { + const TEAM_ID = 987654 + const COHORT_ID = 5150 + const SNAPSHOT_VERSION = '2026-08-26 10:00:00.000000' + const STALE_VERSION = '2026-08-26 09:00:00.000000' + + let hub: Hub + let consumer: CdpCohortMembershipConsumer + let producer: KafkaProducerWrapper + + afterEach(async () => { + await consumer?.stop() + await producer?.disconnect() + await closeHub(hub) + }) + + it('deletes rows the completed reconcile run did not re-assert', async () => { + const runId = new UUIDT().toString() + const stillMember = new UUIDT().toString() + const droppedOut = new UUIDT().toString() + + hub = await createHub() + await resetBehavioralCohortsDatabase(hub.postgres) + + // Someone who matched before the cohort was edited. No live transition ever arrives for + // them, so without a sweep this row claims membership forever. + await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + `INSERT INTO cohort_membership (team_id, cohort_id, person_id, in_cohort, version) + VALUES ($1, $2, $3, true, $4::timestamp)`, + [TEAM_ID, COHORT_ID, droppedOut, STALE_VERSION], + 'seedStaleRow' + ) + + consumer = new CdpCohortMembershipConsumer( + { ...hub, COHORT_MEMBERSHIP_SWEEP_ENABLED: true, COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS: 1000 }, + createCdpConsumerDeps(hub) + ) + await consumer.start() + + producer = await KafkaProducerWrapper.create(hub.KAFKA_CLIENT_RACK) + + await producer.produce({ + topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, + key: Buffer.from(stillMember), + value: Buffer.from( + JSON.stringify({ + team_id: TEAM_ID, + cohort_id: COHORT_ID, + person_id: stillMember, + last_updated: SNAPSHOT_VERSION, + status: 'entered', + origin: 'reconcile', + run_id: runId, + }) + ), + }) + await producer.flush() + + // The processor produces a marker only after its partition's membership rows are acked. + await waitForExpect(async () => { + const applied = await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + 'SELECT person_id FROM cohort_membership WHERE person_id = $1', + [stillMember], + 'awaitSnapshot' + ) + expect(applied.rows).toHaveLength(1) + }, 60_000) + + for (let partition = 0; partition < COHORT_PARTITION_COUNT; partition++) { + await producer.produce({ + topic: KAFKA_COHORT_RECONCILE_MARKERS, + key: Buffer.from(`${TEAM_ID}:${COHORT_ID}:${runId}:${partition}`), + value: Buffer.from( + JSON.stringify({ + type: 'reconcile_complete', + team_id: TEAM_ID, + cohort_id: COHORT_ID, + partition, + run_id: runId, + last_updated: SNAPSHOT_VERSION, + }) + ), + }) + } + await producer.flush() + + // Both checks sit in the same wait: the rows go before the ledger is marked done, so + // asserting the ledger separately would race the sweep's own last statement. + await waitForExpect(async () => { + const members = await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + 'SELECT person_id FROM cohort_membership WHERE team_id = $1 AND cohort_id = $2', + [TEAM_ID, COHORT_ID], + 'readMembers' + ) + expect(members.rows.map((row: { person_id: string }) => row.person_id)).toEqual([stillMember]) + + const sweeps = await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + 'SELECT marker_bits, status, swept_rows, snapshot_rows FROM cohort_membership_sweeps WHERE run_id = $1', + [runId], + 'readSweep' + ) + expect(sweeps.rows[0]).toMatchObject({ + marker_bits: '-1', + status: 'swept', + swept_rows: '1', + snapshot_rows: '1', + }) + }, 120_000) + }) +}) diff --git a/nodejs/src/cdp/services/cohort-membership/sweeper.service.test.ts b/nodejs/src/cdp/services/cohort-membership/sweeper.service.test.ts new file mode 100644 index 000000000000..3d113c8354b2 --- /dev/null +++ b/nodejs/src/cdp/services/cohort-membership/sweeper.service.test.ts @@ -0,0 +1,252 @@ +import { Message } from 'node-rdkafka' + +import { defaultConfig } from '~/common/config/config' +import { KAFKA_COHORT_RECONCILE_MARKERS } from '~/common/config/kafka-topics' +import { PostgresRouter, PostgresUse } from '~/common/utils/db/postgres' +import { UUIDT } from '~/common/utils/utils' + +import { resetBehavioralCohortsDatabase } from '../../../../tests/helpers/sql' +import { + COHORT_PARTITION_COUNT, + CohortMembershipSweepConfig, + CohortMembershipSweeper, + MembershipWatermarks, +} from './sweeper.service' + +jest.setTimeout(30_000) + +describe('CohortMembershipSweeper', () => { + const TEAM_ID = 4242 + const COHORT_ID = 91204 + const MARKER_VERSION = '2026-05-26 12:30:00.000000' + const THRESHOLD = '2026-05-26 12:00:00.000000' + const BELOW_THRESHOLD = '2026-05-26 11:00:00.000000' + const ABOVE_THRESHOLD = '2026-05-26 13:00:00.000000' + + const config: CohortMembershipSweepConfig = { + COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS: 60000, + // Small enough that the batched delete loop has to iterate. + COHORT_MEMBERSHIP_SWEEP_BATCH_SIZE: 2, + COHORT_MEMBERSHIP_SWEEP_CLAIM_TIMEOUT_MS: 300000, + COHORT_MEMBERSHIP_SWEEP_ABANDON_AFTER_DAYS: 3, + } + + let postgres: PostgresRouter + let sweeper: CohortMembershipSweeper + let watermarks: MembershipWatermarks + let runId: string + + beforeEach(async () => { + postgres = new PostgresRouter(defaultConfig) + await resetBehavioralCohortsDatabase(postgres) + watermarks = { 0: 10 } + runId = new UUIDT().toString() + sweeper = new CohortMembershipSweeper(config, postgres, { + captureMembershipWatermarks: () => Promise.resolve(watermarks), + refreshConsumerProgress: () => Promise.resolve(), + }) + }) + + afterEach(async () => { + await postgres.end() + }) + + // The wire contract the Rust processor pins with golden bytes in + // `rust/cohort-core/src/seed/reconcile.rs`. + const markerMessage = (partition: number, overrides: Record = {}): Message => ({ + value: Buffer.from( + JSON.stringify({ + type: 'reconcile_complete', + team_id: TEAM_ID, + cohort_id: COHORT_ID, + partition, + run_id: runId, + last_updated: MARKER_VERSION, + ...overrides, + }) + ), + topic: KAFKA_COHORT_RECONCILE_MARKERS, + partition: 0, + offset: partition, + timestamp: 0, + size: 0, + key: null, + }) + + const query = async (sql: string, values?: any[]): Promise => { + const result = await postgres.query(PostgresUse.BEHAVIORAL_COHORTS_RW, sql, values, 'test') + return result.rows as T[] + } + + const readSweep = async (): Promise> => { + const rows = await query>( + `SELECT marker_bits, status, swept_rows, membership_hwms, + to_char(min_marker_version, 'YYYY-MM-DD HH24:MI:SS.US') AS min_marker_version + FROM cohort_membership_sweeps WHERE run_id = $1 AND cohort_id = $2`, + [runId, COHORT_ID] + ) + return rows[0] + } + + const insertMembership = async (teamId: number, cohortId: number, version: string | null): Promise => { + const personId = new UUIDT().toString() + await query( + `INSERT INTO cohort_membership (team_id, cohort_id, person_id, in_cohort, version) + VALUES ($1, $2, $3, true, $4::timestamp)`, + [teamId, cohortId, personId, version] + ) + return personId + } + + const survivingPersonIds = async (): Promise => { + const rows = await query<{ person_id: string }>('SELECT person_id FROM cohort_membership ORDER BY person_id') + return rows.map((row) => row.person_id) + } + + /** A run whose markers are all in, with a snapshot minimum already folded in by the consumer. */ + const readyRun = async (): Promise => { + await query( + `INSERT INTO cohort_membership_sweeps + (run_id, cohort_id, team_id, marker_bits, min_marker_version, min_snapshot_version, + membership_hwms, status) + VALUES ($1, $2, $3, -1, $4::timestamp, $5::timestamp, $6::jsonb, 'ready')`, + [runId, COHORT_ID, TEAM_ID, MARKER_VERSION, THRESHOLD, JSON.stringify(watermarks)] + ) + } + + /** Record the consumer as having applied everything the run's watermarks cover. */ + const openGate = async (): Promise => { + for (const [partition, highWatermark] of Object.entries(watermarks)) { + await query(`INSERT INTO cohort_membership_consumer_progress (partition, next_offset) VALUES ($1, $2)`, [ + Number(partition), + highWatermark, + ]) + } + } + + it('should skip unusable markers instead of failing the batch', () => { + const messages: Message[] = [ + markerMessage(0), + { ...markerMessage(1), value: Buffer.from('not json') }, + markerMessage(2, { type: 'some_future_control_message' }), + { ...markerMessage(3), value: null }, + markerMessage(4, { partition: COHORT_PARTITION_COUNT }), + markerMessage(5, { run_id: 'not-a-uuid' }), + ] + + const markers = sweeper.parseMarkers(messages) + + expect(markers.map((marker) => marker.partition)).toEqual([0]) + }) + + it('should union marker bits and turn the run claimable only once the set is complete', async () => { + const partitions = Array.from({ length: COHORT_PARTITION_COUNT }, (_, partition) => partition) + + await sweeper.applyMarkers(sweeper.parseMarkers(partitions.slice(0, -1).map((p) => markerMessage(p)))) + + expect(await readSweep()).toMatchObject({ status: 'collecting', membership_hwms: null }) + + // Redelivery must be a no-op, not an increment. + await sweeper.applyMarkers(sweeper.parseMarkers([markerMessage(5), markerMessage(5)])) + + expect(await readSweep()).toMatchObject({ status: 'collecting' }) + + watermarks = { 0: 10, 3: 4 } + // The last partition sets the sign bit, which is what completes the i64 bitmap. + await sweeper.applyMarkers( + sweeper.parseMarkers([markerMessage(COHORT_PARTITION_COUNT - 1, { last_updated: ABOVE_THRESHOLD })]) + ) + + expect(await readSweep()).toMatchObject({ + marker_bits: '-1', + status: 'ready', + membership_hwms: { 0: 10, 3: 4 }, + min_marker_version: MARKER_VERSION, + }) + }) + + it('should hold the sweep until consumer progress passes every captured watermark', async () => { + watermarks = { 0: 10, 3: 4 } + await readyRun() + const stale = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + + expect(await sweeper.runOnce()).toMatchObject({ swept: 0, blocked: 1 }) + expect(await survivingPersonIds()).toEqual([stale]) + + await query(`INSERT INTO cohort_membership_consumer_progress (partition, next_offset) VALUES (0, 9), (3, 4)`) + + expect(await sweeper.runOnce()).toMatchObject({ swept: 0, blocked: 1 }) + expect(await survivingPersonIds()).toEqual([stale]) + + await query(`UPDATE cohort_membership_consumer_progress SET next_offset = 10 WHERE partition = 0`) + + expect(await sweeper.runOnce()).toMatchObject({ swept: 1, rowsDeleted: 1, blocked: 0 }) + expect(await survivingPersonIds()).toEqual([]) + }) + + it('should delete only this cohort rows older than the run, and stop after one sweep', async () => { + await readyRun() + await openGate() + + const stale1 = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + const stale2 = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + const neverVersioned = await insertMembership(TEAM_ID, COHORT_ID, null) + const asserted = await insertMembership(TEAM_ID, COHORT_ID, THRESHOLD) + const liveChange = await insertMembership(TEAM_ID, COHORT_ID, ABOVE_THRESHOLD) + const otherCohort = await insertMembership(TEAM_ID, COHORT_ID + 1, BELOW_THRESHOLD) + const otherTeam = await insertMembership(TEAM_ID + 1, COHORT_ID, BELOW_THRESHOLD) + + expect(await sweeper.runOnce()).toMatchObject({ swept: 1, rowsDeleted: 3 }) + + const surviving = await survivingPersonIds() + expect(surviving).toEqual([asserted, liveChange, otherCohort, otherTeam].sort((a, b) => a.localeCompare(b))) + expect(surviving).not.toEqual(expect.arrayContaining([stale1, stale2, neverVersioned])) + expect(await readSweep()).toMatchObject({ status: 'swept', swept_rows: '3' }) + + // A second cycle must not re-sweep a finished run. + expect(await sweeper.runOnce()).toMatchObject({ swept: 0, rowsDeleted: 0 }) + }) + + it.each([ + ['it never observed the snapshot it is supposed to have asserted', 'min_snapshot_version = NULL'], + ['its watermark capture came back empty', `membership_hwms = '{}'::jsonb`], + ])('should refuse to sweep a run when %s', async (_label, damage) => { + watermarks = { 0: 10 } + await readyRun() + await query(`INSERT INTO cohort_membership_consumer_progress (partition, next_offset) VALUES (0, 10)`) + await query(`UPDATE cohort_membership_sweeps SET ${damage} WHERE run_id = $1`, [runId]) + + const live = await insertMembership(TEAM_ID, COHORT_ID, ABOVE_THRESHOLD) + const unversioned = await insertMembership(TEAM_ID, COHORT_ID, null) + + expect(await sweeper.runOnce()).toMatchObject({ swept: 0, rowsDeleted: 0 }) + expect(await survivingPersonIds()).toEqual([live, unversioned].sort((a, b) => a.localeCompare(b))) + }) + + it('should reclaim a sweep whose pod died mid-run', async () => { + await readyRun() + await openGate() + const stale = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + await query( + `UPDATE cohort_membership_sweeps + SET status = 'sweeping', claimed_at = CURRENT_TIMESTAMP - INTERVAL '1 hour' + WHERE run_id = $1`, + [runId] + ) + + expect(await sweeper.runOnce()).toMatchObject({ swept: 1, rowsDeleted: 1 }) + expect(await survivingPersonIds()).not.toContain(stale) + }) + + it('should abandon a run whose marker set never completed', async () => { + await query( + `INSERT INTO cohort_membership_sweeps (run_id, cohort_id, team_id, marker_bits, created_at, updated_at) + VALUES ($1, $2, $3, 7, CURRENT_TIMESTAMP - INTERVAL '30 days', + CURRENT_TIMESTAMP - INTERVAL '30 days')`, + [runId, COHORT_ID, TEAM_ID] + ) + + expect(await sweeper.runOnce()).toMatchObject({ abandoned: 1 }) + expect(await readSweep()).toMatchObject({ status: 'abandoned' }) + }) +}) diff --git a/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts b/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts new file mode 100644 index 000000000000..2a1b4f4ad285 --- /dev/null +++ b/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts @@ -0,0 +1,545 @@ +import { Message } from 'node-rdkafka' +import { Counter, Gauge } from 'prom-client' +import { z } from 'zod' + +import { PostgresRouter, PostgresUse } from '~/common/utils/db/postgres' +import { parseJSON } from '~/common/utils/json-parse' +import { logger } from '~/common/utils/logger' + +/** + * One completion marker per processor partition certifies a reconcile run. Kept in step with + * `COHORT_PARTITION_COUNT` in `rust/cohort-core/src/partitioner.rs`, which is compile-time stable + * by design — the marker topic's partition count is load-bearing for the same reason. + */ +export const COHORT_PARTITION_COUNT = 64 + +/** All 64 bits set, read as a signed 64-bit integer. Same convention as the seeder's ledger. */ +const ALL_MARKERS = '-1' + +/** A tick sweeps at most this many runs; the rest wait for the next one. */ +const MAX_RUNS_PER_TICK = 10 + +/** + * Candidates read per tick. Larger than the sweep cap because most candidates are usually waiting + * on the gate, and a run cannot be told apart from a blocked one until it is read — a page the size + * of the sweep cap would let a handful of blocked runs starve everything behind them. + */ +const MAX_CANDIDATES_PER_TICK = 200 + +const ReconcileCompleteMarkerSchema = z.object({ + type: z.literal('reconcile_complete'), + team_id: z.number(), + cohort_id: z.number(), + partition: z + .number() + .int() + .min(0) + .max(COHORT_PARTITION_COUNT - 1), + run_id: z.guid(), + // The producer's fixed-width stamp. Pinned here because it is bound straight into a + // `::timestamp` cast: anything Postgres cannot parse would throw out of the marker batch and + // take the membership consumer down with it. + last_updated: z.string().regex(/^\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{6}$/), +}) + +export type ReconcileCompleteMarker = z.infer + +/** Membership-topic high watermarks at marker completeness, keyed by partition. */ +export type MembershipWatermarks = Record + +export type CohortMembershipSweepConfig = { + COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS: number + COHORT_MEMBERSHIP_SWEEP_BATCH_SIZE: number + COHORT_MEMBERSHIP_SWEEP_CLAIM_TIMEOUT_MS: number + COHORT_MEMBERSHIP_SWEEP_ABANDON_AFTER_DAYS: number +} + +export type SweepCycleResult = { + swept: number + rowsDeleted: number + blocked: number + abandoned: number +} + +type ClaimableRun = { + run_id: string + cohort_id: string + team_id: string + membership_hwms: MembershipWatermarks | null + threshold: string | null +} + +const markersConsumed = new Counter({ + name: 'cdp_cohort_membership_markers_consumed', + help: 'Reconcile completion markers read off the marker topic, by parse outcome', + labelNames: ['outcome'], +}) + +const sweepsExecuted = new Counter({ + name: 'cdp_cohort_membership_sweeps_executed', + help: 'Reconcile runs whose stale rows were swept', +}) + +const rowsSwept = new Counter({ + name: 'cdp_cohort_membership_rows_swept', + help: 'Stale cohort_membership rows deleted by a sweep', +}) + +const sweepsGateBlocked = new Counter({ + name: 'cdp_cohort_membership_sweeps_gate_blocked', + help: 'Sweep attempts held back because the membership consumer had not passed the snapshot', +}) + +const sweepsAbandoned = new Counter({ + name: 'cdp_cohort_membership_sweeps_abandoned', + help: 'Reconcile runs given up on because their marker set never completed', +}) + +const sweepCycles = new Counter({ + name: 'cdp_cohort_membership_sweep_cycles', + help: 'Sweep cycles run, by outcome. A cycle that only ever errors sweeps nothing, silently', + labelNames: ['status'], +}) + +const sweepGateWaitSeconds = new Gauge({ + name: 'cdp_cohort_membership_sweep_gate_wait_seconds', + help: 'Time since the oldest fully-marked run was last written to, while it waits on the gate', +}) + +/** + * Deletes `cohort_membership` rows that a completed reconcile run did not re-assert. + * + * The consumer only ever upserts, so a person who stops matching an edited cohort keeps an + * `in_cohort = true` row forever. A reconcile run replays the cohort's full current membership and + * then certifies itself with one marker per partition; everything the run stamped older than + * itself is by definition no longer a member. + */ +export class CohortMembershipSweeper { + private running = false + private loopPromise: Promise | null = null + private sleepResolve: (() => void) | null = null + + constructor( + private config: CohortMembershipSweepConfig, + private postgres: PostgresRouter, + /** Both read the membership consumer's own client: the two topics may live on different clusters. */ + private kafka: { + captureMembershipWatermarks: () => Promise + refreshConsumerProgress: () => Promise + } + ) {} + + /** + * One bad message must not wedge the sweep pipeline, so parsing skips and counts rather than + * failing the batch the way the membership consumer does. + */ + public parseMarkers(messages: Message[]): ReconcileCompleteMarker[] { + const markers: ReconcileCompleteMarker[] = [] + + for (const message of messages) { + try { + const value = message.value?.toString() + if (!value) { + markersConsumed.inc({ outcome: 'skipped' }) + continue + } + + const parsed = ReconcileCompleteMarkerSchema.safeParse(parseJSON(value)) + if (!parsed.success) { + markersConsumed.inc({ outcome: 'skipped' }) + logger.warn('Skipping unusable reconcile marker', { + errors: parsed.error.issues, + }) + continue + } + + markersConsumed.inc({ outcome: 'parsed' }) + markers.push(parsed.data) + } catch (error) { + markersConsumed.inc({ outcome: 'skipped' }) + logger.warn('Skipping unparseable reconcile marker', { error }) + } + } + + return markers + } + + /** + * Folding a marker in is a bit union, so redelivery changes nothing. The run turns claimable + * only on the transition to a complete set, and only once. + */ + public async applyMarkers(markers: ReconcileCompleteMarker[]): Promise { + for (const marker of markers) { + const result = await this.postgres.query<{ complete: boolean; status: string }>( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + INSERT INTO cohort_membership_sweeps + (run_id, cohort_id, team_id, marker_bits, min_marker_version) + VALUES ($1, $2, $3, (1::bigint << $4), $5::timestamp) + ON CONFLICT (run_id, cohort_id) + DO UPDATE SET + marker_bits = cohort_membership_sweeps.marker_bits | (1::bigint << $4), + min_marker_version = LEAST( + cohort_membership_sweeps.min_marker_version, + EXCLUDED.min_marker_version + ), + updated_at = CURRENT_TIMESTAMP + RETURNING marker_bits = ${ALL_MARKERS} AS complete, status + `, + [marker.run_id, marker.cohort_id, marker.team_id, marker.partition, marker.last_updated], + 'applyCohortReconcileMarker' + ) + + const { complete, status } = result.rows[0] + if (!complete || status !== 'collecting') { + continue + } + + // Every snapshot row of this run was acked before its marker was produced, so the + // watermarks read now bound the whole snapshot on the membership topic. + const watermarks = await this.kafka.captureMembershipWatermarks() + + await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + UPDATE cohort_membership_sweeps + SET status = 'ready', membership_hwms = $1::jsonb, updated_at = CURRENT_TIMESTAMP + WHERE run_id = $2 AND cohort_id = $3 + AND status = 'collecting' AND marker_bits = ${ALL_MARKERS} + `, + [JSON.stringify(watermarks), marker.run_id, marker.cohort_id], + 'readyCohortMembershipSweep' + ) + + logger.info('Reconcile run is fully marked and ready to sweep', { + run_id: marker.run_id, + cohort_id: marker.cohort_id, + team_id: marker.team_id, + }) + } + } + + public async runOnce(): Promise { + const result: SweepCycleResult = { swept: 0, rowsDeleted: 0, blocked: 0, abandoned: 0 } + + await this.kafka.refreshConsumerProgress() + const progress = await this.readConsumerProgress() + + for (const run of await this.claimableRuns()) { + if (!this.gateSatisfied(run, progress)) { + result.blocked += 1 + sweepsGateBlocked.inc() + continue + } + + if (result.swept >= MAX_RUNS_PER_TICK) { + break + } + + const claimed = await this.claim(run) + if (!claimed) { + continue + } + + result.rowsDeleted += await this.sweep(claimed) + result.swept += 1 + sweepsExecuted.inc() + } + + result.abandoned = await this.abandonStuckRuns() + await this.collectGarbage() + await this.recordGateWait() + + sweepCycles.inc({ status: 'ok' }) + return result + } + + /** + * Runs that hold every marker, plus runs whose claiming pod died. Re-sweeping is harmless: the + * delete only ever removes rows below a threshold that does not move. + * + * A run with no snapshot minimum is never claimable. Marker stamps alone cannot bound the + * snapshot: a fast partition finishes and stamps its marker while a slower one is still + * emitting rows, so the earliest marker can sit above rows the run asserted. Sweeping there + * would delete live members — and every unversioned row with them. Such a run waits for + * abandonment instead; the next reconcile heals the cohort. + */ + private async claimableRuns(): Promise { + const result = await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + SELECT run_id, cohort_id, team_id, membership_hwms, + LEAST(min_snapshot_version, min_marker_version) AS threshold + FROM cohort_membership_sweeps + WHERE min_snapshot_version IS NOT NULL + AND (status = 'ready' + OR (status = 'sweeping' AND claimed_at < CURRENT_TIMESTAMP - $1::interval)) + ORDER BY created_at + LIMIT ${MAX_CANDIDATES_PER_TICK} + `, + [`${this.config.COHORT_MEMBERSHIP_SWEEP_CLAIM_TIMEOUT_MS} milliseconds`], + 'findClaimableCohortMembershipSweeps' + ) + + return result.rows + } + + private async readConsumerProgress(): Promise> { + const result = await this.postgres.query<{ partition: number; next_offset: string }>( + PostgresUse.BEHAVIORAL_COHORTS_RW, + 'SELECT partition, next_offset FROM cohort_membership_consumer_progress', + undefined, + 'readCohortMembershipProgress' + ) + + return new Map(result.rows.map((row) => [row.partition, Number(row.next_offset)])) + } + + /** + * Markers are a tiny topic; the membership topic carries the multi-million-row snapshot they + * certify. Sweeping on markers alone would delete rows whose re-asserting messages are still + * queued, blanking a cohort for flag evaluation until the consumer catches up. + */ + private gateSatisfied(run: ClaimableRun, progress: Map): boolean { + const watermarks = Object.entries(run.membership_hwms ?? {}) + + // No watermarks means the capture never happened or came back empty. An empty map would + // otherwise satisfy every partition vacuously, which is the one way this gate can pass + // without proving anything. + if (watermarks.length === 0) { + logger.warn('Refusing to sweep a run with no captured watermarks', { + run_id: run.run_id, + cohort_id: run.cohort_id, + }) + return false + } + + const blocked: number[] = [] + + for (const [partition, highWatermark] of watermarks) { + // Nothing was ever produced to this partition, so there is nothing to wait for. + if (highWatermark <= 0) { + continue + } + + const applied = progress.get(Number(partition)) + if (applied === undefined || applied < highWatermark) { + blocked.push(Number(partition)) + } + } + + if (blocked.length > 0) { + logger.info('Holding a cohort membership sweep until the consumer passes the snapshot', { + run_id: run.run_id, + cohort_id: run.cohort_id, + blocked_partitions: blocked, + }) + return false + } + + return true + } + + private async claim(run: ClaimableRun): Promise { + const result = await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + UPDATE cohort_membership_sweeps + SET status = 'sweeping', claimed_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP + WHERE run_id = $1 AND cohort_id = $2 + AND min_snapshot_version IS NOT NULL + AND (status = 'ready' + OR (status = 'sweeping' AND claimed_at < CURRENT_TIMESTAMP - $3::interval)) + RETURNING run_id, cohort_id, team_id, membership_hwms, + LEAST(min_snapshot_version, min_marker_version) AS threshold + `, + [run.run_id, run.cohort_id, `${this.config.COHORT_MEMBERSHIP_SWEEP_CLAIM_TIMEOUT_MS} milliseconds`], + 'claimCohortMembershipSweep' + ) + + return result.rows[0] ?? null + } + + private async sweep(run: ClaimableRun): Promise { + if (!run.threshold) { + logger.error('Refusing to sweep a run with no threshold', { + run_id: run.run_id, + cohort_id: run.cohort_id, + }) + return 0 + } + + let deleted = 0 + let lastBatch = 0 + + do { + // The predicate is repeated in the DELETE on purpose: under READ COMMITTED the DELETE + // re-checks it against the newest row version, so a row a concurrent upsert just + // refreshed past the threshold is left alone instead of deleted anyway. + const result = await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + WITH candidates AS ( + SELECT id FROM cohort_membership + WHERE team_id = $1 AND cohort_id = $2 + AND (version IS NULL OR version < $3::timestamp) + LIMIT $4 + FOR UPDATE SKIP LOCKED + ) + DELETE FROM cohort_membership m + USING candidates c + WHERE m.id = c.id AND (m.version IS NULL OR m.version < $3::timestamp) + `, + [run.team_id, run.cohort_id, run.threshold, this.config.COHORT_MEMBERSHIP_SWEEP_BATCH_SIZE], + 'sweepCohortMembership' + ) + + lastBatch = result.rowCount ?? 0 + deleted += lastBatch + + // Without this the claim goes stale mid-run and every other pod reclaims the same + // cohort, so N pods page over the same rows. + if (lastBatch > 0) { + await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + `UPDATE cohort_membership_sweeps SET claimed_at = CURRENT_TIMESTAMP + WHERE run_id = $1 AND cohort_id = $2 AND status = 'sweeping'`, + [run.run_id, run.cohort_id], + 'heartbeatCohortMembershipSweep' + ) + } + } while (lastBatch > 0) + + // A zero-delete iteration means either nothing is left or SKIP LOCKED skipped everything. + // Only the first is finished; calling the second `swept` would retire the run with stale + // rows still in place, and nothing revisits a swept run. + const remaining = await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + `SELECT 1 FROM cohort_membership + WHERE team_id = $1 AND cohort_id = $2 AND (version IS NULL OR version < $3::timestamp) + LIMIT 1`, + [run.team_id, run.cohort_id, run.threshold], + 'checkCohortMembershipSweepComplete' + ) + const finished = remaining.rows.length === 0 + + // Guarded on 'sweeping' so a straggler cannot overwrite the pod that reclaimed the run, nor + // resurrect one that was abandoned in between. + await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + UPDATE cohort_membership_sweeps + SET status = $1, swept_rows = $2, claimed_at = NULL, updated_at = CURRENT_TIMESTAMP + WHERE run_id = $3 AND cohort_id = $4 AND status = 'sweeping' + `, + [finished ? 'swept' : 'ready', deleted, run.run_id, run.cohort_id], + 'finishCohortMembershipSweep' + ) + + rowsSwept.inc(deleted) + logger.info('Swept stale cohort membership rows', { + run_id: run.run_id, + cohort_id: run.cohort_id, + team_id: run.team_id, + rows: deleted, + finished, + }) + + return deleted + } + + /** + * A reconcile superseded by a cohort edit emits no marker, so its set can never complete. The + * row is only table hygiene at that point. + */ + private async abandonStuckRuns(): Promise { + const result = await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + UPDATE cohort_membership_sweeps + SET status = 'abandoned', updated_at = CURRENT_TIMESTAMP + WHERE status IN ('collecting', 'ready') + AND updated_at < CURRENT_TIMESTAMP - $1::interval + `, + [`${this.config.COHORT_MEMBERSHIP_SWEEP_ABANDON_AFTER_DAYS} days`], + 'abandonCohortMembershipSweeps' + ) + + const abandoned = result.rowCount ?? 0 + if (abandoned > 0) { + sweepsAbandoned.inc(abandoned) + logger.warn('Abandoned reconcile runs that never completed their marker set', { count: abandoned }) + } + + return abandoned + } + + /** Matches the marker topic's retention, past which a run's markers cannot be redelivered. */ + private async collectGarbage(): Promise { + await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + DELETE FROM cohort_membership_sweeps + WHERE status IN ('swept', 'abandoned') + AND updated_at < CURRENT_TIMESTAMP - INTERVAL '30 days' + `, + undefined, + 'gcCohortMembershipSweeps' + ) + } + + private async recordGateWait(): Promise { + const result = await this.postgres.query<{ wait_seconds: string | null }>( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + SELECT EXTRACT(EPOCH FROM (CURRENT_TIMESTAMP - MIN(updated_at)))::text AS wait_seconds + FROM cohort_membership_sweeps WHERE status = 'ready' + `, + undefined, + 'measureCohortMembershipGateWait' + ) + + sweepGateWaitSeconds.set(Number(result.rows[0]?.wait_seconds ?? 0)) + } + + public start(): void { + if (this.running) { + return + } + + this.running = true + this.loopPromise = this.sweepLoop() + logger.info('CohortMembershipSweeper started', { + intervalMs: this.config.COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS, + }) + } + + public async stop(): Promise { + this.running = false + this.sleepResolve?.() + await this.loopPromise + this.loopPromise = null + } + + private async sweepLoop(): Promise { + while (this.running) { + try { + await this.runOnce() + } catch (error) { + sweepCycles.inc({ status: 'error' }) + logger.error('CohortMembershipSweeper cycle failed', { error: String(error) }) + } + + if (this.running) { + let timer: NodeJS.Timeout | undefined + await new Promise((resolve) => { + this.sleepResolve = resolve + timer = setTimeout(resolve, this.config.COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS) + }) + // stop() resolves the sleep early, which leaves the timer armed and the process held open. + clearTimeout(timer) + this.sleepResolve = null + } + } + } +} diff --git a/nodejs/src/common/config/kafka-topics.ts b/nodejs/src/common/config/kafka-topics.ts index b42edb374310..289c67d05064 100644 --- a/nodejs/src/common/config/kafka-topics.ts +++ b/nodejs/src/common/config/kafka-topics.ts @@ -84,6 +84,9 @@ export const KAFKA_CDP_FUNCTION_OVERFLOW = `${prefix}cdp_function_overflow${suff export const KAFKA_CDP_INTERNAL_EVENTS = `${prefix}cdp_internal_events${suffix}` export const KAFKA_CDP_CLICKHOUSE_BEHAVIORAL_COHORTS_MATCHES = `${prefix}clickhouse_behavioral_cohorts_matches${suffix}` export const KAFKA_COHORT_MEMBERSHIP_CHANGED = `${prefix}cohort_membership_changed${suffix}` +// One completion marker per processor partition, certifying that a reconcile run replayed a +// cohort's full membership. Produced by the cohort-stream-processor (Rust). +export const KAFKA_COHORT_RECONCILE_MARKERS = `${prefix}cohort_reconcile_markers${suffix}` // Cross-partition merge protocol trigger consumed by the cohort-stream-processor (Rust). export const KAFKA_PERSON_MERGE_EVENTS = `${prefix}person_merge_events${suffix}` diff --git a/nodejs/src/common/kafka/consumer/consumer-v1.ts b/nodejs/src/common/kafka/consumer/consumer-v1.ts index 3e9df68736cb..75253872ae47 100644 --- a/nodejs/src/common/kafka/consumer/consumer-v1.ts +++ b/nodejs/src/common/kafka/consumer/consumer-v1.ts @@ -378,6 +378,14 @@ export class KafkaConsumer { return [offsets.lowOffset, offsets.highOffset] } + public async committedOffsets(timeout = 10000): Promise { + if (!this.rdKafkaConsumer.isConnected()) { + return [] + } + + return await promisifyCallback((cb) => this.rdKafkaConsumer.committed(timeout, cb)) + } + public async getPartitionsForTopic(topic: string): Promise { if (!this.rdKafkaConsumer.isConnected()) { throw new Error('Not connected') diff --git a/nodejs/src/common/kafka/consumer/consumer-v2.ts b/nodejs/src/common/kafka/consumer/consumer-v2.ts index 6c5114136d37..7280b0397fcc 100644 --- a/nodejs/src/common/kafka/consumer/consumer-v2.ts +++ b/nodejs/src/common/kafka/consumer/consumer-v2.ts @@ -6,8 +6,10 @@ import { LibrdKafkaError, Message, Metadata, + PartitionMetadata, KafkaConsumer as RdKafkaConsumer, TopicPartitionOffset, + WatermarkOffsets, } from 'node-rdkafka' import { hostname } from 'os' @@ -207,6 +209,46 @@ export class KafkaConsumerV2 { return this.rdKafkaConsumer.isConnected() ? this.rdKafkaConsumer.assignments() : [] } + public async queryWatermarkOffsets(topic: string, partition: number, timeout = 10000): Promise<[number, number]> { + if (!this.rdKafkaConsumer.isConnected()) { + throw new Error('Not connected') + } + + const offsets = await promisifyCallback((cb) => + this.rdKafkaConsumer.queryWatermarkOffsets(topic, partition, timeout, cb) + ).catch((err) => { + captureException(err) + logger.error('🔥', 'Failed to query kafka watermark offsets', err) + throw err + }) + + return [offsets.lowOffset, offsets.highOffset] + } + + public async committedOffsets(timeout = 10000): Promise { + if (!this.rdKafkaConsumer.isConnected()) { + return [] + } + + return await promisifyCallback((cb) => this.rdKafkaConsumer.committed(timeout, cb)) + } + + public async getPartitionsForTopic(topic: string): Promise { + if (!this.rdKafkaConsumer.isConnected()) { + throw new Error('Not connected') + } + + const meta = await promisifyCallback((cb) => this.rdKafkaConsumer.getMetadata({ topic }, cb)).catch( + (err) => { + captureException(err) + logger.error('🔥', 'Failed to get partition metadata', err) + throw err + } + ) + + return meta.topics.find((entry) => entry.name === topic)?.partitions ?? [] + } + public offsetsStore(offsets: TopicPartitionOffset[]): void { // Manual offset path used by SessionRecording. Storing offsets after the consumer // has revoked the partitions is harmless — librdkafka will reject the store internally. diff --git a/nodejs/src/common/kafka/consumer/index.ts b/nodejs/src/common/kafka/consumer/index.ts index b3f7d567bedd..c5908d0b7317 100644 --- a/nodejs/src/common/kafka/consumer/index.ts +++ b/nodejs/src/common/kafka/consumer/index.ts @@ -1,4 +1,4 @@ -import { Message, TopicPartitionOffset } from 'node-rdkafka' +import { Message, PartitionMetadata, TopicPartitionOffset } from 'node-rdkafka' import { HealthCheckResult } from '~/types' @@ -22,6 +22,12 @@ export interface KafkaConsumerInterface { disconnect(): Promise isHealthy(): HealthCheckResult offsetsStore(offsets: TopicPartitionOffset[]): void + /** Broker-side `[low, high]` offsets for a partition. `high` is the next offset to be produced. */ + queryWatermarkOffsets(topic: string, partition: number, timeout?: number): Promise<[number, number]> + /** The group's committed offsets for the currently assigned partitions. */ + committedOffsets(timeout?: number): Promise + /** Broker-reported partitions of a topic, whether or not this consumer is assigned them. */ + getPartitionsForTopic(topic: string): Promise } /** diff --git a/nodejs/tests/helpers/sql.ts b/nodejs/tests/helpers/sql.ts index 0b3a0314fa11..0739e0690dbc 100644 --- a/nodejs/tests/helpers/sql.ts +++ b/nodejs/tests/helpers/sql.ts @@ -143,7 +143,11 @@ export async function resetTestDatabase(): Promise { // Helper function to determine which database a table belongs to function getPostgresUseForTable(table: string): PostgresUse { // Behavioral cohorts tables - if (table === 'cohort_membership') { + if ( + table === 'cohort_membership' || + table === 'cohort_membership_sweeps' || + table === 'cohort_membership_consumer_progress' + ) { return PostgresUse.BEHAVIORAL_COHORTS_RW } @@ -552,7 +556,7 @@ export async function resetBehavioralCohortsDatabase(postgres: PostgresRouter): await assertRouterTargetsTestDatabase(postgres, PostgresUse.BEHAVIORAL_COHORTS_RW) await postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, - 'TRUNCATE TABLE cohort_membership', + 'TRUNCATE TABLE cohort_membership, cohort_membership_sweeps, cohort_membership_consumer_progress', undefined, 'reset-behavioral-cohorts-db' ) From 745715a2a7060899798e904e4fb85975d8db51e3 Mon Sep 17 00:00:00 2001 From: Matheus Veras Batista Date: Thu, 27 Aug 2026 09:23:05 -0300 Subject: [PATCH 2/9] fix(cohorts): harden sweep gate, marker path, and shutdown from review --- .../cdp-cohort-membership.consumer.test.ts | 300 ++++++++++---- .../cdp-cohort-membership.consumer.ts | 157 +++++++- .../cohort-membership/sweeper.service.test.ts | 67 +++- .../cohort-membership/sweeper.service.ts | 369 +++++++++++++----- 4 files changed, 693 insertions(+), 200 deletions(-) diff --git a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts index 18f283b42eb0..4be50283fe87 100644 --- a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts +++ b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts @@ -13,15 +13,26 @@ import { CdpCohortMembershipConsumer } from './cdp-cohort-membership.consumer' describe('CdpCohortMembershipConsumer', () => { let hub: Hub let consumer: CdpCohortMembershipConsumer + // Never started: it exercises the flag-on persistence paths directly, without a marker + // consumer connection. + let sweepConsumer: CdpCohortMembershipConsumer let teamId: number + let cluster: string let personId1: string let personId2: string let personId3: string beforeEach(async () => { hub = await createHub() - consumer = new CdpCohortMembershipConsumer(hub, createCdpConsumerDeps(hub)) + const deps = createCdpConsumerDeps(hub) + consumer = new CdpCohortMembershipConsumer(hub, deps) + sweepConsumer = new CdpCohortMembershipConsumer({ ...hub, COHORT_MEMBERSHIP_SWEEP_ENABLED: true }, deps) teamId = Number.parseInt(new UUIDT().toString().replaceAll('-', '').slice(-7), 16) + // Progress rows key on (cluster, partition) and nothing truncates between tests, so a + // unique cluster isolates each test from parallel workers and from prior runs. + cluster = `test-cluster-${new UUIDT().toString()}` + consumer['membershipCluster'] = cluster + sweepConsumer['membershipCluster'] = cluster personId1 = new UUIDT().toString() personId2 = new UUIDT().toString() personId3 = new UUIDT().toString() @@ -32,6 +43,16 @@ describe('CdpCohortMembershipConsumer', () => { }) describe('end-to-end cohort membership processing', () => { + const readProgress = async (): Promise[]> => { + const result = await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + 'SELECT partition, next_offset FROM cohort_membership_consumer_progress WHERE cluster = $1 ORDER BY partition', + [cluster], + 'testQuery' + ) + return result.rows + } + it('should process entered and left events and write to PostgreSQL correctly', async () => { const testEvents = createCohortMembershipEvents([ { @@ -71,21 +92,21 @@ describe('CdpCohortMembershipConsumer', () => { expect(result.rows).toHaveLength(3) expect(result.rows[0]).toMatchObject({ - team_id: teamId.toString(), + team_id: String(teamId), cohort_id: '456', person_id: personId1, in_cohort: true, }) expect(result.rows[1]).toMatchObject({ - team_id: teamId.toString(), + team_id: String(teamId), cohort_id: '456', person_id: personId2, in_cohort: true, }) expect(result.rows[2]).toMatchObject({ - team_id: teamId.toString(), + team_id: String(teamId), cohort_id: '457', person_id: personId3, in_cohort: false, @@ -174,39 +195,54 @@ describe('CdpCohortMembershipConsumer', () => { expect(new Date(thirdTimestamp).getTime()).toBeGreaterThan(new Date(secondTimestamp).getTime()) }) - it('should deduplicate batch entries for the same (team_id, cohort_id, person_id), keeping last in Kafka order', async () => { - const testEvents = createCohortMembershipEvents([ - { - person_id: personId1, - cohort_id: 456, - team_id: teamId, - status: 'entered', - }, - { - person_id: personId1, - cohort_id: 456, - team_id: teamId, - status: 'left', - }, - ]) + it.each([ + ['keeping last in Kafka order for equal versions', {}, {}, false], + // The SQL guard only ever sees the batch's surviving entry, so the in-memory pick has + // to agree with last-writer-wins: a newer version beats a later offset. + [ + 'keeping the highest version regardless of order', + { last_updated: '2026-05-26 13:00:00.000000' }, + { last_updated: '2026-05-26 12:00:00.000000' }, + true, + ], + ])( + 'should deduplicate batch entries for the same (team_id, cohort_id, person_id), %s', + async (_label, firstTags, secondTags, expectedInCohort) => { + const testEvents = createCohortMembershipEvents([ + { + person_id: personId1, + cohort_id: 456, + team_id: teamId, + status: 'entered', + ...firstTags, + }, + { + person_id: personId1, + cohort_id: 456, + team_id: teamId, + status: 'left', + ...secondTags, + }, + ]) - const messages = testEvents.map((event, index) => - createKafkaMessage(event, { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: index }) - ) + const messages = testEvents.map((event, index) => + createKafkaMessage(event, { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: index }) + ) - const cohortMembershipChanges = consumer['_parseAndValidateBatch'](messages) - await consumer['persistCohortMembershipChanges'](cohortMembershipChanges) + const cohortMembershipChanges = consumer['_parseAndValidateBatch'](messages) + await consumer['persistCohortMembershipChanges'](cohortMembershipChanges) - const result = await hub.postgres.query( - PostgresUse.BEHAVIORAL_COHORTS_RW, - 'SELECT * FROM cohort_membership WHERE team_id = $1 AND person_id = $2 AND cohort_id = $3', - [teamId, personId1, 456], - 'testQuery' - ) + const result = await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + 'SELECT * FROM cohort_membership WHERE team_id = $1 AND person_id = $2 AND cohort_id = $3', + [teamId, personId1, 456], + 'testQuery' + ) - expect(result.rows).toHaveLength(1) - expect(result.rows[0].in_cohort).toBe(false) - }) + expect(result.rows).toHaveLength(1) + expect(result.rows[0].in_cohort).toBe(expectedInCohort) + } + ) it.each([ ['seed' as const, new UUIDT().toString()], @@ -216,7 +252,7 @@ describe('CdpCohortMembershipConsumer', () => { createCohortMembershipEvent({ person_id: personId1, cohort_id: 456, - team_id: 1, + team_id: teamId, origin, run_id: runId, }), @@ -235,7 +271,7 @@ describe('CdpCohortMembershipConsumer', () => { createCohortMembershipEvent({ person_id: personId1, cohort_id: 456, - team_id: 1, + team_id: teamId, status: 'entered', origin: 'snapshot', run_id: runId, @@ -250,7 +286,7 @@ describe('CdpCohortMembershipConsumer', () => { expect(change).toMatchObject({ person_id: personId1, cohort_id: 456, - team_id: 1, + team_id: teamId, status: 'entered', }) }) @@ -264,8 +300,8 @@ describe('CdpCohortMembershipConsumer', () => { const result = await hub.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, `SELECT in_cohort, to_char(version, 'YYYY-MM-DD HH24:MI:SS.US') AS version - FROM cohort_membership WHERE team_id = 1 AND person_id = $1 AND cohort_id = $2`, - [personId1, cohortId], + FROM cohort_membership WHERE team_id = $1 AND person_id = $2 AND cohort_id = $3`, + [teamId, personId1, cohortId], 'testQuery' ) return result.rows[0].version @@ -280,14 +316,14 @@ describe('CdpCohortMembershipConsumer', () => { createCohortMembershipEvent({ person_id: personId1, cohort_id: 456, - team_id: 1, + team_id: teamId, last_updated: VERSION, ...tags, }), { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: 0 } ) - await consumer['handleBatch']([message]) + await sweepConsumer['handleBatch']([message]) expect(await readVersion(456)).toEqual(VERSION) }) @@ -297,15 +333,19 @@ describe('CdpCohortMembershipConsumer', () => { ['an equal version is applied', VERSION, VERSION, false, VERSION], ['a newer version is applied', VERSION, NEWER_VERSION, false, NEWER_VERSION], ['any version beats a versionless row', undefined, OLDER_VERSION, false, OLDER_VERSION], + // A change that lost its version (sanitized off-format stamp) is still a membership + // transition: it must apply while leaving the row's stamp in place, not silently + // no-op against the NULL comparison. + ['a versionless change applies without moving the version', VERSION, undefined, false, VERSION], ])( 'should apply last-writer-wins on replay: %s', async (_label, storedVersion, incomingVersion, expectedInCohort, expectedVersion) => { - await consumer['handleBatch']([ + await sweepConsumer['handleBatch']([ createKafkaMessage( createCohortMembershipEvent({ person_id: personId1, cohort_id: 456, - team_id: 1, + team_id: teamId, status: 'entered', last_updated: storedVersion, }), @@ -313,12 +353,12 @@ describe('CdpCohortMembershipConsumer', () => { ), ]) - await consumer['handleBatch']([ + await sweepConsumer['handleBatch']([ createKafkaMessage( createCohortMembershipEvent({ person_id: personId1, cohort_id: 456, - team_id: 1, + team_id: teamId, status: 'left', last_updated: incomingVersion, }), @@ -329,8 +369,8 @@ describe('CdpCohortMembershipConsumer', () => { const result = await hub.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, `SELECT in_cohort, to_char(version, 'YYYY-MM-DD HH24:MI:SS.US') AS version - FROM cohort_membership WHERE team_id = 1 AND person_id = $1 AND cohort_id = 456`, - [personId1], + FROM cohort_membership WHERE team_id = $1 AND person_id = $2 AND cohort_id = 456`, + [teamId, personId1], 'testQuery' ) @@ -346,7 +386,7 @@ describe('CdpCohortMembershipConsumer', () => { createCohortMembershipEvent({ person_id: personId, cohort_id: 456, - team_id: 1, + team_id: teamId, origin: 'reconcile', run_id: runId, last_updated: lastUpdated, @@ -354,7 +394,7 @@ describe('CdpCohortMembershipConsumer', () => { { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: 0 } ) - await consumer['handleBatch']([ + await sweepConsumer['handleBatch']([ reconcileRow(personId1, NEWER_VERSION), reconcileRow(personId2, VERSION), // A live transition rides the same topic and must not count as an assertion. @@ -362,7 +402,7 @@ describe('CdpCohortMembershipConsumer', () => { createCohortMembershipEvent({ person_id: personId3, cohort_id: 456, - team_id: 1, + team_id: teamId, last_updated: OLDER_VERSION, }), { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: 0 } @@ -383,7 +423,7 @@ describe('CdpCohortMembershipConsumer', () => { expect(await readSweep()).toEqual([ { - team_id: '1', + team_id: String(teamId), marker_bits: '0', status: 'collecting', snapshot_rows: '2', @@ -391,11 +431,11 @@ describe('CdpCohortMembershipConsumer', () => { }, ]) - await consumer['handleBatch']([reconcileRow(personId3, OLDER_VERSION)]) + await sweepConsumer['handleBatch']([reconcileRow(personId3, OLDER_VERSION)]) expect(await readSweep()).toEqual([ { - team_id: '1', + team_id: String(teamId), marker_bits: '0', status: 'collecting', snapshot_rows: '3', @@ -405,24 +445,17 @@ describe('CdpCohortMembershipConsumer', () => { }) it('should advance consumer progress to the next offset without ever regressing it', async () => { - const readProgress = async () => { - const result = await hub.postgres.query( - PostgresUse.BEHAVIORAL_COHORTS_RW, - 'SELECT partition, next_offset FROM cohort_membership_consumer_progress ORDER BY partition', - undefined, - 'testQuery' - ) - return result.rows - } - const rowOnPartition = (partition: number, offset: number) => - createKafkaMessage(createCohortMembershipEvent({ person_id: new UUIDT().toString(), team_id: 1 }), { - topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, - partition, - offset, - }) + createKafkaMessage( + createCohortMembershipEvent({ person_id: new UUIDT().toString(), team_id: teamId }), + { + topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, + partition, + offset, + } + ) - await consumer['handleBatch']([rowOnPartition(3, 10), rowOnPartition(3, 11), rowOnPartition(7, 5)]) + await sweepConsumer['handleBatch']([rowOnPartition(3, 10), rowOnPartition(3, 11), rowOnPartition(7, 5)]) expect(await readProgress()).toEqual([ { partition: 3, next_offset: '12' }, @@ -430,7 +463,7 @@ describe('CdpCohortMembershipConsumer', () => { ]) // A rebalance can replay from an older committed offset; the gate must not walk back. - await consumer['handleBatch']([rowOnPartition(3, 0)]) + await sweepConsumer['handleBatch']([rowOnPartition(3, 0)]) expect(await readProgress()).toEqual([ { partition: 3, next_offset: '12' }, @@ -438,6 +471,122 @@ describe('CdpCohortMembershipConsumer', () => { ]) }) + it('should keep the pre-sweep write shape while the sweep flag is off', async () => { + // The migrations are a prerequisite of the flag flip, not of the deploy, so a flag-off + // batch must not touch the version column or the sweep tables. + const runId = new UUIDT().toString() + const message = createKafkaMessage( + createCohortMembershipEvent({ + person_id: personId1, + cohort_id: 456, + team_id: teamId, + origin: 'reconcile', + run_id: runId, + last_updated: VERSION, + }), + { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: 0 } + ) + + await consumer['handleBatch']([message]) + + const membership = await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + 'SELECT in_cohort, version FROM cohort_membership WHERE team_id = $1 AND person_id = $2', + [teamId, personId1], + 'testQuery' + ) + expect(membership.rows).toEqual([{ in_cohort: true, version: null }]) + + const sweeps = await hub.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + 'SELECT 1 FROM cohort_membership_sweeps WHERE run_id = $1', + [runId], + 'testQuery' + ) + expect(sweeps.rows).toHaveLength(0) + expect(await readProgress()).toHaveLength(0) + }) + + it('should sanitize a version the producer contract does not allow instead of failing the batch', () => { + const message = createKafkaMessage( + createCohortMembershipEvent({ + person_id: personId1, + cohort_id: 456, + team_id: teamId, + last_updated: '2026-05-26T12:34:56Z', + }), + { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: 0 } + ) + + const [change] = sweepConsumer['_parseAndValidateBatch']([message]) + + expect(change.last_updated).toBeUndefined() + }) + + it('should contain a marker batch failure instead of taking the consumer down', async () => { + // The marker batch's offsets commit either way, and a throw would escape to the + // consumer loop and stop membership ingestion with it. + sweepConsumer['sweeper'] = { + parseMarkers: () => { + throw new Error('marker path failure') + }, + } as any + + await expect(sweepConsumer['handleMarkerBatch']([])).resolves.toBeUndefined() + }) + + it('should refuse to capture watermarks from a partial partition list', async () => { + await sweepConsumer['upsertConsumerProgress'](PostgresUse.BEHAVIORAL_COHORTS_RW, [ + { partition: 0, nextOffset: 5 }, + { partition: 1, nextOffset: 7 }, + { partition: 2, nextOffset: 9 }, + ]) + + const stubConsumer = { + getPartitionsForTopic: jest.fn(), + queryWatermarkOffsets: jest.fn((_topic: string, partition: number) => + Promise.resolve([0, 10 + partition]) + ), + } + sweepConsumer['kafkaConsumer'] = stubConsumer as any + + // Three partitions have durable consumer progress on this cluster, so a broker + // answering with two is a partial list, and capturing it would let the gate pass while + // the missing partition still holds unconsumed snapshot rows. + stubConsumer.getPartitionsForTopic.mockResolvedValue([{ id: 0 }, { id: 1 }]) + await expect(sweepConsumer['captureMembershipWatermarks']()).rejects.toThrow('Partial partition metadata') + + stubConsumer.getPartitionsForTopic.mockResolvedValue([{ id: 0 }, { id: 1 }, { id: 2 }]) + await expect(sweepConsumer['captureMembershipWatermarks']()).resolves.toEqual({ + 0: 10, + 1: 11, + 2: 12, + }) + }) + + it('should refresh progress from committed offsets, skipping invalid ones and never regressing', async () => { + const stubConsumer = { committedOffsets: jest.fn() } + sweepConsumer['kafkaConsumer'] = stubConsumer as any + + // -1001 is rdkafka's "no committed offset" sentinel; another topic's offsets must not + // leak into the membership gate. + stubConsumer.committedOffsets.mockResolvedValue([ + { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, partition: 0, offset: 5 }, + { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, partition: 1, offset: -1001 }, + { topic: 'some_other_topic', partition: 2, offset: 9 }, + ]) + await sweepConsumer['refreshConsumerProgressFromCommittedOffsets']() + + expect(await readProgress()).toEqual([{ partition: 0, next_offset: '5' }]) + + stubConsumer.committedOffsets.mockResolvedValue([ + { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, partition: 0, offset: 3 }, + ]) + await sweepConsumer['refreshConsumerProgressFromCommittedOffsets']() + + expect(await readProgress()).toEqual([{ partition: 0, next_offset: '5' }]) + }) + it('should reject entire batch when invalid messages are present', async () => { const validEvent = { person_id: personId1, @@ -485,20 +634,13 @@ describe('CdpCohortMembershipConsumer', () => { // A non-integer cohort_id passes schema validation but Postgres rejects it, so the // membership write fails while the progress write on its own would have succeeded. const message = createKafkaMessage( - createCohortMembershipEvent({ person_id: personId1, team_id: 1, cohort_id: 1.5 }), + createCohortMembershipEvent({ person_id: personId1, team_id: teamId, cohort_id: 1.5 }), { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, partition: 9, offset: 42 } ) - await expect(consumer['handleBatch']([message])).rejects.toThrow() - - const progress = await hub.postgres.query( - PostgresUse.BEHAVIORAL_COHORTS_RW, - 'SELECT * FROM cohort_membership_consumer_progress', - undefined, - 'testQuery' - ) + await expect(sweepConsumer['handleBatch']([message])).rejects.toThrow() - expect(progress.rows).toHaveLength(0) + expect(await readProgress()).toHaveLength(0) }) it('should not produce side effects when database insertion fails', async () => { diff --git a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts index 11776e5fa954..ea6121adfd17 100644 --- a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts +++ b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts @@ -2,6 +2,7 @@ import { Message } from 'node-rdkafka' import { z } from 'zod' import { KAFKA_COHORT_MEMBERSHIP_CHANGED, KAFKA_COHORT_RECONCILE_MARKERS } from '~/common/config/kafka-topics' +import { getKafkaConfigFromEnv } from '~/common/kafka/config' import { KafkaConsumerInterface, RdKafkaConsumerConfig, createKafkaConsumer } from '~/common/kafka/consumer' import { instrumentFn } from '~/common/tracing/tracing-utils' import { PostgresUse, TransactionClient } from '~/common/utils/db/postgres' @@ -10,7 +11,11 @@ import { logger } from '~/common/utils/logger' import { HealthCheckResult } from '../../types' import type { CdpConfig } from '../config' -import { CohortMembershipSweeper, MembershipWatermarks } from '../services/cohort-membership/sweeper.service' +import { + CohortMembershipSweeper, + MembershipWatermarks, + PRODUCER_VERSION_FORMAT, +} from '../services/cohort-membership/sweeper.service' import { CdpConsumerBase, CdpConsumerBaseConfig, CdpConsumerBaseDeps } from './cdp-base.consumer' export type CdpCohortMembershipConsumerConfig = CdpConsumerBaseConfig & @@ -26,6 +31,17 @@ export type CdpCohortMembershipConsumerConfig = CdpConsumerBaseConfig & // Markers from runs that predate the sweeper are useless — the next reconcile heals anyway. const startAtLatest = { ['auto.offset.reset' as keyof RdKafkaConsumerConfig]: 'latest' as never } +/** + * The broker list the membership consumer resolves, mirroring the consumer's own config + * resolution. Progress offsets are only comparable against watermarks from the same cluster, so + * this keys every progress row: after a cluster move the old rows stop matching, and the gate + * starts from no progress instead of comparing the old cluster's large offsets against the new + * cluster's small watermarks. + */ +function membershipClusterId(): string { + return String(getKafkaConfigFromEnv('CONSUMER')['metadata.broker.list'] ?? 'kafka:9092') +} + // `origin` is a loose string, not an enum: nothing branches on it, and rejecting a value the // processor added later would throw the whole batch and crash-loop on redelivery. const CohortMembershipChangeSchema = z.object({ @@ -78,6 +94,9 @@ function partitionProgressFromMessages(messages: Message[]): PartitionProgress[] * * Versions are compared as strings: the producer's format is fixed-width, so lexicographic order * is chronological order. A row carrying no version bounds nothing and is left out. + * + * The result is sorted by (run_id, cohort_id) because it feeds one multi-row upsert whose row + * locks are held to commit: two pods writing the same runs in different orders would deadlock. */ function collectSnapshotMarks(changes: CohortMembershipChange[]): SnapshotMark[] { const marks = new Map() @@ -116,7 +135,7 @@ function collectSnapshotMarks(changes: CohortMembershipChange[]): SnapshotMark[] } } - return Array.from(marks.values()) + return Array.from(marks.values()).sort((a, b) => a.runId.localeCompare(b.runId) || a.cohortId - b.cohortId) } export class CdpCohortMembershipConsumer extends CdpConsumerBase { @@ -124,6 +143,7 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase this.captureMembershipWatermarks(), refreshConsumerProgress: () => this.refreshConsumerProgressFromCommittedOffsets(), }) @@ -179,6 +200,26 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase( + PostgresUse.BEHAVIORAL_COHORTS_RW, + 'SELECT COUNT(*) AS count FROM cohort_membership_consumer_progress WHERE cluster = $1', + [this.membershipCluster], + 'countCohortMembershipProgressPartitions' + ) + + const knownPartitions = Number(known.rows[0]?.count ?? 0) + if (metadata.length < knownPartitions) { + throw new Error( + `Partial partition metadata for ${KAFKA_COHORT_MEMBERSHIP_CHANGED}: ` + + `broker reported ${metadata.length}, consumer progress covers ${knownPartitions}` + ) + } + const results = await Promise.all( metadata.map(async ({ id: partition }) => { const [, high] = await this.kafkaConsumer.queryWatermarkOffsets( @@ -201,6 +242,11 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase() for (const change of changes) { - deduped.set(`${change.team_id}:${change.cohort_id}:${change.person_id}`, change) + const key = `${change.team_id}:${change.cohort_id}:${change.person_id}` + const existing = deduped.get(key) + if (!existing || (change.last_updated ?? '') >= (existing.last_updated ?? '')) { + deduped.set(key, change) + } + } + + if (!this.config.COHORT_MEMBERSHIP_SWEEP_ENABLED) { + if (deduped.size > 0) { + await this.upsertMembershipRowsWithoutVersion(Array.from(deduped.values())) + } + return } const marks = collectSnapshotMarks(changes) @@ -287,6 +346,35 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { + const values: any[] = [] + const placeholders: string[] = [] + let paramIndex = 1 + + for (const change of changes) { + placeholders.push( + `($${paramIndex}, $${paramIndex + 1}, $${paramIndex + 2}, $${paramIndex + 3}, CURRENT_TIMESTAMP)` + ) + values.push(change.team_id, change.cohort_id, change.person_id, change.status === 'entered') + paramIndex += 4 + } + + await this.deps.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + INSERT INTO cohort_membership (team_id, cohort_id, person_id, in_cohort, last_updated) + VALUES ${placeholders.join(', ')} + ON CONFLICT (team_id, cohort_id, person_id) + DO UPDATE SET + in_cohort = EXCLUDED.in_cohort, + last_updated = CURRENT_TIMESTAMP + `, + values, + 'batchUpsertCohortMembership' + ) + } + /** * Folds this batch's contribution into each run's ledger row. The row is created here when the * snapshot arrives before the run's completion markers, which is the normal ordering. @@ -339,17 +427,17 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { const cohortMembershipChanges = this._parseAndValidateBatch(messages) - await this.persistCohortMembershipChanges(cohortMembershipChanges, partitionProgressFromMessages(messages)) + const progress = this.config.COHORT_MEMBERSHIP_SWEEP_ENABLED ? partitionProgressFromMessages(messages) : [] + await this.persistCohortMembershipChanges(cohortMembershipChanges, progress) } public override async start(): Promise { @@ -427,23 +533,50 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { const sweeper = this.sweeper if (!sweeper) { return } - await sweeper.applyMarkers(sweeper.parseMarkers(messages)) + try { + await sweeper.applyMarkers(sweeper.parseMarkers(messages)) + } catch (error) { + logger.error('Failed to process a reconcile marker batch', { + error: String(error), + count: messages.length, + }) + } } public override async stop(): Promise { logger.info('💤', `Stopping ${this.name}...`) await this.sweeper?.stop() - await Promise.all([this.kafkaConsumer.disconnect(), this.markerKafkaConsumer?.disconnect()]) + // The marker path reads watermarks off the membership client, so that client disconnects + // only after the marker consumer has fully drained. + await this.markerKafkaConsumer?.disconnect() + await this.kafkaConsumer.disconnect() // IMPORTANT: super always comes last await super.stop() diff --git a/nodejs/src/cdp/services/cohort-membership/sweeper.service.test.ts b/nodejs/src/cdp/services/cohort-membership/sweeper.service.test.ts index 3d113c8354b2..e3ddcd615a89 100644 --- a/nodejs/src/cdp/services/cohort-membership/sweeper.service.test.ts +++ b/nodejs/src/cdp/services/cohort-membership/sweeper.service.test.ts @@ -31,6 +31,8 @@ describe('CohortMembershipSweeper', () => { COHORT_MEMBERSHIP_SWEEP_ABANDON_AFTER_DAYS: 3, } + const CLUSTER = 'test-cluster' + let postgres: PostgresRouter let sweeper: CohortMembershipSweeper let watermarks: MembershipWatermarks @@ -42,6 +44,7 @@ describe('CohortMembershipSweeper', () => { watermarks = { 0: 10 } runId = new UUIDT().toString() sweeper = new CohortMembershipSweeper(config, postgres, { + cluster: CLUSTER, captureMembershipWatermarks: () => Promise.resolve(watermarks), refreshConsumerProgress: () => Promise.resolve(), }) @@ -108,19 +111,23 @@ describe('CohortMembershipSweeper', () => { await query( `INSERT INTO cohort_membership_sweeps (run_id, cohort_id, team_id, marker_bits, min_marker_version, min_snapshot_version, - membership_hwms, status) - VALUES ($1, $2, $3, -1, $4::timestamp, $5::timestamp, $6::jsonb, 'ready')`, + membership_hwms, status, ready_at) + VALUES ($1, $2, $3, -1, $4::timestamp, $5::timestamp, $6::jsonb, 'ready', CURRENT_TIMESTAMP)`, [runId, COHORT_ID, TEAM_ID, MARKER_VERSION, THRESHOLD, JSON.stringify(watermarks)] ) } + const insertProgress = async (partition: number, nextOffset: number): Promise => { + await query( + `INSERT INTO cohort_membership_consumer_progress (cluster, partition, next_offset) VALUES ($1, $2, $3)`, + [CLUSTER, partition, nextOffset] + ) + } + /** Record the consumer as having applied everything the run's watermarks cover. */ const openGate = async (): Promise => { for (const [partition, highWatermark] of Object.entries(watermarks)) { - await query(`INSERT INTO cohort_membership_consumer_progress (partition, next_offset) VALUES ($1, $2)`, [ - Number(partition), - highWatermark, - ]) + await insertProgress(Number(partition), highWatermark) } } @@ -151,12 +158,23 @@ describe('CohortMembershipSweeper', () => { expect(await readSweep()).toMatchObject({ status: 'collecting' }) + // An incomplete set must not be promoted by the loop either. + await sweeper.runOnce() + + expect(await readSweep()).toMatchObject({ status: 'collecting', membership_hwms: null }) + watermarks = { 0: 10, 3: 4 } - // The last partition sets the sign bit, which is what completes the i64 bitmap. + // The last partition sets the sign bit, which is what completes the i64 bitmap. The + // watermark capture and the 'ready' transition happen on the next sweep tick, not in the + // marker batch, so a transient Kafka failure retries instead of killing the pod. await sweeper.applyMarkers( sweeper.parseMarkers([markerMessage(COHORT_PARTITION_COUNT - 1, { last_updated: ABOVE_THRESHOLD })]) ) + expect(await readSweep()).toMatchObject({ marker_bits: '-1', status: 'collecting', membership_hwms: null }) + + await sweeper.runOnce() + expect(await readSweep()).toMatchObject({ marker_bits: '-1', status: 'ready', @@ -173,7 +191,8 @@ describe('CohortMembershipSweeper', () => { expect(await sweeper.runOnce()).toMatchObject({ swept: 0, blocked: 1 }) expect(await survivingPersonIds()).toEqual([stale]) - await query(`INSERT INTO cohort_membership_consumer_progress (partition, next_offset) VALUES (0, 9), (3, 4)`) + await insertProgress(0, 9) + await insertProgress(3, 4) expect(await sweeper.runOnce()).toMatchObject({ swept: 0, blocked: 1 }) expect(await survivingPersonIds()).toEqual([stale]) @@ -190,6 +209,9 @@ describe('CohortMembershipSweeper', () => { const stale1 = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) const stale2 = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + const stale3 = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + // Only versioned rows are sweepable: a NULL was written either before the feature or by a + // rolled-back image asserting current membership, and the two cannot be told apart here. const neverVersioned = await insertMembership(TEAM_ID, COHORT_ID, null) const asserted = await insertMembership(TEAM_ID, COHORT_ID, THRESHOLD) const liveChange = await insertMembership(TEAM_ID, COHORT_ID, ABOVE_THRESHOLD) @@ -199,8 +221,10 @@ describe('CohortMembershipSweeper', () => { expect(await sweeper.runOnce()).toMatchObject({ swept: 1, rowsDeleted: 3 }) const surviving = await survivingPersonIds() - expect(surviving).toEqual([asserted, liveChange, otherCohort, otherTeam].sort((a, b) => a.localeCompare(b))) - expect(surviving).not.toEqual(expect.arrayContaining([stale1, stale2, neverVersioned])) + expect(surviving).toEqual( + [neverVersioned, asserted, liveChange, otherCohort, otherTeam].sort((a, b) => a.localeCompare(b)) + ) + expect(surviving).not.toEqual(expect.arrayContaining([stale1, stale2, stale3])) expect(await readSweep()).toMatchObject({ status: 'swept', swept_rows: '3' }) // A second cycle must not re-sweep a finished run. @@ -213,7 +237,7 @@ describe('CohortMembershipSweeper', () => { ])('should refuse to sweep a run when %s', async (_label, damage) => { watermarks = { 0: 10 } await readyRun() - await query(`INSERT INTO cohort_membership_consumer_progress (partition, next_offset) VALUES (0, 10)`) + await insertProgress(0, 10) await query(`UPDATE cohort_membership_sweeps SET ${damage} WHERE run_id = $1`, [runId]) const live = await insertMembership(TEAM_ID, COHORT_ID, ABOVE_THRESHOLD) @@ -238,7 +262,7 @@ describe('CohortMembershipSweeper', () => { expect(await survivingPersonIds()).not.toContain(stale) }) - it('should abandon a run whose marker set never completed', async () => { + it('should abandon a run whose marker set never completed, but not a gate-blocked ready run', async () => { await query( `INSERT INTO cohort_membership_sweeps (run_id, cohort_id, team_id, marker_bits, created_at, updated_at) VALUES ($1, $2, $3, 7, CURRENT_TIMESTAMP - INTERVAL '30 days', @@ -246,7 +270,26 @@ describe('CohortMembershipSweeper', () => { [runId, COHORT_ID, TEAM_ID] ) + // A fully-proven run held back only by consumer lag must survive the short timeout: it + // sweeps whenever the gate opens, and discarding it would leave the cohort stale. + const blockedRunId = new UUIDT().toString() + await query( + `INSERT INTO cohort_membership_sweeps + (run_id, cohort_id, team_id, marker_bits, min_marker_version, min_snapshot_version, + membership_hwms, status, ready_at, created_at, updated_at) + VALUES ($1, $2, $3, -1, $4::timestamp, $5::timestamp, $6::jsonb, 'ready', + CURRENT_TIMESTAMP - INTERVAL '5 days', CURRENT_TIMESTAMP - INTERVAL '5 days', + CURRENT_TIMESTAMP - INTERVAL '5 days')`, + [blockedRunId, COHORT_ID + 1, TEAM_ID, MARKER_VERSION, THRESHOLD, JSON.stringify({ 0: 10 })] + ) + expect(await sweeper.runOnce()).toMatchObject({ abandoned: 1 }) expect(await readSweep()).toMatchObject({ status: 'abandoned' }) + + const blocked = await query<{ status: string }>( + `SELECT status FROM cohort_membership_sweeps WHERE run_id = $1`, + [blockedRunId] + ) + expect(blocked[0].status).toEqual('ready') }) }) diff --git a/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts b/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts index 2a1b4f4ad285..a147895849f9 100644 --- a/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts +++ b/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts @@ -5,17 +5,30 @@ import { z } from 'zod' import { PostgresRouter, PostgresUse } from '~/common/utils/db/postgres' import { parseJSON } from '~/common/utils/json-parse' import { logger } from '~/common/utils/logger' +import { UUIDT } from '~/common/utils/utils' /** - * One completion marker per processor partition certifies a reconcile run. Kept in step with - * `COHORT_PARTITION_COUNT` in `rust/cohort-core/src/partitioner.rs`, which is compile-time stable - * by design — the marker topic's partition count is load-bearing for the same reason. + * One completion marker per processor partition certifies a reconcile run. Must match the + * processor's `cohort_partition_count` (`rust/cohort-stream-processor/src/config.rs`), which is 64 + * in production but env-overridable; on an environment running fewer partitions the bitmap never + * completes and every sweep is a no-op, surfaced only by runs aging into abandonment. */ export const COHORT_PARTITION_COUNT = 64 +/** + * The producer's fixed-width version stamp, so lexicographic order is chronological order. Pinned + * wherever a stamp is bound into a `::timestamp` cast: anything Postgres cannot parse would throw + * out of the batch and take the consumer down with it. + */ +export const PRODUCER_VERSION_FORMAT = /^\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{6}$/ + /** All 64 bits set, read as a signed 64-bit integer. Same convention as the seeder's ledger. */ const ALL_MARKERS = '-1' +/** The marker topic's retention. Past this a run's markers cannot be redelivered, so a run still + * unswept by then is unrecoverable and its ledger row is only hygiene. */ +const MARKER_RETENTION_DAYS = 30 + /** A tick sweeps at most this many runs; the rest wait for the next one. */ const MAX_RUNS_PER_TICK = 10 @@ -36,10 +49,7 @@ const ReconcileCompleteMarkerSchema = z.object({ .min(0) .max(COHORT_PARTITION_COUNT - 1), run_id: z.guid(), - // The producer's fixed-width stamp. Pinned here because it is bound straight into a - // `::timestamp` cast: anything Postgres cannot parse would throw out of the marker batch and - // take the membership consumer down with it. - last_updated: z.string().regex(/^\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{6}$/), + last_updated: z.string().regex(PRODUCER_VERSION_FORMAT), }) export type ReconcileCompleteMarker = z.infer @@ -67,6 +77,7 @@ type ClaimableRun = { team_id: string membership_hwms: MembershipWatermarks | null threshold: string | null + claim_token: string | null } const markersConsumed = new Counter({ @@ -92,7 +103,8 @@ const sweepsGateBlocked = new Counter({ const sweepsAbandoned = new Counter({ name: 'cdp_cohort_membership_sweeps_abandoned', - help: 'Reconcile runs given up on because their marker set never completed', + help: 'Reconcile runs given up on, by the state they were stuck in', + labelNames: ['reason'], }) const sweepCycles = new Counter({ @@ -103,27 +115,31 @@ const sweepCycles = new Counter({ const sweepGateWaitSeconds = new Gauge({ name: 'cdp_cohort_membership_sweep_gate_wait_seconds', - help: 'Time since the oldest fully-marked run was last written to, while it waits on the gate', + help: 'Time since the oldest gate-blocked run completed its marker set', }) /** * Deletes `cohort_membership` rows that a completed reconcile run did not re-assert. * * The consumer only ever upserts, so a person who stops matching an edited cohort keeps an - * `in_cohort = true` row forever. A reconcile run replays the cohort's full current membership and - * then certifies itself with one marker per partition; everything the run stamped older than - * itself is by definition no longer a member. + * `in_cohort = true` row forever. A reconcile run emits one row per person holding a `cf_stage2` + * register row for the cohort and then certifies itself with one marker per partition; persons + * with no register row get nothing, and their versioned rows older than the run are what the + * sweep deletes. */ export class CohortMembershipSweeper { private running = false + private stopping = false private loopPromise: Promise | null = null private sleepResolve: (() => void) | null = null constructor( private config: CohortMembershipSweepConfig, private postgres: PostgresRouter, - /** Both read the membership consumer's own client: the two topics may live on different clusters. */ + /** All three describe the membership consumer's own client: the marker topic may live on another cluster. */ private kafka: { + /** Broker identity the progress rows are keyed by; offsets are only comparable within one cluster. */ + cluster: string captureMembershipWatermarks: () => Promise refreshConsumerProgress: () => Promise } @@ -165,56 +181,89 @@ export class CohortMembershipSweeper { } /** - * Folding a marker in is a bit union, so redelivery changes nothing. The run turns claimable - * only on the transition to a complete set, and only once. + * Folding a marker in is a bit union, so redelivery changes nothing. This path runs inside the + * marker consumer's batch, whose offsets commit whether or not each fold landed, so a failed + * fold is caught and counted rather than rethrown: a rethrow would take the membership consumer + * in the same process down with it. A lost fold leaves the run one bit short until it ages into + * abandonment, which `sweepsAbandoned` surfaces, and the next reconcile heals the cohort. */ public async applyMarkers(markers: ReconcileCompleteMarker[]): Promise { for (const marker of markers) { - const result = await this.postgres.query<{ complete: boolean; status: string }>( - PostgresUse.BEHAVIORAL_COHORTS_RW, - ` - INSERT INTO cohort_membership_sweeps - (run_id, cohort_id, team_id, marker_bits, min_marker_version) - VALUES ($1, $2, $3, (1::bigint << $4), $5::timestamp) - ON CONFLICT (run_id, cohort_id) - DO UPDATE SET - marker_bits = cohort_membership_sweeps.marker_bits | (1::bigint << $4), - min_marker_version = LEAST( - cohort_membership_sweeps.min_marker_version, - EXCLUDED.min_marker_version - ), - updated_at = CURRENT_TIMESTAMP - RETURNING marker_bits = ${ALL_MARKERS} AS complete, status - `, - [marker.run_id, marker.cohort_id, marker.team_id, marker.partition, marker.last_updated], - 'applyCohortReconcileMarker' - ) - - const { complete, status } = result.rows[0] - if (!complete || status !== 'collecting') { - continue + try { + await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + INSERT INTO cohort_membership_sweeps + (run_id, cohort_id, team_id, marker_bits, min_marker_version) + VALUES ($1, $2, $3, (1::bigint << $4), $5::timestamp) + ON CONFLICT (run_id, cohort_id) + DO UPDATE SET + marker_bits = cohort_membership_sweeps.marker_bits | (1::bigint << $4), + min_marker_version = LEAST( + cohort_membership_sweeps.min_marker_version, + EXCLUDED.min_marker_version + ), + updated_at = CURRENT_TIMESTAMP + `, + [marker.run_id, marker.cohort_id, marker.team_id, marker.partition, marker.last_updated], + 'applyCohortReconcileMarker' + ) + } catch (error) { + markersConsumed.inc({ outcome: 'apply_failed' }) + logger.error('Failed to apply a reconcile marker', { + error: String(error), + run_id: marker.run_id, + cohort_id: marker.cohort_id, + partition: marker.partition, + }) } + } + } + + /** + * Turns fully-marked runs claimable. This lives in the sweep loop rather than the marker batch + * because capturing watermarks talks to Kafka and can fail transiently; the loop retries every + * tick until it lands, while a throw inside the marker batch would kill the pod. + */ + private async promoteMarkedRuns(): Promise { + const marked = await this.postgres.query<{ run_id: string; cohort_id: string; team_id: string }>( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + SELECT run_id, cohort_id, team_id FROM cohort_membership_sweeps + WHERE status = 'collecting' AND marker_bits = ${ALL_MARKERS} + LIMIT ${MAX_CANDIDATES_PER_TICK} + `, + undefined, + 'findMarkedCohortMembershipSweeps' + ) + + if (marked.rows.length === 0) { + return + } - // Every snapshot row of this run was acked before its marker was produced, so the - // watermarks read now bound the whole snapshot on the membership topic. - const watermarks = await this.kafka.captureMembershipWatermarks() + // Every snapshot row of these runs was acked before its last marker was produced, so + // watermarks read now bound each whole snapshot on the membership topic. Reading later + // than the completion only raises the bar the gate waits for, never lowers it. + const watermarks = await this.kafka.captureMembershipWatermarks() + for (const run of marked.rows) { await this.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, ` UPDATE cohort_membership_sweeps - SET status = 'ready', membership_hwms = $1::jsonb, updated_at = CURRENT_TIMESTAMP + SET status = 'ready', membership_hwms = $1::jsonb, + ready_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP WHERE run_id = $2 AND cohort_id = $3 AND status = 'collecting' AND marker_bits = ${ALL_MARKERS} `, - [JSON.stringify(watermarks), marker.run_id, marker.cohort_id], + [JSON.stringify(watermarks), run.run_id, run.cohort_id], 'readyCohortMembershipSweep' ) logger.info('Reconcile run is fully marked and ready to sweep', { - run_id: marker.run_id, - cohort_id: marker.cohort_id, - team_id: marker.team_id, + run_id: run.run_id, + cohort_id: run.cohort_id, + team_id: run.team_id, }) } } @@ -222,10 +271,15 @@ export class CohortMembershipSweeper { public async runOnce(): Promise { const result: SweepCycleResult = { swept: 0, rowsDeleted: 0, blocked: 0, abandoned: 0 } + await this.promoteMarkedRuns() await this.kafka.refreshConsumerProgress() const progress = await this.readConsumerProgress() for (const run of await this.claimableRuns()) { + if (this.stopping) { + break + } + if (!this.gateSatisfied(run, progress)) { result.blocked += 1 sweepsGateBlocked.inc() @@ -241,7 +295,12 @@ export class CohortMembershipSweeper { continue } - result.rowsDeleted += await this.sweep(claimed) + const deleted = await this.sweep(claimed) + if (deleted === null) { + continue + } + + result.rowsDeleted += deleted result.swept += 1 sweepsExecuted.inc() } @@ -261,15 +320,17 @@ export class CohortMembershipSweeper { * A run with no snapshot minimum is never claimable. Marker stamps alone cannot bound the * snapshot: a fast partition finishes and stamps its marker while a slower one is still * emitting rows, so the earliest marker can sit above rows the run asserted. Sweeping there - * would delete live members — and every unversioned row with them. Such a run waits for - * abandonment instead; the next reconcile heals the cohort. + * would delete live members. Such a run waits for abandonment instead. This also excludes a + * run whose snapshot was legitimately empty (a cohort edited so nobody matches): it can never + * prove a snapshot minimum, so its stale rows outlive every reconcile until they are cleaned + * up out of band. */ private async claimableRuns(): Promise { const result = await this.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, ` SELECT run_id, cohort_id, team_id, membership_hwms, - LEAST(min_snapshot_version, min_marker_version) AS threshold + LEAST(min_snapshot_version, min_marker_version) AS threshold, claim_token FROM cohort_membership_sweeps WHERE min_snapshot_version IS NOT NULL AND (status = 'ready' @@ -287,8 +348,8 @@ export class CohortMembershipSweeper { private async readConsumerProgress(): Promise> { const result = await this.postgres.query<{ partition: number; next_offset: string }>( PostgresUse.BEHAVIORAL_COHORTS_RW, - 'SELECT partition, next_offset FROM cohort_membership_consumer_progress', - undefined, + 'SELECT partition, next_offset FROM cohort_membership_consumer_progress WHERE cluster = $1', + [this.kafka.cluster], 'readCohortMembershipProgress' ) @@ -341,98 +402,171 @@ export class CohortMembershipSweeper { } private async claim(run: ClaimableRun): Promise { + const claimToken = new UUIDT().toString() + const result = await this.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, ` UPDATE cohort_membership_sweeps - SET status = 'sweeping', claimed_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP + SET status = 'sweeping', claimed_at = CURRENT_TIMESTAMP, claim_token = $4, + updated_at = CURRENT_TIMESTAMP WHERE run_id = $1 AND cohort_id = $2 AND min_snapshot_version IS NOT NULL AND (status = 'ready' OR (status = 'sweeping' AND claimed_at < CURRENT_TIMESTAMP - $3::interval)) RETURNING run_id, cohort_id, team_id, membership_hwms, - LEAST(min_snapshot_version, min_marker_version) AS threshold + LEAST(min_snapshot_version, min_marker_version) AS threshold, claim_token `, - [run.run_id, run.cohort_id, `${this.config.COHORT_MEMBERSHIP_SWEEP_CLAIM_TIMEOUT_MS} milliseconds`], + [ + run.run_id, + run.cohort_id, + `${this.config.COHORT_MEMBERSHIP_SWEEP_CLAIM_TIMEOUT_MS} milliseconds`, + claimToken, + ], 'claimCohortMembershipSweep' ) return result.rows[0] ?? null } - private async sweep(run: ClaimableRun): Promise { + /** Returns the deleted row count, or null when the run was refused rather than swept. */ + private async sweep(run: ClaimableRun): Promise { if (!run.threshold) { + // Unreachable while claim() requires a snapshot minimum; kept so a future claim change + // cannot silently delete below no bound. Abandoning gives the row a terminal state + // instead of leaving it wedged in 'sweeping'. logger.error('Refusing to sweep a run with no threshold', { run_id: run.run_id, cohort_id: run.cohort_id, }) - return 0 + await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + UPDATE cohort_membership_sweeps + SET status = 'abandoned', claimed_at = NULL, claim_token = NULL, + updated_at = CURRENT_TIMESTAMP + WHERE run_id = $1 AND cohort_id = $2 AND status = 'sweeping' AND claim_token = $3 + `, + [run.run_id, run.cohort_id, run.claim_token], + 'abandonThresholdlessCohortMembershipSweep' + ) + sweepsAbandoned.inc({ reason: 'no_threshold' }) + return null } + // The threshold string was parsed to millisecond precision on the way out of Postgres, so + // it can sit up to 999 microseconds below the stored minimum. Truncation only ever lowers + // it, which makes the sweep under-delete, never over-delete; the next reconcile deletes + // the sliver it leaves behind. let deleted = 0 let lastBatch = 0 + // Keyset cursor over the (team_id, cohort_id, version) index order. Deleted index entries + // stay visible to a scan until VACUUM, so restarting each page from the range start would + // walk past every previously-deleted entry again. + let cursorVersion = '-infinity' + let cursorId = '0' do { // The predicate is repeated in the DELETE on purpose: under READ COMMITTED the DELETE // re-checks it against the newest row version, so a row a concurrent upsert just - // refreshed past the threshold is left alone instead of deleted anyway. - const result = await this.postgres.query( + // refreshed past the threshold is left alone instead of deleted anyway. A refreshed + // row also leaves the cursor range entirely, so skipping past it loses nothing. + const result = await this.postgres.query<{ version: string; id: string }>( PostgresUse.BEHAVIORAL_COHORTS_RW, ` WITH candidates AS ( SELECT id FROM cohort_membership WHERE team_id = $1 AND cohort_id = $2 - AND (version IS NULL OR version < $3::timestamp) + AND version < $3::timestamp + AND (version, id) > ($5::timestamp, $6::bigint) + ORDER BY version, id LIMIT $4 FOR UPDATE SKIP LOCKED ) DELETE FROM cohort_membership m USING candidates c - WHERE m.id = c.id AND (m.version IS NULL OR m.version < $3::timestamp) + WHERE m.id = c.id AND m.version < $3::timestamp + RETURNING to_char(m.version, 'YYYY-MM-DD HH24:MI:SS.US') AS version, m.id `, - [run.team_id, run.cohort_id, run.threshold, this.config.COHORT_MEMBERSHIP_SWEEP_BATCH_SIZE], + [ + run.team_id, + run.cohort_id, + run.threshold, + this.config.COHORT_MEMBERSHIP_SWEEP_BATCH_SIZE, + cursorVersion, + cursorId, + ], 'sweepCohortMembership' ) lastBatch = result.rowCount ?? 0 deleted += lastBatch + for (const row of result.rows) { + if ( + row.version > cursorVersion || + (row.version === cursorVersion && BigInt(row.id) > BigInt(cursorId)) + ) { + cursorVersion = row.version + cursorId = row.id + } + } + // Without this the claim goes stale mid-run and every other pod reclaims the same - // cohort, so N pods page over the same rows. + // cohort, so N pods page over the same rows. Zero rows updated means the claim is no + // longer ours (reclaimed after a stall, or abandoned); the new owner finishes the run. if (lastBatch > 0) { - await this.postgres.query( + const heartbeat = await this.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, `UPDATE cohort_membership_sweeps SET claimed_at = CURRENT_TIMESTAMP - WHERE run_id = $1 AND cohort_id = $2 AND status = 'sweeping'`, - [run.run_id, run.cohort_id], + WHERE run_id = $1 AND cohort_id = $2 AND status = 'sweeping' AND claim_token = $3`, + [run.run_id, run.cohort_id, run.claim_token], 'heartbeatCohortMembershipSweep' ) - } - } while (lastBatch > 0) - // A zero-delete iteration means either nothing is left or SKIP LOCKED skipped everything. - // Only the first is finished; calling the second `swept` would retire the run with stale - // rows still in place, and nothing revisits a swept run. - const remaining = await this.postgres.query( - PostgresUse.BEHAVIORAL_COHORTS_RW, - `SELECT 1 FROM cohort_membership - WHERE team_id = $1 AND cohort_id = $2 AND (version IS NULL OR version < $3::timestamp) - LIMIT 1`, - [run.team_id, run.cohort_id, run.threshold], - 'checkCohortMembershipSweepComplete' - ) - const finished = remaining.rows.length === 0 + if ((heartbeat.rowCount ?? 0) === 0) { + logger.warn('Lost the sweep claim mid-run, yielding to the new owner', { + run_id: run.run_id, + cohort_id: run.cohort_id, + rows: deleted, + }) + rowsSwept.inc(deleted) + return deleted + } + } + } while (lastBatch > 0 && !this.stopping) + + let finished = false + if (this.stopping && lastBatch > 0) { + // Interrupted by shutdown; the run goes back to 'ready' and resumes on the next claim. + finished = false + } else { + // A zero-delete iteration means either nothing is left or SKIP LOCKED skipped + // everything. Only the first is finished; calling the second `swept` would retire the + // run with stale rows still in place, and nothing revisits a swept run. + const remaining = await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + `SELECT 1 FROM cohort_membership + WHERE team_id = $1 AND cohort_id = $2 AND version < $3::timestamp + LIMIT 1`, + [run.team_id, run.cohort_id, run.threshold], + 'checkCohortMembershipSweepComplete' + ) + finished = remaining.rows.length === 0 + } - // Guarded on 'sweeping' so a straggler cannot overwrite the pod that reclaimed the run, nor - // resurrect one that was abandoned in between. + // The claim token fences stragglers: a pod whose claim timed out and was reclaimed no + // longer matches, so it cannot overwrite the new owner's state nor resurrect a run that + // was abandoned in between. swept_rows accumulates across partial passes. await this.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, ` UPDATE cohort_membership_sweeps - SET status = $1, swept_rows = $2, claimed_at = NULL, updated_at = CURRENT_TIMESTAMP - WHERE run_id = $3 AND cohort_id = $4 AND status = 'sweeping' + SET status = $1, swept_rows = COALESCE(swept_rows, 0) + $2, + claimed_at = NULL, claim_token = NULL, updated_at = CURRENT_TIMESTAMP + WHERE run_id = $3 AND cohort_id = $4 AND status = 'sweeping' AND claim_token = $5 `, - [finished ? 'swept' : 'ready', deleted, run.run_id, run.cohort_id], + [finished ? 'swept' : 'ready', deleted, run.run_id, run.cohort_id, run.claim_token], 'finishCohortMembershipSweep' ) @@ -449,29 +583,55 @@ export class CohortMembershipSweeper { } /** - * A reconcile superseded by a cohort edit emits no marker, so its set can never complete. The - * row is only table hygiene at that point. + * Two distinct give-up paths, kept apart because only one is routine. A reconcile superseded + * by a cohort edit emits no marker, so its `collecting` row can never complete; that is + * expected traffic. A `ready` run is fully proven and only waiting on consumer lag, and a + * `sweeping` run only ever fails transiently, so both get the marker-retention horizon rather + * than the short timeout: abandoning them sooner would discard runs that would have swept. + * + * These filters and the GC below scan without an index; the ledger is bounded by the GC to a + * month of reconcile runs, so a sequential scan stays trivial. */ private async abandonStuckRuns(): Promise { - const result = await this.postgres.query( + const unmarked = await this.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, ` UPDATE cohort_membership_sweeps SET status = 'abandoned', updated_at = CURRENT_TIMESTAMP - WHERE status IN ('collecting', 'ready') + WHERE status = 'collecting' AND updated_at < CURRENT_TIMESTAMP - $1::interval `, [`${this.config.COHORT_MEMBERSHIP_SWEEP_ABANDON_AFTER_DAYS} days`], - 'abandonCohortMembershipSweeps' + 'abandonUnmarkedCohortMembershipSweeps' ) - const abandoned = result.rowCount ?? 0 - if (abandoned > 0) { - sweepsAbandoned.inc(abandoned) - logger.warn('Abandoned reconcile runs that never completed their marker set', { count: abandoned }) + const unmarkedCount = unmarked.rowCount ?? 0 + if (unmarkedCount > 0) { + sweepsAbandoned.inc({ reason: 'markers_incomplete' }, unmarkedCount) + logger.warn('Abandoned reconcile runs that never completed their marker set', { count: unmarkedCount }) } - return abandoned + const unswept = await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + UPDATE cohort_membership_sweeps + SET status = 'abandoned', claim_token = NULL, updated_at = CURRENT_TIMESTAMP + WHERE status IN ('ready', 'sweeping') + AND created_at < CURRENT_TIMESTAMP - INTERVAL '${MARKER_RETENTION_DAYS} days' + `, + undefined, + 'abandonUnsweptCohortMembershipSweeps' + ) + + const unsweptCount = unswept.rowCount ?? 0 + if (unsweptCount > 0) { + sweepsAbandoned.inc({ reason: 'never_swept' }, unsweptCount) + logger.warn('Abandoned fully-marked reconcile runs whose sweep never completed', { + count: unsweptCount, + }) + } + + return unmarkedCount + unsweptCount } /** Matches the marker topic's retention, past which a run's markers cannot be redelivered. */ @@ -481,18 +641,30 @@ export class CohortMembershipSweeper { ` DELETE FROM cohort_membership_sweeps WHERE status IN ('swept', 'abandoned') - AND updated_at < CURRENT_TIMESTAMP - INTERVAL '30 days' + AND updated_at < CURRENT_TIMESTAMP - INTERVAL '${MARKER_RETENTION_DAYS} days' `, undefined, 'gcCohortMembershipSweeps' ) + + // Progress rows keyed to a broker cluster the consumer left stop updating and would + // otherwise sit forever. + await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + DELETE FROM cohort_membership_consumer_progress + WHERE updated_at < CURRENT_TIMESTAMP - INTERVAL '${MARKER_RETENTION_DAYS} days' + `, + undefined, + 'gcCohortMembershipProgress' + ) } private async recordGateWait(): Promise { const result = await this.postgres.query<{ wait_seconds: string | null }>( PostgresUse.BEHAVIORAL_COHORTS_RW, ` - SELECT EXTRACT(EPOCH FROM (CURRENT_TIMESTAMP - MIN(updated_at)))::text AS wait_seconds + SELECT EXTRACT(EPOCH FROM (CURRENT_TIMESTAMP - MIN(ready_at)))::text AS wait_seconds FROM cohort_membership_sweeps WHERE status = 'ready' `, undefined, @@ -515,6 +687,9 @@ export class CohortMembershipSweeper { } public async stop(): Promise { + // `stopping` breaks the paging loops inside a running cycle; without it, shutdown waits + // for up to MAX_RUNS_PER_TICK full cohort sweeps to page to completion. + this.stopping = true this.running = false this.sleepResolve?.() await this.loopPromise From 763bf2bae980df8a7e1c6de82e951bdf8ad1823b Mon Sep 17 00:00:00 2001 From: Matheus Veras Batista Date: Thu, 27 Aug 2026 09:28:46 -0300 Subject: [PATCH 3/9] chore(cohorts): document US-only scoping of the sweep flag and marker topic --- nodejs/src/cdp/config.ts | 7 +++++++ nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts | 5 ++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/nodejs/src/cdp/config.ts b/nodejs/src/cdp/config.ts index 363ac8d1a747..756f6b578b25 100644 --- a/nodejs/src/cdp/config.ts +++ b/nodejs/src/cdp/config.ts @@ -355,6 +355,13 @@ export function getDefaultCdpConfig(): CdpConfig { // Stays off until the processor writes membership to the live topic and the whole fleet // persists row versions. Sweeping before either would delete rows nothing re-asserts. + // + // The cohort stream pipeline (cohort-stream-processor and its topics, including the + // marker topic) is deployed US-only, like the rest of the cohort streaming system. This + // flag follows that scoping: environments without the pipeline keep it off, and the + // marker topic must never be created just to let the flag turn on — with no processor + // producing markers there, the sweeper would idle forever. start() fails fast when the + // flag is set and the marker topic does not exist. COHORT_MEMBERSHIP_SWEEP_ENABLED: false, COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS: 60000, COHORT_MEMBERSHIP_SWEEP_BATCH_SIZE: 1000, diff --git a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts index ea6121adfd17..80293dd6d5f2 100644 --- a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts +++ b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts @@ -536,7 +536,10 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase Date: Thu, 27 Aug 2026 10:30:41 -0300 Subject: [PATCH 4/9] fix(cohorts): apply versionless changes and contain promote failures --- .../cdp-cohort-membership.consumer.ts | 22 +++++++++++++++---- .../cohort-membership/sweeper.service.ts | 13 ++++++++++- 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts index 80293dd6d5f2..a13a759d077f 100644 --- a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts +++ b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts @@ -1,4 +1,5 @@ import { Message } from 'node-rdkafka' +import { Counter } from 'prom-client' import { z } from 'zod' import { KAFKA_COHORT_MEMBERSHIP_CHANGED, KAFKA_COHORT_RECONCILE_MARKERS } from '~/common/config/kafka-topics' @@ -42,6 +43,11 @@ function membershipClusterId(): string { return String(getKafkaConfigFromEnv('CONSUMER')['metadata.broker.list'] ?? 'kafka:9092') } +const offFormatVersions = new Counter({ + name: 'cdp_cohort_membership_off_format_versions', + help: 'Membership changes whose version stamp broke the producer contract and was degraded to versionless', +}) + // `origin` is a loose string, not an enum: nothing branches on it, and rejecting a value the // processor added later would throw the whole batch and crash-loop on redelivery. const CohortMembershipChangeSchema = z.object({ @@ -306,7 +312,12 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase= version` evaluates to NULL against + * a versioned row and the change is silently dropped, leaving the member stuck. */ private async upsertMembershipRows(tx: TransactionClient, changes: CohortMembershipChange[]): Promise { const values: any[] = [] @@ -338,8 +349,10 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase= cohort_membership.version + version = COALESCE(EXCLUDED.version, cohort_membership.version) + WHERE cohort_membership.version IS NULL + OR EXCLUDED.version IS NULL + OR EXCLUDED.version >= cohort_membership.version `, values, 'batchUpsertCohortMembership' @@ -479,11 +492,12 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { const result: SweepCycleResult = { swept: 0, rowsDeleted: 0, blocked: 0, abandoned: 0 } - await this.promoteMarkedRuns() + // Promotion talks to Kafka and its partial-metadata guard can fail permanently (a topic + // recreated with fewer partitions). Letting the throw abort the cycle would also skip the + // hygiene steps below, including the progress GC that eventually clears exactly that + // cause, so a promotion failure is contained to its own outcome. + try { + await this.promoteMarkedRuns() + } catch (error) { + sweepCycles.inc({ status: 'promote_error' }) + logger.error('Failed to promote fully-marked reconcile runs', { error: String(error) }) + } + await this.kafka.refreshConsumerProgress() const progress = await this.readConsumerProgress() @@ -680,6 +690,7 @@ export class CohortMembershipSweeper { } this.running = true + this.stopping = false this.loopPromise = this.sweepLoop() logger.info('CohortMembershipSweeper started', { intervalMs: this.config.COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS, From 1ba95e5d2da36528b9e7b08d8f6094a9bde65906 Mon Sep 17 00:00:00 2001 From: Matheus Veras Batista Date: Thu, 27 Aug 2026 10:43:21 -0300 Subject: [PATCH 5/9] chore(cohorts): run the sweeper suite in the serial jest lane --- .../{sweeper.service.test.ts => sweeper.service.serial.test.ts} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename nodejs/src/cdp/services/cohort-membership/{sweeper.service.test.ts => sweeper.service.serial.test.ts} (100%) diff --git a/nodejs/src/cdp/services/cohort-membership/sweeper.service.test.ts b/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts similarity index 100% rename from nodejs/src/cdp/services/cohort-membership/sweeper.service.test.ts rename to nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts From d88bd705e300af7b97a6765bf524344bc9f19db3 Mon Sep 17 00:00:00 2001 From: Matheus Veras Batista Date: Thu, 27 Aug 2026 10:56:09 -0300 Subject: [PATCH 6/9] fix(cohorts): mock the consumer interface read methods --- nodejs/tests/helpers/mocks/consumer.mock.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/nodejs/tests/helpers/mocks/consumer.mock.ts b/nodejs/tests/helpers/mocks/consumer.mock.ts index 8bb3ff70f257..4a4dba9e30a5 100644 --- a/nodejs/tests/helpers/mocks/consumer.mock.ts +++ b/nodejs/tests/helpers/mocks/consumer.mock.ts @@ -5,6 +5,9 @@ const mockKafkaConsumer = (): jest.Mocked => ({ disconnect: jest.fn().mockResolvedValue(undefined), isHealthy: jest.fn(), offsetsStore: jest.fn(), + queryWatermarkOffsets: jest.fn().mockResolvedValue([0, 0]), + committedOffsets: jest.fn().mockResolvedValue([]), + getPartitionsForTopic: jest.fn().mockResolvedValue([]), }) jest.mock('~/common/kafka/consumer', () => ({ From ba202153d3b0be743e8a4d7bff5700a60ad06d08 Mon Sep 17 00:00:00 2001 From: Matheus Veras Batista Date: Thu, 27 Aug 2026 14:56:43 -0300 Subject: [PATCH 7/9] fix(cohorts): follow the sentinel version and stamp watermark provenance --- .../cdp-cohort-membership.consumer.test.ts | 9 ++-- .../cdp-cohort-membership.consumer.ts | 35 +++++++++++----- .../sweeper.service.serial.test.ts | 41 +++++++++++-------- .../cohort-membership/sweeper.service.ts | 35 +++++++++++++++- 4 files changed, 88 insertions(+), 32 deletions(-) diff --git a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts index 4be50283fe87..f44bb89217ba 100644 --- a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts +++ b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts @@ -334,8 +334,8 @@ describe('CdpCohortMembershipConsumer', () => { ['a newer version is applied', VERSION, NEWER_VERSION, false, NEWER_VERSION], ['any version beats a versionless row', undefined, OLDER_VERSION, false, OLDER_VERSION], // A change that lost its version (sanitized off-format stamp) is still a membership - // transition: it must apply while leaving the row's stamp in place, not silently - // no-op against the NULL comparison. + // transition: it must apply while leaving the row's stamp in place, instead of losing + // to every stamped row and no-opping. ['a versionless change applies without moving the version', VERSION, undefined, false, VERSION], ])( 'should apply last-writer-wins on replay: %s', @@ -491,11 +491,12 @@ describe('CdpCohortMembershipConsumer', () => { const membership = await hub.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, - 'SELECT in_cohort, version FROM cohort_membership WHERE team_id = $1 AND person_id = $2', + `SELECT in_cohort, version = '-infinity'::timestamp AS unversioned + FROM cohort_membership WHERE team_id = $1 AND person_id = $2`, [teamId, personId1], 'testQuery' ) - expect(membership.rows).toEqual([{ in_cohort: true, version: null }]) + expect(membership.rows).toEqual([{ in_cohort: true, unversioned: true }]) const sweeps = await hub.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, diff --git a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts index a13a759d077f..2f078ad18dd7 100644 --- a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts +++ b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts @@ -29,9 +29,17 @@ export type CdpCohortMembershipConsumerConfig = CdpConsumerBaseConfig & | 'COHORT_MEMBERSHIP_SWEEP_ABANDON_AFTER_DAYS' > -// Markers from runs that predate the sweeper are useless — the next reconcile heals anyway. +// Markers from runs that predate the sweeper are useless, because the next reconcile heals anyway. const startAtLatest = { ['auto.offset.reset' as keyof RdKafkaConsumerConfig]: 'latest' as never } +/** + * The stamp for a membership row the pipeline never versioned. Must stay equal to the column + * default in `20260826000001_add_version_to_cohort_membership.sql`: every row written before that + * migration reads as this value, and the sweep deletes below a threshold, so a sentinel that does + * not sort below every producer stamp would leave stale rows behind. + */ +const VERSIONLESS = '-infinity' + /** * The broker list the membership consumer resolves, mirroring the consumer's own config * resolution. Progress offsets are only comparable against watermarks from the same cluster, so @@ -165,6 +173,7 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase this.captureMembershipWatermarks(), refreshConsumerProgress: () => this.refreshConsumerProgressFromCommittedOffsets(), }) @@ -312,12 +321,13 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase= version` evaluates to NULL against - * a versioned row and the change is silently dropped, leaving the member stuck. + * transition: it applies unconditionally (the pre-version behavior) while GREATEST keeps the + * row's existing stamp. Without both, an unordered change would lose to every stamped row and + * be silently dropped, leaving the member stuck. */ private async upsertMembershipRows(tx: TransactionClient, changes: CohortMembershipChange[]): Promise { const values: any[] = [] @@ -335,7 +345,7 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase= cohort_membership.version `, values, @@ -359,7 +368,13 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { const values: any[] = [] const placeholders: string[] = [] diff --git a/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts b/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts index e3ddcd615a89..a519e05345bd 100644 --- a/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts +++ b/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts @@ -22,6 +22,8 @@ describe('CohortMembershipSweeper', () => { const THRESHOLD = '2026-05-26 12:00:00.000000' const BELOW_THRESHOLD = '2026-05-26 11:00:00.000000' const ABOVE_THRESHOLD = '2026-05-26 13:00:00.000000' + /** What the column defaults to, so it is what every row written before the feature reads as. */ + const UNVERSIONED = '-infinity' const config: CohortMembershipSweepConfig = { COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS: 60000, @@ -32,6 +34,7 @@ describe('CohortMembershipSweeper', () => { } const CLUSTER = 'test-cluster' + const TOPIC = 'cohort_membership_changed' let postgres: PostgresRouter let sweeper: CohortMembershipSweeper @@ -45,6 +48,7 @@ describe('CohortMembershipSweeper', () => { runId = new UUIDT().toString() sweeper = new CohortMembershipSweeper(config, postgres, { cluster: CLUSTER, + topic: TOPIC, captureMembershipWatermarks: () => Promise.resolve(watermarks), refreshConsumerProgress: () => Promise.resolve(), }) @@ -83,7 +87,7 @@ describe('CohortMembershipSweeper', () => { const readSweep = async (): Promise> => { const rows = await query>( - `SELECT marker_bits, status, swept_rows, membership_hwms, + `SELECT marker_bits, status, swept_rows, membership_hwms, membership_cluster, membership_topic, to_char(min_marker_version, 'YYYY-MM-DD HH24:MI:SS.US') AS min_marker_version FROM cohort_membership_sweeps WHERE run_id = $1 AND cohort_id = $2`, [runId, COHORT_ID] @@ -91,7 +95,7 @@ describe('CohortMembershipSweeper', () => { return rows[0] } - const insertMembership = async (teamId: number, cohortId: number, version: string | null): Promise => { + const insertMembership = async (teamId: number, cohortId: number, version: string): Promise => { const personId = new UUIDT().toString() await query( `INSERT INTO cohort_membership (team_id, cohort_id, person_id, in_cohort, version) @@ -111,9 +115,10 @@ describe('CohortMembershipSweeper', () => { await query( `INSERT INTO cohort_membership_sweeps (run_id, cohort_id, team_id, marker_bits, min_marker_version, min_snapshot_version, - membership_hwms, status, ready_at) - VALUES ($1, $2, $3, -1, $4::timestamp, $5::timestamp, $6::jsonb, 'ready', CURRENT_TIMESTAMP)`, - [runId, COHORT_ID, TEAM_ID, MARKER_VERSION, THRESHOLD, JSON.stringify(watermarks)] + membership_hwms, membership_cluster, membership_topic, status, ready_at) + VALUES ($1, $2, $3, -1, $4::timestamp, $5::timestamp, $6::jsonb, $7, $8, 'ready', + CURRENT_TIMESTAMP)`, + [runId, COHORT_ID, TEAM_ID, MARKER_VERSION, THRESHOLD, JSON.stringify(watermarks), CLUSTER, TOPIC] ) } @@ -179,6 +184,8 @@ describe('CohortMembershipSweeper', () => { marker_bits: '-1', status: 'ready', membership_hwms: { 0: 10, 3: 4 }, + membership_cluster: CLUSTER, + membership_topic: TOPIC, min_marker_version: MARKER_VERSION, }) }) @@ -210,22 +217,20 @@ describe('CohortMembershipSweeper', () => { const stale1 = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) const stale2 = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) const stale3 = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) - // Only versioned rows are sweepable: a NULL was written either before the feature or by a - // rolled-back image asserting current membership, and the two cannot be told apart here. - const neverVersioned = await insertMembership(TEAM_ID, COHORT_ID, null) + // A row written before the feature sorts below every producer stamp, so the run that did + // not re-assert it deletes it. Those rows are most of what the first sweep has to remove. + const preFeature = await insertMembership(TEAM_ID, COHORT_ID, UNVERSIONED) const asserted = await insertMembership(TEAM_ID, COHORT_ID, THRESHOLD) const liveChange = await insertMembership(TEAM_ID, COHORT_ID, ABOVE_THRESHOLD) const otherCohort = await insertMembership(TEAM_ID, COHORT_ID + 1, BELOW_THRESHOLD) const otherTeam = await insertMembership(TEAM_ID + 1, COHORT_ID, BELOW_THRESHOLD) - expect(await sweeper.runOnce()).toMatchObject({ swept: 1, rowsDeleted: 3 }) + expect(await sweeper.runOnce()).toMatchObject({ swept: 1, rowsDeleted: 4 }) const surviving = await survivingPersonIds() - expect(surviving).toEqual( - [neverVersioned, asserted, liveChange, otherCohort, otherTeam].sort((a, b) => a.localeCompare(b)) - ) - expect(surviving).not.toEqual(expect.arrayContaining([stale1, stale2, stale3])) - expect(await readSweep()).toMatchObject({ status: 'swept', swept_rows: '3' }) + expect(surviving).toEqual([asserted, liveChange, otherCohort, otherTeam].sort((a, b) => a.localeCompare(b))) + expect(surviving).not.toEqual(expect.arrayContaining([stale1, stale2, stale3, preFeature])) + expect(await readSweep()).toMatchObject({ status: 'swept', swept_rows: '4' }) // A second cycle must not re-sweep a finished run. expect(await sweeper.runOnce()).toMatchObject({ swept: 0, rowsDeleted: 0 }) @@ -234,6 +239,10 @@ describe('CohortMembershipSweeper', () => { it.each([ ['it never observed the snapshot it is supposed to have asserted', 'min_snapshot_version = NULL'], ['its watermark capture came back empty', `membership_hwms = '{}'::jsonb`], + // Offsets from one feed prove nothing about another, so a run that outlived a broker + // repoint or a topic change must not sweep on the watermarks it captured before it. + ['its watermarks came from another cluster', `membership_cluster = 'moved-cluster'`], + ['its watermarks came from another topic', `membership_topic = 'moved_topic'`], ])('should refuse to sweep a run when %s', async (_label, damage) => { watermarks = { 0: 10 } await readyRun() @@ -241,10 +250,10 @@ describe('CohortMembershipSweeper', () => { await query(`UPDATE cohort_membership_sweeps SET ${damage} WHERE run_id = $1`, [runId]) const live = await insertMembership(TEAM_ID, COHORT_ID, ABOVE_THRESHOLD) - const unversioned = await insertMembership(TEAM_ID, COHORT_ID, null) + const stale = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) expect(await sweeper.runOnce()).toMatchObject({ swept: 0, rowsDeleted: 0 }) - expect(await survivingPersonIds()).toEqual([live, unversioned].sort((a, b) => a.localeCompare(b))) + expect(await survivingPersonIds()).toEqual([live, stale].sort((a, b) => a.localeCompare(b))) }) it('should reclaim a sweep whose pod died mid-run', async () => { diff --git a/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts b/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts index c859a133741f..49f7c9101c9e 100644 --- a/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts +++ b/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts @@ -76,6 +76,8 @@ type ClaimableRun = { cohort_id: string team_id: string membership_hwms: MembershipWatermarks | null + membership_cluster: string | null + membership_topic: string | null threshold: string | null claim_token: string | null } @@ -140,6 +142,8 @@ export class CohortMembershipSweeper { private kafka: { /** Broker identity the progress rows are keyed by; offsets are only comparable within one cluster. */ cluster: string + /** Membership topic the watermarks and the progress offsets both describe. */ + topic: string captureMembershipWatermarks: () => Promise refreshConsumerProgress: () => Promise } @@ -252,11 +256,12 @@ export class CohortMembershipSweeper { ` UPDATE cohort_membership_sweeps SET status = 'ready', membership_hwms = $1::jsonb, + membership_cluster = $4, membership_topic = $5, ready_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP WHERE run_id = $2 AND cohort_id = $3 AND status = 'collecting' AND marker_bits = ${ALL_MARKERS} `, - [JSON.stringify(watermarks), run.run_id, run.cohort_id], + [JSON.stringify(watermarks), run.run_id, run.cohort_id, this.kafka.cluster, this.kafka.topic], 'readyCohortMembershipSweep' ) @@ -340,6 +345,7 @@ export class CohortMembershipSweeper { PostgresUse.BEHAVIORAL_COHORTS_RW, ` SELECT run_id, cohort_id, team_id, membership_hwms, + membership_cluster, membership_topic, LEAST(min_snapshot_version, min_marker_version) AS threshold, claim_token FROM cohort_membership_sweeps WHERE min_snapshot_version IS NOT NULL @@ -385,6 +391,23 @@ export class CohortMembershipSweeper { return false } + // An offset only means something against the feed it was read from. If the consumer's + // brokers were repointed or the processor's output topic changed after this run captured + // its watermarks, the two sides of the comparison describe different feeds and the gate + // proves nothing. Refusing holds the run until it ages into abandonment, and the next + // reconcile produces a run stamped with the current feed. + if (run.membership_cluster !== this.kafka.cluster || run.membership_topic !== this.kafka.topic) { + logger.warn('Refusing to sweep a run whose watermarks came from another feed', { + run_id: run.run_id, + cohort_id: run.cohort_id, + captured_cluster: run.membership_cluster, + captured_topic: run.membership_topic, + current_cluster: this.kafka.cluster, + current_topic: this.kafka.topic, + }) + return false + } + const blocked: number[] = [] for (const [partition, highWatermark] of watermarks) { @@ -425,6 +448,7 @@ export class CohortMembershipSweeper { AND (status = 'ready' OR (status = 'sweeping' AND claimed_at < CURRENT_TIMESTAMP - $3::interval)) RETURNING run_id, cohort_id, team_id, membership_hwms, + membership_cluster, membership_topic, LEAST(min_snapshot_version, min_marker_version) AS threshold, claim_token `, [ @@ -473,6 +497,11 @@ export class CohortMembershipSweeper { // Keyset cursor over the (team_id, cohort_id, version) index order. Deleted index entries // stay visible to a scan until VACUUM, so restarting each page from the range start would // walk past every previously-deleted entry again. + // + // The cursor starts at the same sentinel the column defaults to, and `>` on the row + // constructor still advances through those rows by id. `to_char` returns NULL for an + // infinite timestamp, so the RETURNING clause restores the sentinel text: without it the + // cursor would take a NULL and the next page would match nothing. let cursorVersion = '-infinity' let cursorId = '0' @@ -496,7 +525,9 @@ export class CohortMembershipSweeper { DELETE FROM cohort_membership m USING candidates c WHERE m.id = c.id AND m.version < $3::timestamp - RETURNING to_char(m.version, 'YYYY-MM-DD HH24:MI:SS.US') AS version, m.id + RETURNING COALESCE( + to_char(m.version, 'YYYY-MM-DD HH24:MI:SS.US'), '-infinity' + ) AS version, m.id `, [ run.team_id, From de4dc599a58507d0140d19e32711a60c3ea5ffc8 Mon Sep 17 00:00:00 2001 From: Matheus Veras Batista Date: Thu, 27 Aug 2026 15:54:02 -0300 Subject: [PATCH 8/9] fix(cohorts): apply review feedback on the membership sweeper --- nodejs/src/cdp/config.ts | 9 +- .../cdp-cohort-membership.consumer.test.ts | 66 ++++- .../cdp-cohort-membership.consumer.ts | 216 +++++++------- ...p-hogflow-subscription-matcher.consumer.ts | 4 +- nodejs/src/cdp/dev/cohort-sweep-e2e.test.ts | 7 +- .../sweeper.service.serial.test.ts | 94 +++++- .../cohort-membership/sweeper.service.ts | 280 +++++++++++------- .../common/kafka/consumer/consumer-v1.test.ts | 16 + .../src/common/kafka/consumer/consumer-v1.ts | 11 +- nodejs/src/common/kafka/consumer/index.ts | 8 + 10 files changed, 483 insertions(+), 228 deletions(-) diff --git a/nodejs/src/cdp/config.ts b/nodejs/src/cdp/config.ts index 756f6b578b25..8423d0a4ba28 100644 --- a/nodejs/src/cdp/config.ts +++ b/nodejs/src/cdp/config.ts @@ -186,7 +186,11 @@ export type CdpConfig = ClickhouseConfig & { CYCLOTRON_NODE_RESCHEDULE_MAX_CHUNKS_PER_CALL: number CYCLOTRON_NODE_RESCHEDULE_CHUNK_SLEEP_MS: number - // Mark-and-sweep deletion of cohort_membership rows a completed reconcile run did not re-assert + // Mark-and-sweep deletion of cohort_membership rows a completed reconcile run did not re-assert. + // Two flags so a rolling restart never mixes sweeping pods with pods that write no versions: + // version writes turn on first and off last, sweeping turns on last and off first. A pod that + // writes without versions preserves a row's old stamp, which a sweeping pod then reads as stale. + COHORT_MEMBERSHIP_VERSION_WRITES_ENABLED: boolean COHORT_MEMBERSHIP_SWEEP_ENABLED: boolean COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS: number COHORT_MEMBERSHIP_SWEEP_BATCH_SIZE: number @@ -353,6 +357,9 @@ export function getDefaultCdpConfig(): CdpConfig { CYCLOTRON_NODE_RESCHEDULE_MAX_CHUNKS_PER_CALL: 20, CYCLOTRON_NODE_RESCHEDULE_CHUNK_SLEEP_MS: 100, + // Stays off until the schema migrations have applied; a pod reaching an unmigrated + // database must degrade to the pre-sweep write shape, not crash on a missing column. + COHORT_MEMBERSHIP_VERSION_WRITES_ENABLED: false, // Stays off until the processor writes membership to the live topic and the whole fleet // persists row versions. Sweeping before either would delete rows nothing re-asserts. // diff --git a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts index f44bb89217ba..00edf97fef8e 100644 --- a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts +++ b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts @@ -26,7 +26,10 @@ describe('CdpCohortMembershipConsumer', () => { hub = await createHub() const deps = createCdpConsumerDeps(hub) consumer = new CdpCohortMembershipConsumer(hub, deps) - sweepConsumer = new CdpCohortMembershipConsumer({ ...hub, COHORT_MEMBERSHIP_SWEEP_ENABLED: true }, deps) + sweepConsumer = new CdpCohortMembershipConsumer( + { ...hub, COHORT_MEMBERSHIP_VERSION_WRITES_ENABLED: true, COHORT_MEMBERSHIP_SWEEP_ENABLED: true }, + deps + ) teamId = Number.parseInt(new UUIDT().toString().replaceAll('-', '').slice(-7), 16) // Progress rows key on (cluster, partition) and nothing truncates between tests, so a // unique cluster isolates each test from parallel workers and from prior runs. @@ -205,6 +208,14 @@ describe('CdpCohortMembershipConsumer', () => { { last_updated: '2026-05-26 12:00:00.000000' }, true, ], + // A versionless change applies unconditionally in SQL, so it also has to win the + // in-batch pick, even against a versioned entry earlier in Kafka order. + [ + 'keeping a later versionless change over an earlier versioned one', + { last_updated: '2026-05-26 13:00:00.000000' }, + {}, + false, + ], ])( 'should deduplicate batch entries for the same (team_id, cohort_id, person_id), %s', async (_label, firstTags, secondTags, expectedInCohort) => { @@ -442,6 +453,33 @@ describe('CdpCohortMembershipConsumer', () => { min_snapshot_version: OLDER_VERSION, }, ]) + + // A reconcile assertion that lost its version still applies to its row while keeping + // the row's old stamp, so it has to collapse the run's minimum to the sentinel + // (rendered as NULL by to_char): otherwise the run's own sweep would delete the + // person it just asserted. + await sweepConsumer['handleBatch']([ + createKafkaMessage( + createCohortMembershipEvent({ + person_id: personId1, + cohort_id: 456, + team_id: teamId, + origin: 'reconcile', + run_id: runId, + }), + { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: 0 } + ), + ]) + + expect(await readSweep()).toEqual([ + { + team_id: String(teamId), + marker_bits: '0', + status: 'collecting', + snapshot_rows: '4', + min_snapshot_version: null, + }, + ]) }) it('should advance consumer progress to the next offset without ever regressing it', async () => { @@ -471,9 +509,9 @@ describe('CdpCohortMembershipConsumer', () => { ]) }) - it('should keep the pre-sweep write shape while the sweep flag is off', async () => { - // The migrations are a prerequisite of the flag flip, not of the deploy, so a flag-off - // batch must not touch the version column or the sweep tables. + it('should keep the pre-sweep write shape while version writes are off', async () => { + // The migrations are a prerequisite of the version-writes flip, not of the deploy, so + // a flag-off batch must not touch the version column or the sweep tables. const runId = new UUIDT().toString() const message = createKafkaMessage( createCohortMembershipEvent({ @@ -537,12 +575,6 @@ describe('CdpCohortMembershipConsumer', () => { }) it('should refuse to capture watermarks from a partial partition list', async () => { - await sweepConsumer['upsertConsumerProgress'](PostgresUse.BEHAVIORAL_COHORTS_RW, [ - { partition: 0, nextOffset: 5 }, - { partition: 1, nextOffset: 7 }, - { partition: 2, nextOffset: 9 }, - ]) - const stubConsumer = { getPartitionsForTopic: jest.fn(), queryWatermarkOffsets: jest.fn((_topic: string, partition: number) => @@ -551,10 +583,22 @@ describe('CdpCohortMembershipConsumer', () => { } sweepConsumer['kafkaConsumer'] = stubConsumer as any + // With no progress rows for this cluster there is no floor to validate the metadata + // against, so any truncated list would pass as the full set: capture fails closed. + stubConsumer.getPartitionsForTopic.mockResolvedValue([{ id: 0 }, { id: 1 }]) + await expect(sweepConsumer['captureMembershipWatermarks']()).rejects.toThrow( + 'No recorded consumer progress' + ) + + await sweepConsumer['upsertConsumerProgress'](PostgresUse.BEHAVIORAL_COHORTS_RW, [ + { partition: 0, nextOffset: 5 }, + { partition: 1, nextOffset: 7 }, + { partition: 2, nextOffset: 9 }, + ]) + // Three partitions have durable consumer progress on this cluster, so a broker // answering with two is a partial list, and capturing it would let the gate pass while // the missing partition still holds unconsumed snapshot rows. - stubConsumer.getPartitionsForTopic.mockResolvedValue([{ id: 0 }, { id: 1 }]) await expect(sweepConsumer['captureMembershipWatermarks']()).rejects.toThrow('Partial partition metadata') stubConsumer.getPartitionsForTopic.mockResolvedValue([{ id: 0 }, { id: 1 }, { id: 2 }]) diff --git a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts index 2f078ad18dd7..e124f0b8a6fe 100644 --- a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts +++ b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts @@ -4,7 +4,7 @@ import { z } from 'zod' import { KAFKA_COHORT_MEMBERSHIP_CHANGED, KAFKA_COHORT_RECONCILE_MARKERS } from '~/common/config/kafka-topics' import { getKafkaConfigFromEnv } from '~/common/kafka/config' -import { KafkaConsumerInterface, RdKafkaConsumerConfig, createKafkaConsumer } from '~/common/kafka/consumer' +import { KafkaConsumerInterface, START_AT_LATEST, createKafkaConsumer } from '~/common/kafka/consumer' import { instrumentFn } from '~/common/tracing/tracing-utils' import { PostgresUse, TransactionClient } from '~/common/utils/db/postgres' import { parseJSON } from '~/common/utils/json-parse' @@ -22,6 +22,7 @@ import { CdpConsumerBase, CdpConsumerBaseConfig, CdpConsumerBaseDeps } from './c export type CdpCohortMembershipConsumerConfig = CdpConsumerBaseConfig & Pick< CdpConfig, + | 'COHORT_MEMBERSHIP_VERSION_WRITES_ENABLED' | 'COHORT_MEMBERSHIP_SWEEP_ENABLED' | 'COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS' | 'COHORT_MEMBERSHIP_SWEEP_BATCH_SIZE' @@ -29,9 +30,6 @@ export type CdpCohortMembershipConsumerConfig = CdpConsumerBaseConfig & | 'COHORT_MEMBERSHIP_SWEEP_ABANDON_AFTER_DAYS' > -// Markers from runs that predate the sweeper are useless, because the next reconcile heals anyway. -const startAtLatest = { ['auto.offset.reset' as keyof RdKafkaConsumerConfig]: 'latest' as never } - /** * The stamp for a membership row the pipeline never versioned. Must stay equal to the column * default in `20260826000001_add_version_to_cohort_membership.sql`: every row written before that @@ -107,7 +105,13 @@ function partitionProgressFromMessages(messages: Message[]): PartitionProgress[] * every version its run re-asserted, so the run's minimum is folded across batches in Postgres. * * Versions are compared as strings: the producer's format is fixed-width, so lexicographic order - * is chronological order. A row carrying no version bounds nothing and is left out. + * is chronological order. + * + * A reconcile change with no version still applies to its row while preserving the row's old + * stamp (see upsertMembershipRows), so skipping it here would leave the run's minimum above that + * stamp and the sweep would delete a person its own run just asserted. Folding the sentinel in + * collapses the run's threshold instead: the sweep deletes nothing, and the next reconcile + * retries with a clean run. * * The result is sorted by (run_id, cohort_id) because it feeds one multi-row upsert whose row * locks are held to commit: two pods writing the same runs in different orders would deadlock. @@ -121,13 +125,13 @@ function collectSnapshotMarks(changes: CohortMembershipChange[]): SnapshotMark[] } if (!change.last_updated) { - logger.warn('Reconcile membership change without a version, excluded from run marking', { + logger.warn('Reconcile membership change without a version, collapsing the run threshold', { run_id: change.run_id, cohort_id: change.cohort_id, team_id: change.team_id, }) - continue } + const version = change.last_updated ?? VERSIONLESS const key = `${change.run_id}:${change.cohort_id}` const existing = marks.get(key) @@ -137,15 +141,15 @@ function collectSnapshotMarks(changes: CohortMembershipChange[]): SnapshotMark[] runId: change.run_id, cohortId: change.cohort_id, teamId: change.team_id, - minVersion: change.last_updated, + minVersion: version, rows: 1, }) continue } existing.rows += 1 - if (change.last_updated < existing.minVersion) { - existing.minVersion = change.last_updated + if (version < existing.minVersion) { + existing.minVersion = version } } @@ -166,10 +170,21 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase() for (const change of changes) { const key = `${change.team_id}:${change.cohort_id}:${change.person_id}` const existing = deduped.get(key) - if (!existing || (change.last_updated ?? '') >= (existing.last_updated ?? '')) { + if (!existing || !change.last_updated || change.last_updated >= (existing.last_updated ?? '')) { deduped.set(key, change) } } - if (!this.config.COHORT_MEMBERSHIP_SWEEP_ENABLED) { + if (!this.config.COHORT_MEMBERSHIP_VERSION_WRITES_ENABLED) { if (deduped.size > 0) { await this.upsertMembershipRowsWithoutVersion(Array.from(deduped.values())) } @@ -330,31 +358,12 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { - const values: any[] = [] - const placeholders: string[] = [] - let paramIndex = 1 - - for (const change of changes) { - placeholders.push( - `($${paramIndex}, $${paramIndex + 1}, $${paramIndex + 2}, $${paramIndex + 3}, CURRENT_TIMESTAMP, $${ - paramIndex + 4 - }::timestamp)` - ) - values.push( - change.team_id, - change.cohort_id, - change.person_id, - change.status === 'entered', - change.last_updated ?? VERSIONLESS - ) - paramIndex += 5 - } - await this.deps.postgres.query( tx, ` INSERT INTO cohort_membership (team_id, cohort_id, person_id, in_cohort, last_updated, version) - VALUES ${placeholders.join(', ')} + SELECT unnest($1::bigint[]), unnest($2::bigint[]), unnest($3::uuid[]), + unnest($4::boolean[]), CURRENT_TIMESTAMP, unnest($5::timestamp[]) ON CONFLICT (team_id, cohort_id, person_id) DO UPDATE SET in_cohort = EXCLUDED.in_cohort, @@ -363,42 +372,45 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase= cohort_membership.version `, - values, + [ + changes.map((change) => change.team_id), + changes.map((change) => change.cohort_id), + changes.map((change) => change.person_id), + changes.map((change) => change.status === 'entered'), + changes.map((change) => change.last_updated ?? VERSIONLESS), + ], 'batchUpsertCohortMembership' ) } /** - * The flag-off shape: no version column, no transaction, valid against the pre-sweep schema. + * The version-writes-off shape: no version column, no transaction, valid against the + * pre-sweep schema. * - * Once the schema is migrated, rows this path creates take the column default and are therefore - * sweepable. Flip the flag fleet-wide rather than leaving pods mixed: a sweeping pod can delete - * a row a flag-off pod wrote after the reconcile snapshot, and only the next reconcile heals it. + * Once the schema is migrated, rows this path creates take the column default and are + * therefore sweepable, and a re-asserted row keeps whatever stamp it already carried. The + * rollout order is what keeps this path and a sweeping pod from ever coexisting in one fleet: + * version writes turn on first and off last, sweeping turns on last and off first. The + * constructor enforces the per-pod half of that contract. */ private async upsertMembershipRowsWithoutVersion(changes: CohortMembershipChange[]): Promise { - const values: any[] = [] - const placeholders: string[] = [] - let paramIndex = 1 - - for (const change of changes) { - placeholders.push( - `($${paramIndex}, $${paramIndex + 1}, $${paramIndex + 2}, $${paramIndex + 3}, CURRENT_TIMESTAMP)` - ) - values.push(change.team_id, change.cohort_id, change.person_id, change.status === 'entered') - paramIndex += 4 - } - await this.deps.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, ` INSERT INTO cohort_membership (team_id, cohort_id, person_id, in_cohort, last_updated) - VALUES ${placeholders.join(', ')} + SELECT unnest($1::bigint[]), unnest($2::bigint[]), unnest($3::uuid[]), + unnest($4::boolean[]), CURRENT_TIMESTAMP ON CONFLICT (team_id, cohort_id, person_id) DO UPDATE SET in_cohort = EXCLUDED.in_cohort, last_updated = CURRENT_TIMESTAMP `, - values, + [ + changes.map((change) => change.team_id), + changes.map((change) => change.cohort_id), + changes.map((change) => change.person_id), + changes.map((change) => change.status === 'entered'), + ], 'batchUpsertCohortMembership' ) } @@ -408,26 +420,13 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { - const values: any[] = [] - const placeholders: string[] = [] - let paramIndex = 1 - - for (const mark of marks) { - placeholders.push( - `($${paramIndex}, $${paramIndex + 1}, $${paramIndex + 2}, $${paramIndex + 3}::timestamp, $${ - paramIndex + 4 - })` - ) - values.push(mark.runId, mark.cohortId, mark.teamId, mark.minVersion, mark.rows) - paramIndex += 5 - } - await this.deps.postgres.query( tx, ` INSERT INTO cohort_membership_sweeps (run_id, cohort_id, team_id, min_snapshot_version, snapshot_rows) - VALUES ${placeholders.join(', ')} + SELECT unnest($1::uuid[]), unnest($2::bigint[]), unnest($3::bigint[]), + unnest($4::timestamp[]), unnest($5::bigint[]) ON CONFLICT (run_id, cohort_id) DO UPDATE SET min_snapshot_version = LEAST( @@ -437,7 +436,13 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase mark.runId), + marks.map((mark) => mark.cohortId), + marks.map((mark) => mark.teamId), + marks.map((mark) => mark.minVersion), + marks.map((mark) => mark.rows), + ], 'markCohortMembershipSnapshot' ) } @@ -450,21 +455,16 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { - const values: any[] = [] - const placeholders: string[] = [] - let paramIndex = 1 - - for (const entry of progress) { - placeholders.push(`($${paramIndex}, $${paramIndex + 1}, $${paramIndex + 2})`) - values.push(this.membershipCluster, entry.partition, entry.nextOffset) - paramIndex += 3 - } + // Row locks are held to commit, and the batch path and the sweep loop hand partitions in + // in different orders, so every writer has to take them in the same order. GREATEST makes + // the result order-independent, so sorting here is safe. + const sorted = [...progress].sort((a, b) => a.partition - b.partition) await this.deps.postgres.query( target, ` INSERT INTO cohort_membership_consumer_progress (cluster, partition, next_offset) - VALUES ${placeholders.join(', ')} + SELECT $1, unnest($2::int[]), unnest($3::bigint[]) ON CONFLICT (cluster, partition) DO UPDATE SET next_offset = GREATEST( @@ -473,7 +473,7 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase entry.partition), sorted.map((entry) => entry.nextOffset)], 'upsertCohortMembershipProgress' ) } @@ -506,8 +506,11 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { const cohortMembershipChanges = this._parseAndValidateBatch(messages) - const progress = this.config.COHORT_MEMBERSHIP_SWEEP_ENABLED ? partitionProgressFromMessages(messages) : [] + const progress = this.config.COHORT_MEMBERSHIP_VERSION_WRITES_ENABLED + ? partitionProgressFromMessages(messages) + : [] await this.persistCohortMembershipChanges(cohortMembershipChanges, progress) } @@ -557,26 +562,31 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase { - return instrumentFn('cdpCohortMembershipConsumer.handleMarkerBatch', async () => { - await this.handleMarkerBatch(messages) - }) - }) - - // The topic is not auto-created, and a consumer subscribed to a missing topic is a - // green pod that marks nothing: every run would sit `collecting` until abandonment. - // The flag being set is a claim that this environment runs the US-scoped cohort - // pipeline that produces the topic, so hold it to that. The fix for tripping this is - // to unset the flag, not to create the topic: a marker topic with no processor behind - // it turns the loud failure back into the silent one. - const markerPartitions = - await this.markerKafkaConsumer.getPartitionsForTopic(KAFKA_COHORT_RECONCILE_MARKERS) + // A consumer subscribed to a missing topic is a green pod that marks nothing: every + // run would sit `collecting` until abandonment. The flag being set is a claim that + // this environment runs the US-scoped cohort pipeline that produces the topic, so + // hold it to that. The fix for tripping this is to unset the flag, not to create the + // topic: a marker topic with no processor behind it turns the loud failure back into + // the silent one. + // + // The probe reads off the membership client, before the marker consumer connects: + // connect() creates the topic when CONSUMER_AUTO_CREATE_TOPICS is on (the default), + // which would make this check pass exactly where it has to fail. It assumes the two + // topics share brokers, which is today's deployment; a future split fails here + // loudly rather than passing silently. + const markerPartitions = await this.kafkaConsumer.getPartitionsForTopic(KAFKA_COHORT_RECONCILE_MARKERS) if (markerPartitions.length === 0) { throw new Error( `COHORT_MEMBERSHIP_SWEEP_ENABLED is set but ${KAFKA_COHORT_RECONCILE_MARKERS} has no partitions` ) } + await this.markerKafkaConsumer.connect(async (messages) => { + return instrumentFn('cdpCohortMembershipConsumer.handleMarkerBatch', async () => { + await this.handleMarkerBatch(messages) + }) + }) + this.sweeper.start() } } diff --git a/nodejs/src/cdp/consumers/cdp-hogflow-subscription-matcher.consumer.ts b/nodejs/src/cdp/consumers/cdp-hogflow-subscription-matcher.consumer.ts index 1cf8762a0c0b..d82a3822f6ef 100644 --- a/nodejs/src/cdp/consumers/cdp-hogflow-subscription-matcher.consumer.ts +++ b/nodejs/src/cdp/consumers/cdp-hogflow-subscription-matcher.consumer.ts @@ -9,7 +9,7 @@ import { KAFKA_PERSON, KAFKA_PERSON_DISTINCT_ID, } from '~/common/config/kafka-topics' -import { KafkaConsumerInterface, RdKafkaConsumerConfig, createKafkaConsumer } from '~/common/kafka/consumer' +import { KafkaConsumerInterface, START_AT_LATEST, createKafkaConsumer } from '~/common/kafka/consumer' import { InternalCaptureEvent } from '~/common/services/internal-capture' import { instrumentFn, instrumented } from '~/common/tracing/tracing-utils' import { parseJSON } from '~/common/utils/json-parse' @@ -46,7 +46,7 @@ import { counterParseError } from './metrics' // volume) topic backlog. Once the group has committed offsets, both consumer impls resume from them // regardless of this value — so steady-state recovery is committed-offset resume; latest only governs // the first bootstrap and offset-loss edge cases (covered by the deferred lag alerting follow-up). -const startAtLatest = { ['auto.offset.reset' as keyof RdKafkaConsumerConfig]: 'latest' as never } +const startAtLatest = START_AT_LATEST const counterHogflowMatcherCandidatesEvaluated = new Counter({ name: 'cdp_hogflow_matcher_candidates_evaluated', diff --git a/nodejs/src/cdp/dev/cohort-sweep-e2e.test.ts b/nodejs/src/cdp/dev/cohort-sweep-e2e.test.ts index 16f0ca36f926..c17f51229897 100644 --- a/nodejs/src/cdp/dev/cohort-sweep-e2e.test.ts +++ b/nodejs/src/cdp/dev/cohort-sweep-e2e.test.ts @@ -50,7 +50,12 @@ describe('cohort membership mark and sweep, end to end', () => { ) consumer = new CdpCohortMembershipConsumer( - { ...hub, COHORT_MEMBERSHIP_SWEEP_ENABLED: true, COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS: 1000 }, + { + ...hub, + COHORT_MEMBERSHIP_VERSION_WRITES_ENABLED: true, + COHORT_MEMBERSHIP_SWEEP_ENABLED: true, + COHORT_MEMBERSHIP_SWEEP_INTERVAL_MS: 1000, + }, createCdpConsumerDeps(hub) ) await consumer.start() diff --git a/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts b/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts index a519e05345bd..3e9cf57d5d1e 100644 --- a/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts +++ b/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts @@ -271,7 +271,61 @@ describe('CohortMembershipSweeper', () => { expect(await survivingPersonIds()).not.toContain(stale) }) - it('should abandon a run whose marker set never completed, but not a gate-blocked ready run', async () => { + it('should yield a sweep whose claim was taken over mid-run', async () => { + await readyRun() + await openGate() + await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + const unreached = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + + // Take the claim over after the first delete page, as a reclaim after a stall would. The + // heartbeat must notice and yield, instead of paging on and writing the terminal ledger + // row over the new owner's state. + const newOwnerToken = new UUIDT().toString() + const realQuery = postgres.query.bind(postgres) + let pages = 0 + postgres.query = (async (...args: any[]) => { + const result = await (realQuery as any)(...args) + if (args[3] === 'sweepCohortMembership' && ++pages === 1) { + await (realQuery as any)( + PostgresUse.BEHAVIORAL_COHORTS_RW, + 'UPDATE cohort_membership_sweeps SET claim_token = $1 WHERE run_id = $2', + [newOwnerToken, runId], + 'stealClaim' + ) + } + return result + }) as typeof postgres.query + + expect(await sweeper.runOnce()).toMatchObject({ rowsDeleted: 2 }) + expect(await survivingPersonIds()).toEqual([unreached]) + expect(await readSweep()).toMatchObject({ status: 'sweeping', swept_rows: null }) + }) + + it('should send an interrupted sweep back to ready instead of retiring it', async () => { + await readyRun() + await openGate() + await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + const unreached = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + + // Shutdown lands between delete pages. Marking the run swept here would strand the rows + // it never reached, because nothing revisits a swept run. + const realQuery = postgres.query.bind(postgres) + postgres.query = (async (...args: any[]) => { + const result = await (realQuery as any)(...args) + if (args[3] === 'sweepCohortMembership') { + sweeper['stopping'] = true + } + return result + }) as typeof postgres.query + + expect(await sweeper.runOnce()).toMatchObject({ rowsDeleted: 2 }) + expect(await survivingPersonIds()).toEqual([unreached]) + expect(await readSweep()).toMatchObject({ status: 'ready', swept_rows: '2' }) + }) + + it('should abandon unproven and expired runs, but no proven run, even while watermark capture fails', async () => { await query( `INSERT INTO cohort_membership_sweeps (run_id, cohort_id, team_id, marker_bits, created_at, updated_at) VALUES ($1, $2, $3, 7, CURRENT_TIMESTAMP - INTERVAL '30 days', @@ -292,13 +346,41 @@ describe('CohortMembershipSweeper', () => { [blockedRunId, COHORT_ID + 1, TEAM_ID, MARKER_VERSION, THRESHOLD, JSON.stringify({ 0: 10 })] ) - expect(await sweeper.runOnce()).toMatchObject({ abandoned: 1 }) + // A fully-marked run that promotion cannot ready, because watermark capture keeps + // failing, is proven: the short timeout must not take it, and the failure must stay + // contained to promotion so the abandonment below still runs in the same cycle. Only the + // marker-retention horizon retires such a run. + const unpromotedRunId = new UUIDT().toString() + await query( + `INSERT INTO cohort_membership_sweeps (run_id, cohort_id, team_id, marker_bits, created_at, updated_at) + VALUES ($1, $2, $3, -1, CURRENT_TIMESTAMP - INTERVAL '5 days', + CURRENT_TIMESTAMP - INTERVAL '5 days')`, + [unpromotedRunId, COHORT_ID + 2, TEAM_ID] + ) + const expiredRunId = new UUIDT().toString() + await query( + `INSERT INTO cohort_membership_sweeps (run_id, cohort_id, team_id, marker_bits, created_at, updated_at) + VALUES ($1, $2, $3, -1, CURRENT_TIMESTAMP - INTERVAL '31 days', + CURRENT_TIMESTAMP - INTERVAL '31 days')`, + [expiredRunId, COHORT_ID + 3, TEAM_ID] + ) + sweeper['kafka'].captureMembershipWatermarks = () => Promise.reject(new Error('kafka down')) + + expect(await sweeper.runOnce()).toMatchObject({ abandoned: 2 }) expect(await readSweep()).toMatchObject({ status: 'abandoned' }) - const blocked = await query<{ status: string }>( - `SELECT status FROM cohort_membership_sweeps WHERE run_id = $1`, - [blockedRunId] + const statuses = Object.fromEntries( + ( + await query<{ run_id: string; status: string }>( + 'SELECT run_id, status FROM cohort_membership_sweeps WHERE run_id = ANY($1)', + [[blockedRunId, unpromotedRunId, expiredRunId]] + ) + ).map((row) => [row.run_id, row.status]) ) - expect(blocked[0].status).toEqual('ready') + expect(statuses).toEqual({ + [blockedRunId]: 'ready', + [unpromotedRunId]: 'collecting', + [expiredRunId]: 'abandoned', + }) }) }) diff --git a/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts b/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts index 49f7c9101c9e..92539680ddb2 100644 --- a/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts +++ b/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts @@ -5,6 +5,7 @@ import { z } from 'zod' import { PostgresRouter, PostgresUse } from '~/common/utils/db/postgres' import { parseJSON } from '~/common/utils/json-parse' import { logger } from '~/common/utils/logger' +import { retryIfRetriable } from '~/common/utils/retries' import { UUIDT } from '~/common/utils/utils' /** @@ -82,6 +83,27 @@ type ClaimableRun = { claim_token: string | null } +/** + * The rule for which runs may be swept, shared by the tick scan and the claim UPDATE so the two + * cannot drift: only the copy inside the claim actually protects the delete. + * + * A run with no snapshot minimum is never claimable. Marker stamps alone cannot bound the + * snapshot: a fast partition finishes and stamps its marker while a slower one is still emitting + * rows, so the earliest marker can sit above rows the run asserted. Sweeping there would delete + * live members. Such a run waits for abandonment instead. This also excludes a run whose snapshot + * was legitimately empty (a cohort edited so nobody matches): it can never prove a snapshot + * minimum, so its stale rows outlive every reconcile until they are cleaned up out of band. + */ +function claimablePredicate(intervalParam: string): string { + return `min_snapshot_version IS NOT NULL + AND (status = 'ready' + OR (status = 'sweeping' AND claimed_at < CURRENT_TIMESTAMP - ${intervalParam}::interval))` +} + +const CLAIMABLE_PROJECTION = `run_id, cohort_id, team_id, membership_hwms, + membership_cluster, membership_topic, + LEAST(min_snapshot_version, min_marker_version) AS threshold, claim_token` + const markersConsumed = new Counter({ name: 'cdp_cohort_membership_markers_consumed', help: 'Reconcile completion markers read off the marker topic, by parse outcome', @@ -185,40 +207,82 @@ export class CohortMembershipSweeper { } /** - * Folding a marker in is a bit union, so redelivery changes nothing. This path runs inside the - * marker consumer's batch, whose offsets commit whether or not each fold landed, so a failed - * fold is caught and counted rather than rethrown: a rethrow would take the membership consumer - * in the same process down with it. A lost fold leaves the run one bit short until it ages into - * abandonment, which `sweepsAbandoned` surfaces, and the next reconcile heals the cohort. + * The bits are a union and the version stamp is a minimum, so a batch folds per run in memory + * first: one statement per (run_id, cohort_id) instead of one per marker, and a full batch of + * one run's markers is a single write. Redelivery changes nothing either way. + * + * This path runs inside the marker consumer's batch, whose offsets commit whether or not each + * fold landed, so a failed fold is retried and then counted rather than rethrown: a rethrow + * would take the membership consumer in the same process down with it. A lost fold leaves the + * run short until it ages into abandonment, which `sweepsAbandoned` surfaces, and the next + * reconcile heals the cohort. */ public async applyMarkers(markers: ReconcileCompleteMarker[]): Promise { + type FoldedMarkers = { + runId: string + cohortId: number + teamId: number + bits: bigint + minVersion: string + partitions: number[] + } + const folded = new Map() + for (const marker of markers) { + const key = `${marker.run_id}:${marker.cohort_id}` + const existing = folded.get(key) + + if (!existing) { + folded.set(key, { + runId: marker.run_id, + cohortId: marker.cohort_id, + teamId: marker.team_id, + bits: BigInt(1) << BigInt(marker.partition), + minVersion: marker.last_updated, + partitions: [marker.partition], + }) + continue + } + + existing.bits |= BigInt(1) << BigInt(marker.partition) + existing.partitions.push(marker.partition) + if (marker.last_updated < existing.minVersion) { + existing.minVersion = marker.last_updated + } + } + + for (const entry of folded.values()) { + // Postgres bigint is signed, so the last partition's bit has to fold to the sign bit. + const bits = BigInt.asIntN(64, entry.bits).toString() + try { - await this.postgres.query( - PostgresUse.BEHAVIORAL_COHORTS_RW, - ` - INSERT INTO cohort_membership_sweeps - (run_id, cohort_id, team_id, marker_bits, min_marker_version) - VALUES ($1, $2, $3, (1::bigint << $4), $5::timestamp) - ON CONFLICT (run_id, cohort_id) - DO UPDATE SET - marker_bits = cohort_membership_sweeps.marker_bits | (1::bigint << $4), - min_marker_version = LEAST( - cohort_membership_sweeps.min_marker_version, - EXCLUDED.min_marker_version - ), - updated_at = CURRENT_TIMESTAMP - `, - [marker.run_id, marker.cohort_id, marker.team_id, marker.partition, marker.last_updated], - 'applyCohortReconcileMarker' + await retryIfRetriable(() => + this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + INSERT INTO cohort_membership_sweeps + (run_id, cohort_id, team_id, marker_bits, min_marker_version) + VALUES ($1, $2, $3, $4::bigint, $5::timestamp) + ON CONFLICT (run_id, cohort_id) + DO UPDATE SET + marker_bits = cohort_membership_sweeps.marker_bits | $4::bigint, + min_marker_version = LEAST( + cohort_membership_sweeps.min_marker_version, + EXCLUDED.min_marker_version + ), + updated_at = CURRENT_TIMESTAMP + `, + [entry.runId, entry.cohortId, entry.teamId, bits, entry.minVersion], + 'applyCohortReconcileMarker' + ) ) } catch (error) { - markersConsumed.inc({ outcome: 'apply_failed' }) - logger.error('Failed to apply a reconcile marker', { + markersConsumed.inc({ outcome: 'apply_failed' }, entry.partitions.length) + logger.error('Failed to apply reconcile markers', { error: String(error), - run_id: marker.run_id, - cohort_id: marker.cohort_id, - partition: marker.partition, + run_id: entry.runId, + cohort_id: entry.cohortId, + partitions: entry.partitions, }) } } @@ -330,27 +394,16 @@ export class CohortMembershipSweeper { /** * Runs that hold every marker, plus runs whose claiming pod died. Re-sweeping is harmless: the - * delete only ever removes rows below a threshold that does not move. - * - * A run with no snapshot minimum is never claimable. Marker stamps alone cannot bound the - * snapshot: a fast partition finishes and stamps its marker while a slower one is still - * emitting rows, so the earliest marker can sit above rows the run asserted. Sweeping there - * would delete live members. Such a run waits for abandonment instead. This also excludes a - * run whose snapshot was legitimately empty (a cohort edited so nobody matches): it can never - * prove a snapshot minimum, so its stale rows outlive every reconcile until they are cleaned - * up out of band. + * delete only ever removes rows below a threshold that does not move. See claimablePredicate + * for why a run with no snapshot minimum is excluded. */ private async claimableRuns(): Promise { const result = await this.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, ` - SELECT run_id, cohort_id, team_id, membership_hwms, - membership_cluster, membership_topic, - LEAST(min_snapshot_version, min_marker_version) AS threshold, claim_token + SELECT ${CLAIMABLE_PROJECTION} FROM cohort_membership_sweeps - WHERE min_snapshot_version IS NOT NULL - AND (status = 'ready' - OR (status = 'sweeping' AND claimed_at < CURRENT_TIMESTAMP - $1::interval)) + WHERE ${claimablePredicate('$1')} ORDER BY created_at LIMIT ${MAX_CANDIDATES_PER_TICK} `, @@ -444,12 +497,8 @@ export class CohortMembershipSweeper { SET status = 'sweeping', claimed_at = CURRENT_TIMESTAMP, claim_token = $4, updated_at = CURRENT_TIMESTAMP WHERE run_id = $1 AND cohort_id = $2 - AND min_snapshot_version IS NOT NULL - AND (status = 'ready' - OR (status = 'sweeping' AND claimed_at < CURRENT_TIMESTAMP - $3::interval)) - RETURNING run_id, cohort_id, team_id, membership_hwms, - membership_cluster, membership_topic, - LEAST(min_snapshot_version, min_marker_version) AS threshold, claim_token + AND ${claimablePredicate('$3')} + RETURNING ${CLAIMABLE_PROJECTION} `, [ run.run_id, @@ -466,32 +515,49 @@ export class CohortMembershipSweeper { /** Returns the deleted row count, or null when the run was refused rather than swept. */ private async sweep(run: ClaimableRun): Promise { if (!run.threshold) { - // Unreachable while claim() requires a snapshot minimum; kept so a future claim change - // cannot silently delete below no bound. Abandoning gives the row a terminal state - // instead of leaving it wedged in 'sweeping'. - logger.error('Refusing to sweep a run with no threshold', { - run_id: run.run_id, - cohort_id: run.cohort_id, - }) - await this.postgres.query( - PostgresUse.BEHAVIORAL_COHORTS_RW, - ` - UPDATE cohort_membership_sweeps - SET status = 'abandoned', claimed_at = NULL, claim_token = NULL, - updated_at = CURRENT_TIMESTAMP - WHERE run_id = $1 AND cohort_id = $2 AND status = 'sweeping' AND claim_token = $3 - `, - [run.run_id, run.cohort_id, run.claim_token], - 'abandonThresholdlessCohortMembershipSweep' - ) - sweepsAbandoned.inc({ reason: 'no_threshold' }) - return null + return await this.abandonThresholdless(run) } - // The threshold string was parsed to millisecond precision on the way out of Postgres, so - // it can sit up to 999 microseconds below the stored minimum. Truncation only ever lowers - // it, which makes the sweep under-delete, never over-delete; the next reconcile deletes - // the sliver it leaves behind. + const { deleted, claimLost, interrupted } = await this.deleteStaleRowsInPages(run) + rowsSwept.inc(deleted) + if (claimLost) { + // The new claim owner finishes the run and writes its terminal ledger row. + return deleted + } + + const finished = !interrupted && (await this.noRowsBelowThreshold(run)) + await this.finishSweep(run, deleted, finished) + return deleted + } + + /** + * Unreachable while claim() requires a snapshot minimum; kept so a future claim change cannot + * silently delete below no bound. Abandoning gives the row a terminal state instead of + * leaving it wedged in 'sweeping'. + */ + private async abandonThresholdless(run: ClaimableRun): Promise { + logger.error('Refusing to sweep a run with no threshold', { + run_id: run.run_id, + cohort_id: run.cohort_id, + }) + await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + ` + UPDATE cohort_membership_sweeps + SET status = 'abandoned', claimed_at = NULL, claim_token = NULL, + updated_at = CURRENT_TIMESTAMP + WHERE run_id = $1 AND cohort_id = $2 AND status = 'sweeping' AND claim_token = $3 + `, + [run.run_id, run.cohort_id, run.claim_token], + 'abandonThresholdlessCohortMembershipSweep' + ) + sweepsAbandoned.inc({ reason: 'no_threshold' }) + return null + } + + private async deleteStaleRowsInPages( + run: ClaimableRun + ): Promise<{ deleted: number; claimLost: boolean; interrupted: boolean }> { let deleted = 0 let lastBatch = 0 // Keyset cursor over the (team_id, cohort_id, version) index order. Deleted index entries @@ -510,6 +576,11 @@ export class CohortMembershipSweeper { // re-checks it against the newest row version, so a row a concurrent upsert just // refreshed past the threshold is left alone instead of deleted anyway. A refreshed // row also leaves the cursor range entirely, so skipping past it loses nothing. + // + // The threshold bound into `version < $3` was parsed to millisecond precision on the + // way out of Postgres, so it can sit up to 999 microseconds below the stored minimum. + // Truncation only ever lowers it, which makes the sweep under-delete, never + // over-delete; the next reconcile deletes the sliver it leaves behind. const result = await this.postgres.query<{ version: string; id: string }>( PostgresUse.BEHAVIORAL_COHORTS_RW, ` @@ -571,34 +642,38 @@ export class CohortMembershipSweeper { cohort_id: run.cohort_id, rows: deleted, }) - rowsSwept.inc(deleted) - return deleted + return { deleted, claimLost: true, interrupted: false } } } } while (lastBatch > 0 && !this.stopping) - let finished = false - if (this.stopping && lastBatch > 0) { - // Interrupted by shutdown; the run goes back to 'ready' and resumes on the next claim. - finished = false - } else { - // A zero-delete iteration means either nothing is left or SKIP LOCKED skipped - // everything. Only the first is finished; calling the second `swept` would retire the - // run with stale rows still in place, and nothing revisits a swept run. - const remaining = await this.postgres.query( - PostgresUse.BEHAVIORAL_COHORTS_RW, - `SELECT 1 FROM cohort_membership - WHERE team_id = $1 AND cohort_id = $2 AND version < $3::timestamp - LIMIT 1`, - [run.team_id, run.cohort_id, run.threshold], - 'checkCohortMembershipSweepComplete' - ) - finished = remaining.rows.length === 0 - } + // Interrupted by shutdown; the run goes back to 'ready' and resumes on the next claim. + return { deleted, claimLost: false, interrupted: this.stopping && lastBatch > 0 } + } - // The claim token fences stragglers: a pod whose claim timed out and was reclaimed no - // longer matches, so it cannot overwrite the new owner's state nor resurrect a run that - // was abandoned in between. swept_rows accumulates across partial passes. + /** + * A zero-delete iteration means either nothing is left or SKIP LOCKED skipped everything. + * Only the first is finished; calling the second `swept` would retire the run with stale rows + * still in place, and nothing revisits a swept run. + */ + private async noRowsBelowThreshold(run: ClaimableRun): Promise { + const remaining = await this.postgres.query( + PostgresUse.BEHAVIORAL_COHORTS_RW, + `SELECT 1 FROM cohort_membership + WHERE team_id = $1 AND cohort_id = $2 AND version < $3::timestamp + LIMIT 1`, + [run.team_id, run.cohort_id, run.threshold], + 'checkCohortMembershipSweepComplete' + ) + return remaining.rows.length === 0 + } + + /** + * The claim token fences stragglers: a pod whose claim timed out and was reclaimed no longer + * matches, so it cannot overwrite the new owner's state nor resurrect a run that was + * abandoned in between. swept_rows accumulates across partial passes. + */ + private async finishSweep(run: ClaimableRun, deleted: number, finished: boolean): Promise { await this.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, ` @@ -611,7 +686,6 @@ export class CohortMembershipSweeper { 'finishCohortMembershipSweep' ) - rowsSwept.inc(deleted) logger.info('Swept stale cohort membership rows', { run_id: run.run_id, cohort_id: run.cohort_id, @@ -619,16 +693,16 @@ export class CohortMembershipSweeper { rows: deleted, finished, }) - - return deleted } /** * Two distinct give-up paths, kept apart because only one is routine. A reconcile superseded * by a cohort edit emits no marker, so its `collecting` row can never complete; that is - * expected traffic. A `ready` run is fully proven and only waiting on consumer lag, and a - * `sweeping` run only ever fails transiently, so both get the marker-retention horizon rather - * than the short timeout: abandoning them sooner would discard runs that would have swept. + * expected traffic, and the short timeout retires it. Every proven run gets the + * marker-retention horizon instead, because abandoning it sooner would discard a run that + * would have swept: a `ready` run is only waiting on consumer lag, a `sweeping` run only ever + * fails transiently, and a `collecting` run holding every marker is only waiting on watermark + * capture, which the sweep loop retries each tick. * * These filters and the GC below scan without an index; the ledger is bounded by the GC to a * month of reconcile runs, so a sequential scan stays trivial. @@ -640,6 +714,7 @@ export class CohortMembershipSweeper { UPDATE cohort_membership_sweeps SET status = 'abandoned', updated_at = CURRENT_TIMESTAMP WHERE status = 'collecting' + AND marker_bits <> ${ALL_MARKERS} AND updated_at < CURRENT_TIMESTAMP - $1::interval `, [`${this.config.COHORT_MEMBERSHIP_SWEEP_ABANDON_AFTER_DAYS} days`], @@ -657,7 +732,8 @@ export class CohortMembershipSweeper { ` UPDATE cohort_membership_sweeps SET status = 'abandoned', claim_token = NULL, updated_at = CURRENT_TIMESTAMP - WHERE status IN ('ready', 'sweeping') + WHERE (status IN ('ready', 'sweeping') + OR (status = 'collecting' AND marker_bits = ${ALL_MARKERS})) AND created_at < CURRENT_TIMESTAMP - INTERVAL '${MARKER_RETENTION_DAYS} days' `, undefined, diff --git a/nodejs/src/common/kafka/consumer/consumer-v1.test.ts b/nodejs/src/common/kafka/consumer/consumer-v1.test.ts index c1cdfc1f96c7..4d93e5d73263 100644 --- a/nodejs/src/common/kafka/consumer/consumer-v1.test.ts +++ b/nodejs/src/common/kafka/consumer/consumer-v1.test.ts @@ -174,6 +174,22 @@ describe('consumer', () => { ]) }) + // auto.offset.reset is a topic-level librdkafka property, and a value only in the global + // config is ignored, so an override must be mirrored into the topic config (2nd ctor arg). + it.each([ + [ + 'mirrors an auto.offset.reset override into the topic config', + { 'auto.offset.reset': 'latest' } as any, + 'latest', + ], + ['defaults the topic-config auto.offset.reset to earliest with no override', undefined, 'earliest'], + ])('%s', (_label, overrides, expected) => { + const RdKafkaCtor = jest.mocked(RdKafkaConsumer) + new KafkaConsumer({ groupId: 'g', topic: 't' }, overrides) + const topicConfig = RdKafkaCtor.mock.calls.at(-1)![1] as any + expect(topicConfig['auto.offset.reset']).toBe(expected) + }) + describe('background work', () => { /** * NOTE: These tests are pretty verbose but also pretty cool! We are using special wrapped promises diff --git a/nodejs/src/common/kafka/consumer/consumer-v1.ts b/nodejs/src/common/kafka/consumer/consumer-v1.ts index 75253872ae47..9c7f810a88bc 100644 --- a/nodejs/src/common/kafka/consumer/consumer-v1.ts +++ b/nodejs/src/common/kafka/consumer/consumer-v1.ts @@ -520,9 +520,16 @@ export class KafkaConsumer { } private createConsumer(): RdKafkaConsumer { + // auto.offset.reset is a topic-level property, and librdkafka ignores the global-config + // form on our version (see node-rdkafka #984), so the resolved value, including any + // caller override, must be mirrored into the explicit topic config. v2 does the same. + const autoOffsetReset = + (this.consumerConfig['auto.offset.reset' as keyof ConsumerGlobalConfig] as + | 'earliest' + | 'latest' + | undefined) ?? 'earliest' const consumer = new RdKafkaConsumer(this.consumerConfig, { - // Default settings - 'auto.offset.reset': 'earliest', + 'auto.offset.reset': autoOffsetReset, }) consumer.on('event.log', (log) => { diff --git a/nodejs/src/common/kafka/consumer/index.ts b/nodejs/src/common/kafka/consumer/index.ts index c5908d0b7317..134bd14e3d60 100644 --- a/nodejs/src/common/kafka/consumer/index.ts +++ b/nodejs/src/common/kafka/consumer/index.ts @@ -12,6 +12,14 @@ import { KafkaConsumerV2 } from './consumer-v2' export { parseEventHeaders, parseKafkaHeaders } from './consumer-v1' export type { KafkaConsumerConfig, RdKafkaConsumerConfig } from './consumer-v1' +/** + * Overrides `auto.offset.reset` so a fresh consumer group starts at the tip of the topic instead + * of replaying the retained backlog. `auto.offset.reset` is a topic-level librdkafka property, so + * the global config type has no slot for it; both implementations mirror the value into the + * explicit topic config, which is the form librdkafka honors (see node-rdkafka #984). + */ +export const START_AT_LATEST = { ['auto.offset.reset' as keyof RdKafkaConsumerConfig]: 'latest' as never } + /** * The shared surface that both KafkaConsumer (v1) and KafkaConsumerV2 expose to call sites. * Keep this small — every method here pins the migration. Adding a method here forces both From 17feec18e82d4f0cc988e2e83a362345f7221b7e Mon Sep 17 00:00:00 2001 From: Matheus Veras Batista Date: Thu, 27 Aug 2026 22:16:54 -0300 Subject: [PATCH 9/9] fix(cohorts): apply sweep review feedback --- .../cdp-cohort-membership.consumer.test.ts | 39 +++++- .../cdp-cohort-membership.consumer.ts | 54 ++++---- .../sweeper.service.serial.test.ts | 116 ++++++++++++++---- .../cohort-membership/sweeper.service.ts | 26 ++-- 4 files changed, 170 insertions(+), 65 deletions(-) diff --git a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts index 00edf97fef8e..229037905686 100644 --- a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts +++ b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.test.ts @@ -45,6 +45,19 @@ describe('CdpCohortMembershipConsumer', () => { await closeHub(hub) }) + it('should refuse to construct a sweeping consumer whose own write path skips versions', () => { + // The flag ordering (version writes on first and off last, sweep on last and off first) + // is enforced nowhere else: a pod that swept while writing without versions would delete + // rows its own upserts left carrying a stale stamp. + expect( + () => + new CdpCohortMembershipConsumer( + { ...hub, COHORT_MEMBERSHIP_VERSION_WRITES_ENABLED: false, COHORT_MEMBERSHIP_SWEEP_ENABLED: true }, + createCdpConsumerDeps(hub) + ) + ).toThrow('COHORT_MEMBERSHIP_SWEEP_ENABLED requires COHORT_MEMBERSHIP_VERSION_WRITES_ENABLED') + }) + describe('end-to-end cohort membership processing', () => { const readProgress = async (): Promise[]> => { const result = await hub.postgres.query( @@ -199,11 +212,23 @@ describe('CdpCohortMembershipConsumer', () => { }) it.each([ - ['keeping last in Kafka order for equal versions', {}, {}, false], + ['keeping last in Kafka order for equal versions', false, {}, {}, false], + // The version-aware pick only exists on the version-writes path. With the flag off + // there is no SQL guard for it to agree with, so Kafka order must win even when the + // stamps run backwards: anything else would change the write shape on deploy, before + // any flag flips. + [ + 'keeping last in Kafka order even for a lower version while version writes are off', + false, + { last_updated: '2026-05-26 13:00:00.000000' }, + { last_updated: '2026-05-26 12:00:00.000000' }, + false, + ], // The SQL guard only ever sees the batch's surviving entry, so the in-memory pick has // to agree with last-writer-wins: a newer version beats a later offset. [ - 'keeping the highest version regardless of order', + 'keeping the highest version regardless of order with version writes on', + true, { last_updated: '2026-05-26 13:00:00.000000' }, { last_updated: '2026-05-26 12:00:00.000000' }, true, @@ -211,14 +236,16 @@ describe('CdpCohortMembershipConsumer', () => { // A versionless change applies unconditionally in SQL, so it also has to win the // in-batch pick, even against a versioned entry earlier in Kafka order. [ - 'keeping a later versionless change over an earlier versioned one', + 'keeping a later versionless change over an earlier versioned one with version writes on', + true, { last_updated: '2026-05-26 13:00:00.000000' }, {}, false, ], ])( 'should deduplicate batch entries for the same (team_id, cohort_id, person_id), %s', - async (_label, firstTags, secondTags, expectedInCohort) => { + async (_label, withVersionWrites, firstTags, secondTags, expectedInCohort) => { + const target = withVersionWrites ? sweepConsumer : consumer const testEvents = createCohortMembershipEvents([ { person_id: personId1, @@ -240,8 +267,8 @@ describe('CdpCohortMembershipConsumer', () => { createKafkaMessage(event, { topic: KAFKA_COHORT_MEMBERSHIP_CHANGED, offset: index }) ) - const cohortMembershipChanges = consumer['_parseAndValidateBatch'](messages) - await consumer['persistCohortMembershipChanges'](cohortMembershipChanges) + const cohortMembershipChanges = target['_parseAndValidateBatch'](messages) + await target['persistCohortMembershipChanges'](cohortMembershipChanges) const result = await hub.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, diff --git a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts index e124f0b8a6fe..5a6cfe1210d9 100644 --- a/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts +++ b/nodejs/src/cdp/consumers/cdp-cohort-membership.consumer.ts @@ -16,6 +16,7 @@ import { CohortMembershipSweeper, MembershipWatermarks, PRODUCER_VERSION_FORMAT, + VERSIONLESS, } from '../services/cohort-membership/sweeper.service' import { CdpConsumerBase, CdpConsumerBaseConfig, CdpConsumerBaseDeps } from './cdp-base.consumer' @@ -30,20 +31,12 @@ export type CdpCohortMembershipConsumerConfig = CdpConsumerBaseConfig & | 'COHORT_MEMBERSHIP_SWEEP_ABANDON_AFTER_DAYS' > -/** - * The stamp for a membership row the pipeline never versioned. Must stay equal to the column - * default in `20260826000001_add_version_to_cohort_membership.sql`: every row written before that - * migration reads as this value, and the sweep deletes below a threshold, so a sentinel that does - * not sort below every producer stamp would leave stale rows behind. - */ -const VERSIONLESS = '-infinity' - /** * The broker list the membership consumer resolves, mirroring the consumer's own config - * resolution. Progress offsets are only comparable against watermarks from the same cluster, so - * this keys every progress row: after a cluster move the old rows stop matching, and the gate - * starts from no progress instead of comparing the old cluster's large offsets against the new - * cluster's small watermarks. + * resolution. Progress offsets are only comparable against watermarks from the same feed, so + * this keys every progress row together with the topic: after a cluster move or a topic change + * the old rows stop matching, and the gate starts from no progress instead of comparing another + * feed's large offsets against the current feed's small watermarks. */ function membershipClusterId(): string { return String(getKafkaConfigFromEnv('CONSUMER')['metadata.broker.list'] ?? 'kafka:9092') @@ -237,8 +230,8 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase( PostgresUse.BEHAVIORAL_COHORTS_RW, - 'SELECT COUNT(*) AS count FROM cohort_membership_consumer_progress WHERE cluster = $1', - [this.membershipCluster], + 'SELECT COUNT(*) AS count FROM cohort_membership_consumer_progress WHERE cluster = $1 AND topic = $2', + [this.membershipCluster, KAFKA_COHORT_MEMBERSHIP_CHANGED], 'countCohortMembershipProgressPartitions' ) @@ -299,15 +292,23 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase() for (const change of changes) { const key = `${change.team_id}:${change.cohort_id}:${change.person_id}` const existing = deduped.get(key) - if (!existing || !change.last_updated || change.last_updated >= (existing.last_updated ?? '')) { + if ( + !existing || + !versionAware || + !change.last_updated || + change.last_updated >= (existing.last_updated ?? '') + ) { deduped.set(key, change) } } @@ -411,7 +412,7 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase change.person_id), changes.map((change) => change.status === 'entered'), ], - 'batchUpsertCohortMembership' + 'batchUpsertCohortMembershipWithoutVersion' ) } @@ -463,9 +464,9 @@ export class CdpCohortMembershipConsumer extends CdpConsumerBase entry.partition), sorted.map((entry) => entry.nextOffset)], + [ + this.membershipCluster, + KAFKA_COHORT_MEMBERSHIP_CHANGED, + sorted.map((entry) => entry.partition), + sorted.map((entry) => entry.nextOffset), + ], 'upsertCohortMembershipProgress' ) } diff --git a/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts b/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts index 3e9cf57d5d1e..c94ad02c7c30 100644 --- a/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts +++ b/nodejs/src/cdp/services/cohort-membership/sweeper.service.serial.test.ts @@ -39,17 +39,20 @@ describe('CohortMembershipSweeper', () => { let postgres: PostgresRouter let sweeper: CohortMembershipSweeper let watermarks: MembershipWatermarks + let captureFailure: Error | null let runId: string beforeEach(async () => { postgres = new PostgresRouter(defaultConfig) await resetBehavioralCohortsDatabase(postgres) watermarks = { 0: 10 } + captureFailure = null runId = new UUIDT().toString() sweeper = new CohortMembershipSweeper(config, postgres, { cluster: CLUSTER, topic: TOPIC, - captureMembershipWatermarks: () => Promise.resolve(watermarks), + captureMembershipWatermarks: () => + captureFailure ? Promise.reject(captureFailure) : Promise.resolve(watermarks), refreshConsumerProgress: () => Promise.resolve(), }) }) @@ -124,8 +127,9 @@ describe('CohortMembershipSweeper', () => { const insertProgress = async (partition: number, nextOffset: number): Promise => { await query( - `INSERT INTO cohort_membership_consumer_progress (cluster, partition, next_offset) VALUES ($1, $2, $3)`, - [CLUSTER, partition, nextOffset] + `INSERT INTO cohort_membership_consumer_progress (cluster, topic, partition, next_offset) + VALUES ($1, $2, $3, $4)`, + [CLUSTER, TOPIC, partition, nextOffset] ) } @@ -136,6 +140,19 @@ describe('CohortMembershipSweeper', () => { } } + /** Run `onPage` after each delete page the sweep issues. */ + const afterEachDeletePage = (onPage: (page: number) => Promise | void): void => { + const realQuery = postgres.query.bind(postgres) + let pages = 0 + postgres.query = (async (...args: any[]) => { + const result = await (realQuery as any)(...args) + if (args[3] === 'sweepCohortMembership') { + await onPage(++pages) + } + return result + }) as typeof postgres.query + } + it('should skip unusable markers instead of failing the batch', () => { const messages: Message[] = [ markerMessage(0), @@ -191,7 +208,10 @@ describe('CohortMembershipSweeper', () => { }) it('should hold the sweep until consumer progress passes every captured watermark', async () => { - watermarks = { 0: 10, 3: 4 } + // Partition 5 never had a message produced to it, so it also never gets a progress row. + // The gate must skip it rather than wait on it: on a topic with idle partitions, waiting + // would block every sweep forever. + watermarks = { 0: 10, 3: 4, 5: 0 } await readyRun() const stale = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) @@ -282,20 +302,14 @@ describe('CohortMembershipSweeper', () => { // heartbeat must notice and yield, instead of paging on and writing the terminal ledger // row over the new owner's state. const newOwnerToken = new UUIDT().toString() - const realQuery = postgres.query.bind(postgres) - let pages = 0 - postgres.query = (async (...args: any[]) => { - const result = await (realQuery as any)(...args) - if (args[3] === 'sweepCohortMembership' && ++pages === 1) { - await (realQuery as any)( - PostgresUse.BEHAVIORAL_COHORTS_RW, - 'UPDATE cohort_membership_sweeps SET claim_token = $1 WHERE run_id = $2', - [newOwnerToken, runId], - 'stealClaim' - ) + afterEachDeletePage(async (page) => { + if (page === 1) { + await query('UPDATE cohort_membership_sweeps SET claim_token = $1 WHERE run_id = $2', [ + newOwnerToken, + runId, + ]) } - return result - }) as typeof postgres.query + }) expect(await sweeper.runOnce()).toMatchObject({ rowsDeleted: 2 }) expect(await survivingPersonIds()).toEqual([unreached]) @@ -311,21 +325,43 @@ describe('CohortMembershipSweeper', () => { // Shutdown lands between delete pages. Marking the run swept here would strand the rows // it never reached, because nothing revisits a swept run. - const realQuery = postgres.query.bind(postgres) - postgres.query = (async (...args: any[]) => { - const result = await (realQuery as any)(...args) - if (args[3] === 'sweepCohortMembership') { - sweeper['stopping'] = true - } - return result - }) as typeof postgres.query + afterEachDeletePage(() => { + sweeper['stopping'] = true + }) expect(await sweeper.runOnce()).toMatchObject({ rowsDeleted: 2 }) expect(await survivingPersonIds()).toEqual([unreached]) expect(await readSweep()).toMatchObject({ status: 'ready', swept_rows: '2' }) }) - it('should abandon unproven and expired runs, but no proven run, even while watermark capture fails', async () => { + it('should send a sweep whose rows were lock-skipped back to ready instead of retiring it', async () => { + await readyRun() + await openGate() + const locked = await insertMembership(TEAM_ID, COHORT_ID, BELOW_THRESHOLD) + + // A row another transaction holds locked is skipped by FOR UPDATE SKIP LOCKED, so the + // delete page comes back empty while a stale row remains. Retiring the run as swept here + // would strand that row, because nothing revisits a swept run. + await postgres.transaction(PostgresUse.BEHAVIORAL_COHORTS_RW, 'holdRowLock', async (tx) => { + await postgres.query( + tx, + 'SELECT id FROM cohort_membership WHERE person_id = $1 FOR UPDATE', + [locked], + 'lockStaleRow' + ) + expect(await sweeper.runOnce()).toMatchObject({ swept: 1, rowsDeleted: 0 }) + }) + + expect(await survivingPersonIds()).toEqual([locked]) + expect(await readSweep()).toMatchObject({ status: 'ready', swept_rows: '0' }) + + // With the lock released, the next cycle finishes the run. + expect(await sweeper.runOnce()).toMatchObject({ swept: 1, rowsDeleted: 1 }) + expect(await survivingPersonIds()).toEqual([]) + expect(await readSweep()).toMatchObject({ status: 'swept', swept_rows: '1' }) + }) + + it('should abandon unproven and expired runs and collect aged garbage, sparing proven runs, even while watermark capture fails', async () => { await query( `INSERT INTO cohort_membership_sweeps (run_id, cohort_id, team_id, marker_bits, created_at, updated_at) VALUES ($1, $2, $3, 7, CURRENT_TIMESTAMP - INTERVAL '30 days', @@ -364,11 +400,37 @@ describe('CohortMembershipSweeper', () => { CURRENT_TIMESTAMP - INTERVAL '31 days')`, [expiredRunId, COHORT_ID + 3, TEAM_ID] ) - sweeper['kafka'].captureMembershipWatermarks = () => Promise.reject(new Error('kafka down')) + + // GC hygiene rides the same cycle: a terminal ledger row and a progress row past the + // marker-retention horizon are deleted, while recent rows survive. The progress bound is + // the risky half: deleting a live partition's row would make the gate block every sweep. + const gcRunId = new UUIDT().toString() + await query( + `INSERT INTO cohort_membership_sweeps + (run_id, cohort_id, team_id, marker_bits, status, created_at, updated_at) + VALUES ($1, $2, $3, -1, 'swept', CURRENT_TIMESTAMP - INTERVAL '31 days', + CURRENT_TIMESTAMP - INTERVAL '31 days')`, + [gcRunId, COHORT_ID + 4, TEAM_ID] + ) + await insertProgress(0, 5) + await query( + `INSERT INTO cohort_membership_consumer_progress (cluster, topic, partition, next_offset, updated_at) + VALUES ($1, $2, 1, 5, CURRENT_TIMESTAMP - INTERVAL '31 days')`, + [CLUSTER, TOPIC] + ) + + captureFailure = new Error('kafka down') expect(await sweeper.runOnce()).toMatchObject({ abandoned: 2 }) expect(await readSweep()).toMatchObject({ status: 'abandoned' }) + expect(await query('SELECT 1 FROM cohort_membership_sweeps WHERE run_id = $1', [gcRunId])).toEqual([]) + expect( + (await query<{ partition: number }>('SELECT partition FROM cohort_membership_consumer_progress')).map( + (row) => row.partition + ) + ).toEqual([0]) + const statuses = Object.fromEntries( ( await query<{ run_id: string; status: string }>( diff --git a/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts b/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts index 92539680ddb2..896fef73b9ed 100644 --- a/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts +++ b/nodejs/src/cdp/services/cohort-membership/sweeper.service.ts @@ -23,6 +23,16 @@ export const COHORT_PARTITION_COUNT = 64 */ export const PRODUCER_VERSION_FORMAT = /^\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{6}$/ +/** + * The stamp for a membership row the pipeline never versioned. Must stay equal to the column + * default in `20260826000001_add_version_to_cohort_membership.sql`: every row written before that + * migration reads as this value, and the sweep deletes below a threshold, so a sentinel that does + * not sort below every producer stamp would leave stale rows behind. The delete pager also starts + * its keyset cursor here, so a value that stopped matching the column default would make the + * first page match nothing. + */ +export const VERSIONLESS = '-infinity' + /** All 64 bits set, read as a signed 64-bit integer. Same convention as the seeder's ledger. */ const ALL_MARKERS = '-1' @@ -162,9 +172,9 @@ export class CohortMembershipSweeper { private postgres: PostgresRouter, /** All three describe the membership consumer's own client: the marker topic may live on another cluster. */ private kafka: { - /** Broker identity the progress rows are keyed by; offsets are only comparable within one cluster. */ + /** Broker identity half of the progress-row key; offsets are only comparable within one cluster. */ cluster: string - /** Membership topic the watermarks and the progress offsets both describe. */ + /** Topic half of the progress-row key; the watermarks and the progress offsets both describe it. */ topic: string captureMembershipWatermarks: () => Promise refreshConsumerProgress: () => Promise @@ -417,8 +427,8 @@ export class CohortMembershipSweeper { private async readConsumerProgress(): Promise> { const result = await this.postgres.query<{ partition: number; next_offset: string }>( PostgresUse.BEHAVIORAL_COHORTS_RW, - 'SELECT partition, next_offset FROM cohort_membership_consumer_progress WHERE cluster = $1', - [this.kafka.cluster], + 'SELECT partition, next_offset FROM cohort_membership_consumer_progress WHERE cluster = $1 AND topic = $2', + [this.kafka.cluster, this.kafka.topic], 'readCohortMembershipProgress' ) @@ -568,7 +578,7 @@ export class CohortMembershipSweeper { // constructor still advances through those rows by id. `to_char` returns NULL for an // infinite timestamp, so the RETURNING clause restores the sentinel text: without it the // cursor would take a NULL and the next page would match nothing. - let cursorVersion = '-infinity' + let cursorVersion: string = VERSIONLESS let cursorId = '0' do { @@ -597,7 +607,7 @@ export class CohortMembershipSweeper { USING candidates c WHERE m.id = c.id AND m.version < $3::timestamp RETURNING COALESCE( - to_char(m.version, 'YYYY-MM-DD HH24:MI:SS.US'), '-infinity' + to_char(m.version, 'YYYY-MM-DD HH24:MI:SS.US'), '${VERSIONLESS}' ) AS version, m.id `, [ @@ -764,8 +774,8 @@ export class CohortMembershipSweeper { 'gcCohortMembershipSweeps' ) - // Progress rows keyed to a broker cluster the consumer left stop updating and would - // otherwise sit forever. + // Progress rows keyed to a feed the consumer left, by cluster move or topic change, stop + // updating and would otherwise sit forever. await this.postgres.query( PostgresUse.BEHAVIORAL_COHORTS_RW, `