fix(projects): retry writes after concurrent membership changes

Author: mzxchandra

packages/db/migrations/0394_project_membership_enforcement.sql

@@ -142,8 +142,15 @@ BEGIN
       IF TG_OP <> 'INSERT' THEN previous_id := OLD.workspace_id; END IF;
       IF TG_OP <> 'DELETE' THEN next_id := NEW.workspace_id; END IF;
     END IF;
-    SELECT array_agg(project_id) INTO owners FROM project_workspace WHERE workspace_id IN (previous_id, next_id);
+    SELECT array_agg(project_id ORDER BY workspace_id) INTO owners FROM project_workspace WHERE workspace_id IN (previous_id, next_id);
     PERFORM project_contract_lock_projects(owners);
+    IF owners IS DISTINCT FROM (
+      SELECT array_agg(project_id ORDER BY workspace_id) FROM project_workspace
+      WHERE workspace_id IN (previous_id, next_id)
+    ) THEN
+      RAISE EXCEPTION 'Environment changed Projects while acquiring its lifecycle lock; retry the transaction'
+        USING ERRCODE = '40001';
+    END IF;
   END IF;
   IF TG_OP = 'DELETE' THEN RETURN OLD; END IF;
   RETURN NEW;

packages/db/scripts/project-contract.integration.ts

@@ -8,6 +8,7 @@ import { backfillProjects } from '@sim/db/project-backfill'
 import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure'
 import { createDeferred } from '@sim/testing/helpers/deferred'
 import { getPostgresErrorCode } from '@sim/utils/errors'
+import { sleep } from '@sim/utils/helpers'
 import { generateId } from '@sim/utils/id'
 import { drizzle } from 'drizzle-orm/postgres-js'
 import { migrate } from 'drizzle-orm/postgres-js/migrator'
@@ -220,6 +221,52 @@ describe('Project expand/backfill/contract against PostgreSQL', () => {
     })
   })
 
+  it('retries a workflow write when its environment moves Projects while the writer waits', async () => {
+    await database(async (sql) => {
+      await seed(sql)
+      await enforce(sql)
+      const moved = createDeferred<void>()
+      const release = createDeferred<void>()
+      const transfer = sql.begin(async (tx) => {
+        await tx`INSERT INTO project (id, name, owner_id) VALUES ('destination', 'Destination', 'owner')`
+        await tx`UPDATE workspace SET forked_from_workspace_id = NULL WHERE id = 'child'`
+        await tx`UPDATE project_workspace SET project_id = 'destination' WHERE workspace_id = 'child'`
+        moved.resolve()
+        await release.promise
+      })
+      await moved.promise
+      const writer = sql
+        .begin(async (tx) => {
+          await tx`SET LOCAL application_name = 'project_contract_stale_writer'`
+          await tx`INSERT INTO workflow VALUES ('racing', 'child', NULL)`
+        })
+        .then(
+          () => null,
+          (error: unknown) => error
+        )
+      try {
+        let waiting = false
+        for (let attempt = 0; attempt < 100; attempt++) {
+          const rows =
+            await sql`SELECT 1 FROM pg_stat_activity WHERE application_name = 'project_contract_stale_writer' AND wait_event_type = 'Lock'`
+          if (rows.length) {
+            waiting = true
+            break
+          }
+          await sleep(10)
+        }
+        expect(waiting).toBe(true)
+      } finally {
+        release.resolve()
+      }
+      await transfer
+      expect(getPostgresErrorCode(await writer)).toBe('40001')
+      expect(await sql`SELECT 1 FROM workflow WHERE id = 'racing'`).toHaveLength(0)
+      await sql`INSERT INTO workflow VALUES ('retried', 'child', NULL)`
+      expect(await sql`SELECT 1 FROM workflow WHERE id = 'retried'`).toHaveLength(1)
+    })
+  })
+
   it.each(['read committed', 'repeatable read'] as const)(
     'prevents concurrent last-environment removal under %s',
     async (isolation) => {