diff --git a/rust/personhog-replica/.sqlx/query-49bcc915e71cec48ea8e91fac80f2e4a004a39a064354ea5ef7dd6ae4e5afb78.json b/rust/personhog-replica/.sqlx/query-0ebe83ceec43b02189e46d1adbc9fc3a70c8b613c7b3bdfdcb053bb84534a57e.json similarity index 83% rename from rust/personhog-replica/.sqlx/query-49bcc915e71cec48ea8e91fac80f2e4a004a39a064354ea5ef7dd6ae4e5afb78.json rename to rust/personhog-replica/.sqlx/query-0ebe83ceec43b02189e46d1adbc9fc3a70c8b613c7b3bdfdcb053bb84534a57e.json index f9a13ecd40ef..52546426dc78 100644 --- a/rust/personhog-replica/.sqlx/query-49bcc915e71cec48ea8e91fac80f2e4a004a39a064354ea5ef7dd6ae4e5afb78.json +++ b/rust/personhog-replica/.sqlx/query-0ebe83ceec43b02189e46d1adbc9fc3a70c8b613c7b3bdfdcb053bb84534a57e.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n SELECT id::bigint as \"id!\", uuid as \"uuid!\", created_at as \"created_at!\"\n FROM posthog_person\n WHERE team_id = $1 AND uuid = ANY($2)\n FOR UPDATE\n ", + "query": "\n SELECT id::bigint as \"id!\", uuid as \"uuid!\", created_at as \"created_at!\"\n FROM posthog_person\n WHERE team_id = $1 AND uuid = ANY($2)\n ORDER BY id\n FOR UPDATE\n ", "describe": { "columns": [ { @@ -31,5 +31,5 @@ false ] }, - "hash": "49bcc915e71cec48ea8e91fac80f2e4a004a39a064354ea5ef7dd6ae4e5afb78" + "hash": "0ebe83ceec43b02189e46d1adbc9fc3a70c8b613c7b3bdfdcb053bb84534a57e" } diff --git a/rust/personhog-replica/.sqlx/query-373f3eccba1f2a9a6627543a0de528f8a52565944012c4d53908164a308870e4.json b/rust/personhog-replica/.sqlx/query-373f3eccba1f2a9a6627543a0de528f8a52565944012c4d53908164a308870e4.json new file mode 100644 index 000000000000..bddf5c61264c --- /dev/null +++ b/rust/personhog-replica/.sqlx/query-373f3eccba1f2a9a6627543a0de528f8a52565944012c4d53908164a308870e4.json @@ -0,0 +1,24 @@ +{ + "db_name": "PostgreSQL", + "query": "\n SELECT EXISTS (\n SELECT 1 FROM posthog_persondistinctid\n WHERE team_id = $1 AND person_id = $2 AND is_deleted = false AND NOT (distinct_id = ANY($3))\n ) as \"exists!\"\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "exists!", + "type_info": "Bool" + } + ], + "parameters": { + "Left": [ + "Int4", + "Int8", + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "373f3eccba1f2a9a6627543a0de528f8a52565944012c4d53908164a308870e4" +} diff --git a/rust/personhog-replica/.sqlx/query-138b31578ae24294373586142495b6fef5a1750e00e4797687861e4d0f1be82a.json b/rust/personhog-replica/.sqlx/query-459ab338d573c90e1cea2116fc4424ca22eabd036425f5a34f73479e5da56cb7.json similarity index 57% rename from rust/personhog-replica/.sqlx/query-138b31578ae24294373586142495b6fef5a1750e00e4797687861e4d0f1be82a.json rename to rust/personhog-replica/.sqlx/query-459ab338d573c90e1cea2116fc4424ca22eabd036425f5a34f73479e5da56cb7.json index b0642aac9ba6..978beb104c6b 100644 --- a/rust/personhog-replica/.sqlx/query-138b31578ae24294373586142495b6fef5a1750e00e4797687861e4d0f1be82a.json +++ b/rust/personhog-replica/.sqlx/query-459ab338d573c90e1cea2116fc4424ca22eabd036425f5a34f73479e5da56cb7.json @@ -1,15 +1,16 @@ { "db_name": "PostgreSQL", - "query": "\n DELETE FROM lifecycle_op\n WHERE op_id = $1 AND team_id = $2 AND op_type = 'delete'\n ", + "query": "\n DELETE FROM lifecycle_op\n WHERE op_id = $1 AND team_id = $2 AND op_type = $3\n ", "describe": { "columns": [], "parameters": { "Left": [ "Uuid", - "Int4" + "Int4", + "Text" ] }, "nullable": [] }, - "hash": "138b31578ae24294373586142495b6fef5a1750e00e4797687861e4d0f1be82a" + "hash": "459ab338d573c90e1cea2116fc4424ca22eabd036425f5a34f73479e5da56cb7" } diff --git a/rust/personhog-replica/.sqlx/query-96ae2fe1f4c10d259e4151924c80ea04ae5796a362f973adce0a925345e13392.json b/rust/personhog-replica/.sqlx/query-96ae2fe1f4c10d259e4151924c80ea04ae5796a362f973adce0a925345e13392.json new file mode 100644 index 000000000000..48098485a409 --- /dev/null +++ b/rust/personhog-replica/.sqlx/query-96ae2fe1f4c10d259e4151924c80ea04ae5796a362f973adce0a925345e13392.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "\n SELECT uuid as \"uuid!\"\n FROM posthog_person\n WHERE team_id = $1 AND id = $2\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "uuid!", + "type_info": "Uuid" + } + ], + "parameters": { + "Left": [ + "Int4", + "Int8" + ] + }, + "nullable": [ + false + ] + }, + "hash": "96ae2fe1f4c10d259e4151924c80ea04ae5796a362f973adce0a925345e13392" +} diff --git a/rust/personhog-replica/.sqlx/query-cd8cc89f4662b7ea42e9f7e58ef192aaf71583c5583b80a60d690f0386a8026a.json b/rust/personhog-replica/.sqlx/query-cd8cc89f4662b7ea42e9f7e58ef192aaf71583c5583b80a60d690f0386a8026a.json new file mode 100644 index 000000000000..dfe5f23521d1 --- /dev/null +++ b/rust/personhog-replica/.sqlx/query-cd8cc89f4662b7ea42e9f7e58ef192aaf71583c5583b80a60d690f0386a8026a.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "\n INSERT INTO lifecycle_op (op_id, op_type, team_id, step, request, completed_at)\n VALUES ($1, 'merge', $2, 'completed', '{\"source\": \"personhog-replica-split\"}'::jsonb, now())\n ", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Uuid", + "Int4" + ] + }, + "nullable": [] + }, + "hash": "cd8cc89f4662b7ea42e9f7e58ef192aaf71583c5583b80a60d690f0386a8026a" +} diff --git a/rust/personhog-replica/.sqlx/query-d536dcec71c04e794120461f6f655577b09524dbc0b751db57d27968cd23ef7f.json b/rust/personhog-replica/.sqlx/query-d536dcec71c04e794120461f6f655577b09524dbc0b751db57d27968cd23ef7f.json new file mode 100644 index 000000000000..84a12f37e2dd --- /dev/null +++ b/rust/personhog-replica/.sqlx/query-d536dcec71c04e794120461f6f655577b09524dbc0b751db57d27968cd23ef7f.json @@ -0,0 +1,25 @@ +{ + "db_name": "PostgreSQL", + "query": "\n INSERT INTO lifecycle_op_person (op_id, team_id, person_id, person_uuid, role, status, mark_active)\n VALUES ($1, $2, $3, $4, 'source', 'marked', true)\n ON CONFLICT (team_id, person_id) WHERE mark_active DO NOTHING\n RETURNING person_id\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "person_id", + "type_info": "Int8" + } + ], + "parameters": { + "Left": [ + "Uuid", + "Int4", + "Int8", + "Uuid" + ] + }, + "nullable": [ + false + ] + }, + "hash": "d536dcec71c04e794120461f6f655577b09524dbc0b751db57d27968cd23ef7f" +} diff --git a/rust/personhog-replica/src/service/mod.rs b/rust/personhog-replica/src/service/mod.rs index d97090f4161d..c925ae68d0b3 100644 --- a/rust/personhog-replica/src/service/mod.rs +++ b/rust/personhog-replica/src/service/mod.rs @@ -1495,9 +1495,6 @@ impl PersonHogReplica for PersonHogReplicaService { .await .map_err(|e| match &e { storage::StorageError::NotFound(msg) => Status::not_found(msg.clone()), - storage::StorageError::FailedPrecondition(msg) => { - Status::failed_precondition(msg.clone()) - } _ => log_and_convert_error(e, "split_person"), })?; diff --git a/rust/personhog-replica/src/storage/postgres/person.rs b/rust/personhog-replica/src/storage/postgres/person.rs index 8d90bfe2f009..ca4eb9f8614f 100644 --- a/rust/personhog-replica/src/storage/postgres/person.rs +++ b/rust/personhog-replica/src/storage/postgres/person.rs @@ -982,6 +982,14 @@ impl PersonLookup for PostgresStorage { // is constant (not per-distinct_id), so the lock window stays short; the // service-layer cap bounds the number of rows locked. let mut tx = self.bulk_primary_pool.begin().await?; + // Fail fast behind an in-flight merge, delete or split; the bound is per wait, not per call. + sqlx::query("SET LOCAL lock_timeout = '2s'") + .execute(&mut *tx) + .await?; + + // The source's mark serializes splits and merges of this person, so two splits cannot both pass the + // keep-one check below. Claim it before any row lock; the manual detach_distinct_id job takes none. + let mark_op_id = claim_split_mark(&mut tx, team_id, person_id).await?; // No FOR UPDATE on the source person: deletes lock PDI rows before person // rows, so locking the person first here would invert that order and risk @@ -1033,6 +1041,27 @@ impl PersonLookup for PostgresStorage { ))); } + // A fresh statement after the claim, so it sees every split that committed while this one waited. + // An id derived to the source's own uuid counts as leaving. + let keeps_a_distinct_id: bool = sqlx::query_scalar!( + r#" + SELECT EXISTS ( + SELECT 1 FROM posthog_persondistinctid + WHERE team_id = $1 AND person_id = $2 AND is_deleted = false AND NOT (distinct_id = ANY($3)) + ) as "exists!" + "#, + team_id as i32, + person_id, + distinct_ids_to_split + ) + .fetch_one(&mut *tx) + .await?; + if !keeps_a_distinct_id { + return Err(StorageError::FailedPrecondition(format!( + "the split would leave person_id={person_id} (team_id={team_id}) with no live distinct id" + ))); + } + let new_person_version = person_version + SPLIT_VERSION_OFFSET; // Build all arrays in request order so the response order matches the // input order (part of the RPC contract). @@ -1054,12 +1083,13 @@ impl PersonLookup for PostgresStorage { // re-split, or a tombstone left by a merge or delete of the person // that once owned the distinct id. No ON CONFLICT — the partitioned // table has a unique index, not a unique constraint, so ON CONFLICT - // inference doesn't work. + // inference doesn't work. Lock in id order, as the tombstone RPC does, so the two cannot deadlock. let existing_persons = sqlx::query!( r#" SELECT id::bigint as "id!", uuid as "uuid!", created_at as "created_at!" FROM posthog_person WHERE team_id = $1 AND uuid = ANY($2) + ORDER BY id FOR UPDATE "#, team_id as i32, @@ -1165,6 +1195,7 @@ impl PersonLookup for PostgresStorage { ))); } + release_lifecycle_marks(&mut tx, team_id, mark_op_id, "merge").await?; tx.commit().await?; let results: Vec = dids @@ -1618,7 +1649,7 @@ async fn tombstone_persons_by_uuids( tombstones.extend(with_tombstoned_distinct_ids(&mut tx, team_id, already).await?); if let Some(op_id) = mark_op_id { - release_delete_marks(&mut tx, team_id, op_id).await?; + release_lifecycle_marks(&mut tx, team_id, op_id, "delete").await?; } tx.commit().await?; tombstones.sort_by(|a, b| a.uuid.cmp(&b.uuid)); @@ -1692,21 +1723,83 @@ async fn claim_delete_marks( Ok((Some(op_id), person_ids)) } -async fn release_delete_marks( +/// Marks the split source as a 'merge' op, the type ingestion uses to move distinct ids. +async fn claim_split_mark( tx: &mut Transaction<'_, Postgres>, team_id: i64, - op_id: Uuid, -) -> StorageResult<()> { + person_id: i64, +) -> StorageResult { + let person_uuid = sqlx::query_scalar!( + r#" + SELECT uuid as "uuid!" + FROM posthog_person + WHERE team_id = $1 AND id = $2 + "#, + team_id as i32, + person_id + ) + .fetch_optional(&mut **tx) + .await? + .ok_or_else(|| StorageError::NotFound(format!("person_id={person_id} (team_id={team_id})")))?; + + let op_id = Uuid::new_v4(); sqlx::query!( r#" - DELETE FROM lifecycle_op - WHERE op_id = $1 AND team_id = $2 AND op_type = 'delete' + INSERT INTO lifecycle_op (op_id, op_type, team_id, step, request, completed_at) + VALUES ($1, 'merge', $2, 'completed', '{"source": "personhog-replica-split"}'::jsonb, now()) "#, op_id, team_id as i32 ) .execute(&mut **tx) .await?; + let claimed = sqlx::query_scalar!( + r#" + INSERT INTO lifecycle_op_person (op_id, team_id, person_id, person_uuid, role, status, mark_active) + VALUES ($1, $2, $3, $4, 'source', 'marked', true) + ON CONFLICT (team_id, person_id) WHERE mark_active DO NOTHING + RETURNING person_id + "#, + op_id, + team_id as i32, + person_id, + person_uuid + ) + .fetch_optional(&mut **tx) + .await?; + if claimed.is_none() { + return Err(StorageError::FailedPrecondition(format!( + "person_id={person_id} (team_id={team_id}) is claimed by another lifecycle operation" + ))); + } + Ok(op_id) +} + +/// Deletes the op inside the claiming transaction, so committed state never holds its marks. +async fn release_lifecycle_marks( + tx: &mut Transaction<'_, Postgres>, + team_id: i64, + op_id: Uuid, + op_type: &str, +) -> StorageResult<()> { + let deleted = sqlx::query!( + r#" + DELETE FROM lifecycle_op + WHERE op_id = $1 AND team_id = $2 AND op_type = $3 + "#, + op_id, + team_id as i32, + op_type + ) + .execute(&mut **tx) + .await? + .rows_affected(); + // Committing without the delete would leave the marks held for good. + if deleted != 1 { + return Err(StorageError::Query(format!( + "releasing lifecycle op {op_id} deleted {deleted} rows (team_id={team_id})" + ))); + } Ok(()) } diff --git a/rust/personhog-replica/tests/storage_tests.rs b/rust/personhog-replica/tests/storage_tests.rs index 45205de5abbf..752a3e97a89d 100644 --- a/rust/personhog-replica/tests/storage_tests.rs +++ b/rust/personhog-replica/tests/storage_tests.rs @@ -620,10 +620,7 @@ async fn wait_for_a_mark_claim_blocked_by(pool: &sqlx::PgPool, holder_pid: i32) if waiting > 0 { return; } - assert!( - Instant::now() < deadline, - "the delete never waited on the held mark" - ); + assert!(Instant::now() < deadline, "nothing waited on the held mark"); tokio::time::sleep(Duration::from_millis(10)).await; } } @@ -2871,6 +2868,168 @@ async fn test_split_person_creates_new_persons(#[case] num_splits: usize) { .expect("Lookup should succeed") .expect("Keeper DID should still resolve to original"); assert_eq!(original_by_did.id, person.id); + assert_eq!(lifecycle_op_count(&ctx.pool, ctx.team_id).await, 0); + + ctx.cleanup().await.ok(); +} + +async fn live_distinct_ids(pool: &sqlx::PgPool, team_id: i64, person_id: i64) -> Vec { + sqlx::query_scalar( + "SELECT distinct_id FROM posthog_persondistinctid + WHERE team_id = $1 AND person_id = $2 AND NOT is_deleted ORDER BY distinct_id", + ) + .bind(team_id as i32) + .bind(person_id) + .fetch_all(pool) + .await + .unwrap() +} + +fn assert_failed_precondition( + result: Result, +) { + assert!( + matches!( + result, + Err(personhog_replica::storage::StorageError::FailedPrecondition(_)) + ), + "expected FailedPrecondition, got {result:?}" + ); +} + +#[tokio::test] +async fn test_split_person_rejects_moving_every_distinct_id() { + let ctx = TestContext::new().await; + let person = ctx.insert_person("split_all_a", None).await.unwrap(); + ctx.add_distinct_id_to_person(person.id, "split_all_b") + .await + .unwrap(); + + let result = ctx + .storage + .split_person( + ctx.team_id, + person.id, + &["split_all_a".to_string(), "split_all_b".to_string()], + ) + .await; + + assert_failed_precondition(result); + assert_eq!( + live_distinct_ids(&ctx.pool, ctx.team_id, person.id).await, + vec!["split_all_a", "split_all_b"] + ); + assert_eq!(lifecycle_op_count(&ctx.pool, ctx.team_id).await, 0); + + ctx.cleanup().await.ok(); +} + +#[tokio::test] +async fn test_split_person_waits_for_a_concurrent_move_and_keeps_the_last_distinct_id() { + let ctx = TestContext::new().await; + let person = ctx.insert_person("race_keep_a", None).await.unwrap(); + ctx.add_distinct_id_to_person(person.id, "race_keep_b") + .await + .unwrap(); + let other = ctx.insert_person("race_other", None).await.unwrap(); + + // A concurrent split or merge, in flight: it holds the person's mark and has moved "a" away. + let mut mover = ctx.pool.begin().await.unwrap(); + let mover_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *mover) + .await + .unwrap(); + let op_id = claim_merge_mark(&mut mover, ctx.team_id, &person).await; + sqlx::query( + "UPDATE posthog_persondistinctid SET person_id = $1, version = version + 1 + WHERE team_id = $2 AND distinct_id = 'race_keep_a'", + ) + .bind(other.id) + .bind(ctx.team_id as i32) + .execute(&mut *mover) + .await + .unwrap(); + sqlx::query("DELETE FROM lifecycle_op WHERE op_id = $1") + .bind(op_id) + .execute(&mut *mover) + .await + .unwrap(); + + let storage = ctx.storage.clone(); + let team_id = ctx.team_id; + let person_id = person.id; + let split = tokio::spawn(async move { + storage + .split_person(team_id, person_id, &["race_keep_b".to_string()]) + .await + }); + + // Commit the move only once the split waits on the mark, so the split cannot finish first. + wait_for_a_mark_claim_blocked_by(&ctx.pool, mover_pid).await; + mover.commit().await.unwrap(); + + assert_failed_precondition(split.await.unwrap()); + assert_eq!( + live_distinct_ids(&ctx.pool, ctx.team_id, person.id).await, + vec!["race_keep_b"] + ); + assert_eq!(lifecycle_op_count(&ctx.pool, ctx.team_id).await, 0); + + ctx.cleanup().await.ok(); +} + +#[tokio::test] +async fn test_split_person_times_out_and_rolls_back_behind_a_long_held_mark() { + let ctx = TestContext::new().await; + let person = ctx.insert_person("split_slow_a", None).await.unwrap(); + ctx.add_distinct_id_to_person(person.id, "split_slow_b") + .await + .unwrap(); + let mut holder = ctx.pool.begin().await.unwrap(); + claim_merge_mark(&mut holder, ctx.team_id, &person).await; + + let result = ctx + .storage + .split_person(ctx.team_id, person.id, &["split_slow_b".to_string()]) + .await; + holder.rollback().await.unwrap(); + + match result { + Err(err) => assert_eq!(err.error_type(), "lock_timeout"), + Ok(results) => panic!("expected a lock timeout, got {results:?}"), + } + assert_eq!( + live_distinct_ids(&ctx.pool, ctx.team_id, person.id).await, + vec!["split_slow_a", "split_slow_b"] + ); + assert_eq!(lifecycle_op_count(&ctx.pool, ctx.team_id).await, 0); + + ctx.cleanup().await.ok(); +} + +#[tokio::test] +async fn test_split_person_fails_when_a_lifecycle_op_holds_the_person() { + let ctx = TestContext::new().await; + let person = ctx.insert_person("split_held_a", None).await.unwrap(); + ctx.add_distinct_id_to_person(person.id, "split_held_b") + .await + .unwrap(); + let mut conn = ctx.pool.acquire().await.unwrap(); + claim_merge_mark(&mut conn, ctx.team_id, &person).await; + drop(conn); + + let result = ctx + .storage + .split_person(ctx.team_id, person.id, &["split_held_b".to_string()]) + .await; + + assert_failed_precondition(result); + assert_eq!( + live_distinct_ids(&ctx.pool, ctx.team_id, person.id).await, + vec!["split_held_a", "split_held_b"] + ); + // Only the holder's op: the split's own claim rolled back. + assert_eq!(lifecycle_op_count(&ctx.pool, ctx.team_id).await, 1); ctx.cleanup().await.ok(); } @@ -2898,9 +3057,9 @@ async fn test_split_person_version_arithmetic() { .await .unwrap(); - // Insert PDIs with non-zero versions + // Insert PDIs with non-zero versions; the split must leave one on the person. let pdi_version: i64 = 7; - for did in &["v_did_a", "v_did_b"] { + for did in &["v_did_a", "v_did_b", "v_did_keep"] { sqlx::query( r#"INSERT INTO posthog_persondistinctid (distinct_id, person_id, team_id, version) VALUES ($1, $2, $3, $4)"#,