From 238bc09bed57dcd0e4887615f3c57a580eb26d9e Mon Sep 17 00:00:00 2001
From: Alvaro Herrera <alvherre@alvh.no-ip.org>
Date: Mon, 22 Apr 2024 11:32:04 +0200
Subject: [PATCH v2 1/2] Acquire locks on children before recursing

ALTER TABLE ADD CONSTRAINT was missing this, as evidenced by assertion
failures.  While at it, create a small routine to encapsulate the weird
find_all_inheritors() call.

Reported-by: Alexander Lakhin <exclusion@gmail.com>
Discussion: https://postgr.es/m/5b4cd32f-1d5b-c080-c688-31736bbcd739@gmail.com
---
 src/backend/commands/tablecmds.c | 40 +++++++++++++++++++++-----------
 1 file changed, 27 insertions(+), 13 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 3556240c8e..9058a0de7a 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -427,6 +427,7 @@ static void ATSimplePermissions(AlterTableType cmdtype, Relation rel, int allowe
 static void ATSimpleRecursion(List **wqueue, Relation rel,
 							  AlterTableCmd *cmd, bool recurse, LOCKMODE lockmode,
 							  AlterTableUtilityContext *context);
+static void ATLockAllDescendants(Oid relid, LOCKMODE lockmode);
 static void ATCheckPartitionsNotInUse(Relation rel, LOCKMODE lockmode);
 static void ATTypedTableRecursion(List **wqueue, Relation rel, AlterTableCmd *cmd,
 								  LOCKMODE lockmode,
@@ -1621,9 +1622,7 @@ RemoveRelations(DropStmt *drop)
 		 * will lock those objects in the other order.
 		 */
 		if (state.actual_relkind == RELKIND_PARTITIONED_INDEX)
-			(void) find_all_inheritors(state.heapOid,
-									   state.heap_lockmode,
-									   NULL);
+			ATLockAllDescendants(state.heapOid, state.heap_lockmode);
 
 		/* OK, we're ready to delete this one */
 		obj.classId = RelationRelationId;
@@ -4979,10 +4978,12 @@ ATPrepCmd(List **wqueue, Relation rel, AlterTableCmd *cmd,
 			break;
 		case AT_AddConstraint:	/* ADD CONSTRAINT */
 			ATSimplePermissions(cmd->subtype, rel, ATT_TABLE | ATT_FOREIGN_TABLE);
-			/* Recursion occurs during execution phase */
-			/* No command-specific prep needed except saving recurse flag */
 			if (recurse)
+			{
+				/* if recursing, set flag and lock all descendants */
 				cmd->recurse = true;
+				ATLockAllDescendants(RelationGetRelid(rel), lockmode);
+			}
 			pass = AT_PASS_ADD_CONSTR;
 			break;
 		case AT_AddIndexConstraint: /* ADD CONSTRAINT USING INDEX */
@@ -6721,6 +6722,21 @@ ATSimpleRecursion(List **wqueue, Relation rel,
 	}
 }
 
+/*
+ * ATLockAllDescendants
+ *
+ * Acquire lock on all descendant relations of the given relation.
+ */
+static void
+ATLockAllDescendants(Oid relid, LOCKMODE lockmode)
+{
+	/*
+	 * This is only used in DDL code, so it doesn't matter that we leak the
+	 * list storage; it'll be gone soon enough.
+	 */
+	(void) find_all_inheritors(relid, lockmode, NULL);
+}
+
 /*
  * Obtain list of partitions of the given table, locking them all at the given
  * lockmode and ensuring that they all pass CheckTableNotInUse.
@@ -9370,10 +9386,9 @@ ATPrepAddPrimaryKey(List **wqueue, Relation rel, AlterTableCmd *cmd,
 
 	/*
 	 * Acquire locks all the way down the hierarchy.  The recursion to lower
-	 * levels occurs at execution time as necessary, so we don't need to do it
-	 * here, and we don't need the returned list either.
+	 * levels occurs at execution time as necessary.
 	 */
-	(void) find_all_inheritors(RelationGetRelid(rel), lockmode, NULL);
+	ATLockAllDescendants(RelationGetRelid(rel), lockmode);
 
 	/*
 	 * Construct the list of constraints that we need to add to each child
@@ -9819,10 +9834,10 @@ ATAddCheckNNConstraint(List **wqueue, AlteredTableInfo *tab, Relation rel,
 	/*
 	 * Propagate to children as appropriate.  Unlike most other ALTER
 	 * routines, we have to do this one level of recursion at a time; we can't
-	 * use find_all_inheritors to do it in one pass.
+	 * use find_all_inheritors to do it in one pass.  We have all locks
+	 * already, however.
 	 */
-	children =
-		find_inheritance_children(RelationGetRelid(rel), lockmode);
+	children = find_inheritance_children(RelationGetRelid(rel), NoLock);
 
 	/*
 	 * Check if ONLY was specified with ALTER TABLE.  If so, allow the
@@ -11258,8 +11273,7 @@ CloneFkReferencing(List **wqueue, Relation parentRel, Relation partRel)
 		 */
 		pkrel = table_open(constrForm->confrelid, ShareRowExclusiveLock);
 		if (pkrel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
-			(void) find_all_inheritors(RelationGetRelid(pkrel),
-									   ShareRowExclusiveLock, NULL);
+			ATLockAllDescendants(RelationGetRelid(pkrel), ShareRowExclusiveLock);
 
 		DeconstructFkConstraintRow(tuple, &numfks, conkey, confkey,
 								   conpfeqop, conppeqop, conffeqop,
-- 
2.39.2