Repository: PostHog/posthog
Pull request: feat(ingestion): run all analytic events through populateTeamDataStep
## Problem

Closes #14327, part of removing the Postgres dependency from capture.py

## Changes

- Run all events through the `populateTeamDataStep`, not just events without a team_id. This allows us to:
  - confirm that the Postgres load will be OK once we fully roll it out (and rollback without getting in a pickle)
  - control the rollout from a single point (dropping the team_id field in capture)
- If a `team_id` field is present, trust it but report inconsistencies with the new `ingestion_team_resolution_checks_total` counter
- Emit a new `ingestion_event_dropped_total` counter, that could be used by other parts of the pipeline (dropping events from apps)
- Tune the team cache settings, only cache negative lookups for 5 minutes
- Don't drop the token from the event, mirroring previous behaviour


## How did you test this code?

unit and end2end tests

Problem statement:
Lightweigth capture 1/4: always resolve team_id at p-s and compare result
Let's start doing it as soon as possible to build trust on the performance and accuracy of the resolution code in plugin-server.

- [ ] Always hit postgres for resolution, even if `team_id` is found in the event
- [ ] If `team_id` is found in the event, compare the lookup result, and increment a prometheus metric with ok/error status. Log differences if found, and keep the value from capture
- [ ] Let's review the caching logic doesn't get into edge cases when tokens are rotated.


## Why now?

This has been identified as something not to port in a rewrite of capture. Regardless of whether capture v2 plans move forward in Q1, this will help with:

- making sure we don't DLQ & timeout when PG is unhealthy
- allowing infra to have a less stressful PG upgrade (version 11 goes EOL in October)
- scaling down posthog-events, as the team resolution is a 20%-ish part of the endpoint latency

Hints from issue discussion:
whether it's at the django or plugin-server level, for flags we already have a token->team cache in redis, which is somewhat battle-tested. You can re-use this too for the prom metric and then decide if it's good enough to switch :) 
- if plugin-server-ingestion has lag while a user rotates their token, in-flights events will be dropped  -> update documentation
- invalidated tokens will stay up for up to 2 minutes (the plugin-server caching TTL), we could later on add faster invalidation (django sents a kakfa message to p-s), but out of scope for now.

-> Update doc to reflect these two behaviours.

Commit message: ingestEvent: run all analytic events through runLightweightCaptureEndpointEventPipeline

Patch to review:
diff --git a/plugin-server/src/main/ingestion-queues/batch-processing/each-batch-ingestion.ts b/plugin-server/src/main/ingestion-queues/batch-processing/each-batch-ingestion.ts
index 3a7e114df8dab..b8530371aa3fa 100644
--- a/plugin-server/src/main/ingestion-queues/batch-processing/each-batch-ingestion.ts
+++ b/plugin-server/src/main/ingestion-queues/batch-processing/each-batch-ingestion.ts
@@ -1,4 +1,3 @@
-import { PluginEvent } from '@posthog/plugin-scaffold'
 import { EachBatchPayload, KafkaMessage } from 'kafkajs'
 
 import { KAFKA_SESSION_RECORDING_EVENTS } from '../../../config/kafka-topics'
@@ -61,22 +60,10 @@ export async function ingestEvent(
 
     checkAndPause?.()
 
-    // Eventually no events will have a team_id as that will be determined and
-    // populated in the plugin server rather than the capture endpoint. However,
-    // we support both paths during the transitional period.
-    if (event.team_id) {
-        server.statsd?.increment('kafka_queue_ingest_event_hit', {
-            pipeline: 'runEventPipeline',
-            team_id: event.team_id.toString(),
-        })
-        // we've confirmed team_id exists so can assert event as PluginEvent
-        await workerMethods.runEventPipeline(event as PluginEvent)
-    } else {
-        server.statsd?.increment('kafka_queue_ingest_event_hit', {
-            pipeline: 'runLightweightCaptureEndpointEventPipeline',
-        })
-        await workerMethods.runLightweightCaptureEndpointEventPipeline(event)
-    }
+    server.statsd?.increment('kafka_queue_ingest_event_hit', {
+        pipeline: 'runLightweightCaptureEndpointEventPipeline',
+    })
+    await workerMethods.runLightweightCaptureEndpointEventPipeline(event)
 
     server.statsd?.timing('kafka_queue.each_event', eachEventStartTimer)
 
diff --git a/plugin-server/src/main/ingestion-queues/metrics.ts b/plugin-server/src/main/ingestion-queues/metrics.ts
index 93c0d98a595d5..c3077d045b0ba 100644
--- a/plugin-server/src/main/ingestion-queues/metrics.ts
+++ b/plugin-server/src/main/ingestion-queues/metrics.ts
@@ -1,9 +1,15 @@
 // Metrics that make sense across all Kafka consumers
 
-import { Gauge } from 'prom-client'
+import { Counter, Gauge } from 'prom-client'
 
 export const latestOffsetTimestampGauge = new Gauge({
     name: 'latest_processed_timestamp_ms',
     help: 'Timestamp of the latest offset that has been committed.',
     labelNames: ['topic', 'partition', 'groupId'],
 })
+
+export const eventDroppedCounter = new Counter({
+    name: 'ingestion_event_dropped_total',
+    help: 'Count of events dropped by the ingestion pipeline, by type and cause.',
+    labelNames: ['event_type', 'drop_cause'],
+})
diff --git a/plugin-server/src/worker/ingestion/event-pipeline/populateTeamDataStep.ts b/plugin-server/src/worker/ingestion/event-pipeline/populateTeamDataStep.ts
index a436689062cc7..12a61887b1140 100644
--- a/plugin-server/src/worker/ingestion/event-pipeline/populateTeamDataStep.ts
+++ b/plugin-server/src/worker/ingestion/event-pipeline/populateTeamDataStep.ts
@@ -1,28 +1,69 @@
 import { PluginEvent } from '@posthog/plugin-scaffold'
+import { Counter } from 'prom-client'
 
+import { eventDroppedCounter } from '../../../main/ingestion-queues/metrics'
 import { PipelineEvent } from '../../../types'
 import { EventPipelineRunner } from './runner'
 
-/* 
-This step populates event.team_id and deletes event.ip if needed.
-If the event already has a team_id we will not run this step and 
-the capture endpoint will have handled this process. This is 
-temporary as this step will become the default for all events
-when we fully remove this functionality from the capture endpoint.
-*/
+export const inconsistentTeamCounter = new Counter({
+    name: 'ingestion_inconsistent_team_resolution_total',
+    help: 'Count of events with an inconsistent team resolution between capture and plugin-server, should be zero.',
+    labelNames: ['token', 'captured_team_id', 'resolved_team_id'],
+})
+
 export async function populateTeamDataStep(
     runner: EventPipelineRunner,
     event: PipelineEvent
 ): Promise<PluginEvent | null> {
+    /**
+     * Implements team_id resolution and applies the team's ingestion settings (dropping event.ip if requested).
+     *
+     * If the event already has a team_id field set by capture, it is used, but plugin-server still runs
+     * the resolution logic to confirm no inconsistency exists. Once team_id resolution is fully removed
+     * from capture, that section should be resolved, and team_id not trusted anymore.
+     */
+
+    // Events ingested with no token are dropped, they should be blocked by capture
     if (!event.token) {
+        eventDroppedCounter
+            .labels({
+                event_type: 'analytics',
+                drop_cause: 'no_token',
+            })
+            .inc()
         runner.hub.statsd?.increment('dropped_event_with_no_team', { token_set: 'false' })
         return null
     }
 
+    // Team lookup is cached, but will fail if PG is unavailable and the key expired.
+    // We should retry processing this event.
     const team = await runner.hub.teamManager.getTeamByToken(event.token)
 
-    // should we actually throw here?
+    // Short-circuit the logic if capture already resolved the team ID,
+    // just compare the lookup result to confirm data quality.
+    // TODO: remove after lightweight capture is fully rolled-out
+    if (event.team_id) {
+        if (team?.id || event.team_id) {
+            inconsistentTeamCounter
+                .labels({
+                    token: event.token,
+                    captured_team_id: event.team_id,
+                    resolved_team_id: team?.id,
+                })
+                .inc()
+            runner.hub.statsd?.increment('ingestion_inconsistent_team_resolution_total')
+        }
+        return event as PluginEvent
+    }
+
+    // If the token does not resolve to an existing team, drop the events.
     if (!team) {
+        eventDroppedCounter
+            .labels({
+                event_type: 'analytics',
+                drop_cause: 'invalid_token',
+            })
+            .inc()
         runner.hub.statsd?.increment('dropped_event_with_no_team', { token_set: 'true' })
         return null
     }
@@ -33,6 +74,5 @@ export async function populateTeamDataStep(
         ip: team.anonymize_ips ? null : event.ip,
     }
 
-    delete event['token']
     return event as PluginEvent
 }
diff --git a/plugin-server/src/worker/ingestion/team-manager.ts b/plugin-server/src/worker/ingestion/team-manager.ts
index cfa360cecc173..104bf6f99ff8a 100644
--- a/plugin-server/src/worker/ingestion/team-manager.ts
+++ b/plugin-server/src/worker/ingestion/team-manager.ts
@@ -20,15 +20,14 @@ export class TeamManager {
         this.statsd = statsd
 
         this.teamCache = new LRU({
-            max: 10000,
+            max: 10_000,
             maxAge: 2 * ONE_MINUTE,
-            // being explicit about the fact that we want to update
-            // the team cache every 2min, irrespective of the last access
-            updateAgeOnGet: false,
+            updateAgeOnGet: false, // Make default behaviour explicit
         })
         this.tokenToTeamIdCache = new LRU({
-            // TODO: add `maxAge` to ensure we avoid negatively caching teamId as null.
-            max: 100_000,
+            max: 1_000_000, // Entries are small, keep a high limit to reduce risk of bad requests evicting good tokens
+            maxAge: 5 * ONE_MINUTE, // Expiration for negative lookups, positive lookups will expire via teamCache first
+            updateAgeOnGet: false, // Make default behaviour explicit
         })
         this.instanceSiteUrl = serverConfig.SITE_URL || 'unknown'
     }
@@ -50,26 +49,41 @@ export class TeamManager {
     }
 
     public async getTeamByToken(token: string): Promise<Team | null> {
+        /**
+         * Validates and resolves the api token from an incoming event.
+         *
+         * Caching is added to reduce the load on Postgres, not to be resilient
+         * to failures. If PG is unavailable, this function will trow and the
+         * lookup must be retried later.
+         *
+         * Returns null if the token is invalid.
+         */
+
         const cachedTeamId = this.tokenToTeamIdCache.get(token)
 
-        // tokenToTeamIdCache.get returns `undefined` if the value doesn't
-        // exist so we check for the value being `null` as that means we've
-        // explictly cached that the team does not exist
+        // Negative lookups (`null` instead of `undefined`) return fast,
+        // but will be retried after that cache key expires.
+        // A new token can potentially get caught here for up to 5 minutes
+        // if a bad request in the past used that token.
         if (cachedTeamId === null) {
             return null
-        } else if (cachedTeamId) {
+        }
+
+        // Positive lookups hit both tokenToTeamIdCache and teamCache before returning without PG lookup.
+        // A revoked token will still be accepted until the teamCache entry expires (up to 2 minutes)
+        if (cachedTeamId) {
             const cachedTeam = this.teamCache.get(cachedTeamId)
             if (cachedTeam) {
                 return cachedTeam
             }
         }
 
-        const timeout = timeoutGuard(`Still running "fetchTeam". Timeout warning after 30 sec!`)
+        // Query PG if token is not in cache. This will throw if PG is unavailable.
+        const timeout = timeoutGuard(`Still running "fetchTeamByToken". Timeout warning after 30 sec!`)
         try {
             const team = await this.db.fetchTeamByToken(token)
             if (!team) {
-                // explicitly cache a null to avoid
-                // unnecessary lookups in the future
+                // Cache `null` for unknown tokens to reduce PG load, cache TTL will lead to retries later.
                 this.tokenToTeamIdCache.set(token, null)
                 return null
             }
diff --git a/plugin-server/tests/helpers/metrics.ts b/plugin-server/tests/helpers/metrics.ts
new file mode 100644
index 0000000000000..6182945f4d079
--- /dev/null
+++ b/plugin-server/tests/helpers/metrics.ts
@@ -0,0 +1,15 @@
+import { register } from 'prom-client'
+
+export function resetMetrics() {
+    register.resetMetrics()
+}
+
+export async function getMetricValues(metricName: string): Promise<any> {
+    const metrics = await register.getMetricsAsJSON()
+    for (const metric of metrics) {
+        if (metric.name === metricName) {
+            return (metric as any).values
+        }
+    }
+    throw Error(`Metric not found: ${metricName}`)
+}
diff --git a/plugin-server/tests/worker/ingestion/event-pipeline/populateTeamDataStep.test.ts b/plugin-server/tests/worker/ingestion/event-pipeline/populateTeamDataStep.test.ts
index a211ff2d9b346..738d477057823 100644
--- a/plugin-server/tests/worker/ingestion/event-pipeline/populateTeamDataStep.test.ts
+++ b/plugin-server/tests/worker/ingestion/event-pipeline/populateTeamDataStep.test.ts
@@ -1,6 +1,7 @@
 import { PipelineEvent, Team } from '../../../../src/types'
 import { UUIDT } from '../../../../src/utils/utils'
 import { populateTeamDataStep } from '../../../../src/worker/ingestion/event-pipeline/populateTeamDataStep'
+import { getMetricValues, resetMetrics } from '../../../helpers/metrics'
 
 const pipelineEvent: PipelineEvent = {
     event: '$pageview',
@@ -15,8 +16,6 @@ const pipelineEvent: PipelineEvent = {
     token: 'token',
 }
 
-const { token, ...defaultResultEvent } = pipelineEvent
-
 const team: Team = {
     id: 2,
     uuid: 'af95d312-1a0a-4208-b80f-562ddafc9bcd',
@@ -32,6 +31,7 @@ const team: Team = {
 let runner: any
 
 beforeEach(() => {
+    resetMetrics()
     runner = {
         nextStep: (...args: any[]) => args,
         hub: {
@@ -45,20 +45,62 @@ beforeEach(() => {
 describe('populateTeamDataStep()', () => {
     it('event with no token is not processed and the step returns null', async () => {
         const response = await populateTeamDataStep(runner, { ...pipelineEvent, team_id: undefined, token: undefined })
+        expect(response).toEqual(null)
+        expect(await getMetricValues('ingestion_event_dropped_total')).toEqual([
+            {
+                labels: {
+                    drop_cause: 'no_token',
+                    event_type: 'analytics',
+                },
+                value: 1,
+            },
+        ])
+    })
 
+    it('event with an invalid token is not processed and the step returns null', async () => {
+        jest.mocked(runner.hub.teamManager.getTeamByToken).mockReturnValue(null)
+        const response = await populateTeamDataStep(runner, { ...pipelineEvent, team_id: undefined, token: 'unknown' })
         expect(response).toEqual(null)
+        expect(await getMetricValues('ingestion_event_dropped_total')).toEqual([
+            {
+                labels: {
+                    drop_cause: 'invalid_token',
+                    event_type: 'analytics',
+                },
+                value: 1,
+            },
+        ])
     })
 
     it('event with a valid token gets assigned a team_id keeps its ip', async () => {
         const response = await populateTeamDataStep(runner, { ...pipelineEvent, team_id: undefined })
 
-        expect(response).toEqual({ ...defaultResultEvent, team_id: 2, ip: '127.0.0.1' })
+        expect(response).toEqual({ ...pipelineEvent, team_id: 2, ip: '127.0.0.1' })
+        expect(await getMetricValues('ingestion_event_dropped_total')).toEqual([])
     })
 
     it('event with a valid token for a team with anonymize_ips=true gets its ip set to null', async () => {
         jest.mocked(runner.hub.teamManager.getTeamByToken).mockReturnValue({ ...team, anonymize_ips: true })
         const response = await populateTeamDataStep(runner, { ...pipelineEvent, team_id: undefined })
 
-        expect(response).toEqual({ ...defaultResultEvent, team_id: 2, ip: null })
+        expect(response).toEqual({ ...pipelineEvent, team_id: 2, ip: null })
+        expect(await getMetricValues('ingestion_event_dropped_total')).toEqual([])
+    })
+
+    it('team_id from the event is preferred to the one returned by teamManager', async () => {
+        // Temporary behaviour while we progressively rollout this change: trust capture
+        jest.mocked(runner.hub.teamManager.getTeamByToken).mockReturnValue({ ...team, anonymize_ips: true })
+        const response = await populateTeamDataStep(runner, { ...pipelineEvent, team_id: 43, token: 'unknown' })
+        expect(response.team_id).toEqual(43)
+        expect(await getMetricValues('ingestion_inconsistent_team_resolution_total')).toEqual([
+            {
+                labels: {
+                    captured_team_id: 43,
+                    resolved_team_id: 2,
+                    token: 'unknown',
+                },
+                value: 1,
+            },
+        ])
     })
 })
diff --git a/plugin-server/tests/worker/ingestion/team-manager.test.ts b/plugin-server/tests/worker/ingestion/team-manager.test.ts
index a5ef94eb6edf7..ca8932ee00a04 100644
--- a/plugin-server/tests/worker/ingestion/team-manager.test.ts
+++ b/plugin-server/tests/worker/ingestion/team-manager.test.ts
@@ -61,4 +61,65 @@ describe('TeamManager()', () => {
             expect(await teamManager.fetchTeam(-1)).toEqual(null)
         })
     })
+
+    describe('getTeamByToken()', () => {
+        it('caches positive lookups for 2 minutes', async () => {
+            jest.spyOn(global.Date, 'now').mockImplementation(() => new Date('2020-02-27T11:00:05Z').getTime())
+            await hub.db.postgresQuery("UPDATE posthog_team SET api_token = 'my_token'", undefined, 'testTag')
+
+            // Initial lookup hits the DB and returns null
+            jest.spyOn(hub.db, 'postgresQuery')
+            let team = await teamManager.getTeamByToken('my_token')
+            expect(hub.db.postgresQuery).toHaveBeenCalledTimes(1)
+            expect(team!.id).toEqual(2)
+            expect(team!.anonymize_ips).toEqual(false)
+
+            // Settings are updated
+            await hub.db.postgresQuery('UPDATE posthog_team SET anonymize_ips = true', undefined, 'testTag')
+
+            // Second lookup hits the cache and skips the DB lookup, setting is stale
+            jest.spyOn(global.Date, 'now').mockImplementation(() => new Date('2020-02-27T11:01:56Z').getTime())
+            jest.mocked(hub.db.postgresQuery).mockClear()
+            team = await teamManager.getTeamByToken('my_token')
+            expect(hub.db.postgresQuery).toHaveBeenCalledTimes(0)
+            expect(team!.id).toEqual(2)
+            expect(team!.anonymize_ips).toEqual(false)
+
+            // Setting change take effect after cache expiration
+            jest.spyOn(global.Date, 'now').mockImplementation(() => new Date('2020-02-27T11:25:06Z').getTime())
+            jest.mocked(hub.db.postgresQuery).mockClear()
+            team = await teamManager.getTeamByToken('my_token')
+            expect(hub.db.postgresQuery).toHaveBeenCalledTimes(1)
+            expect(team!.id).toEqual(2)
+            expect(team!.anonymize_ips).toEqual(true)
+        })
+
+        it('caches negative lookups for 5 minutes', async () => {
+            jest.spyOn(global.Date, 'now').mockImplementation(() => new Date('2020-02-27T11:00:05Z').getTime())
+
+            // Initial lookup hits the DB and returns null
+            jest.spyOn(hub.db, 'postgresQuery')
+            expect(await teamManager.getTeamByToken('unknown')).toEqual(null)
+            expect(hub.db.postgresQuery).toHaveBeenCalledTimes(1)
+
+            // Second lookup hits the cache and skips the DB lookup
+            jest.spyOn(global.Date, 'now').mockImplementation(() => new Date('2020-02-27T11:03:06Z').getTime())
+            jest.mocked(hub.db.postgresQuery).mockClear()
+            expect(await teamManager.getTeamByToken('unknown')).toEqual(null)
+            expect(hub.db.postgresQuery).toHaveBeenCalledTimes(0)
+
+            // Hit the DB on cache expiration
+            jest.spyOn(global.Date, 'now').mockImplementation(() => new Date('2020-02-27T11:05:06Z').getTime())
+            jest.mocked(hub.db.postgresQuery).mockClear()
+            expect(await teamManager.getTeamByToken('unknown')).toEqual(null)
+            expect(hub.db.postgresQuery).toHaveBeenCalledTimes(1)
+        })
+
+        it('throws on postgres errors', async () => {
+            hub.db.postgresQuery = jest.fn().mockRejectedValue(new Error('PG unavailable'))
+            await expect(async () => {
+                await teamManager.getTeamByToken('another')
+            }).rejects.toThrow('PG unavailable')
+        })
+    })
 })


Code review comments:
diff_hunk: @@ -1,28 +1,69 @@
 import { PluginEvent } from '@posthog/plugin-scaffold'
+import { Counter } from 'prom-client'
 
+import { eventDroppedCounter } from '../../../main/ingestion-queues/metrics'
 import { PipelineEvent } from '../../../types'
 import { EventPipelineRunner } from './runner'
 
-/* 
-This step populates event.team_id and deletes event.ip if needed.
-If the event already has a team_id we will not run this step and 
-the capture endpoint will have handled this process. This is 
-temporary as this step will become the default for all events
-when we fully remove this functionality from the capture endpoint.
-*/
+export const inconsistentTeamCounter = new Counter({
+    name: 'ingestion_inconsistent_team_resolution_total',
+    help: 'Count of events with an inconsistent team resolution between capture and plugin-server, should be zero.',
+    labelNames: ['token', 'captured_team_id', 'resolved_team_id'],
+})
+
 export async function populateTeamDataStep(
     runner: EventPipelineRunner,
     event: PipelineEvent
 ): Promise<PluginEvent | null> {
+    /**
+     * Implements team_id resolution and applies the team's ingestion settings (dropping event.ip if requested).
+     *
+     * If the event already has a team_id field set by capture, it is used, but plugin-server still runs
+     * the resolution logic to confirm no inconsistency exists. Once team_id resolution is fully removed
+     * from capture, that section should be resolved, and team_id not trusted anymore.
+     */
+
+    // Events ingested with no token are dropped, they should be blocked by capture
     if (!event.token) {
+        eventDroppedCounter
+            .labels({
+                event_type: 'analytics',
+                drop_cause: 'no_token',
+            })
+            .inc()
         runner.hub.statsd?.increment('dropped_event_with_no_team', { token_set: 'false' })
         return null
     }
 
+    // Team lookup is cached, but will fail if PG is unavailable and the key expired.
+    // We should retry processing this event.
     const team = await runner.hub.teamManager.getTeamByToken(event.token)
 
-    // should we actually throw here?
+    // Short-circuit the logic if capture already resolved the team ID,
+    // just compare the lookup result to confirm data quality.
+    // TODO: remove after lightweight capture is fully rolled-out
+    if (event.team_id) {
+        if (team?.id || event.team_id) {
+            inconsistentTeamCounter
+                .labels({
path: plugin-server/src/worker/ingestion/event-pipeline/populateTeamDataStep.ts
text: @author:
High cardinality, but expect to have zero contexts (hopefully?).
We could play it safer and not add the labels for now, only add them if we find it's not zero?

@user1:
Assuming we expect to see nothing here and if we do, then we can revert this fast so seems fine 🤷‍♀️ 
What's the failure mode: we kill prometheus?

@author:
Indeed, failure mode is we kill prometheus and lose observability and alerting 😅 
Gonna remove the label just to be sure.