From 03a15471b0f5a6776331d430cca2df619047271b Mon Sep 17 00:00:00 2001
From: Corey Huinker <corey.huinker@gmail.com>
Date: Mon, 20 Jan 2025 16:18:51 -0500
Subject: [PATCH v40] Lock table first when setting index relation statistics.

Jian He reported [1] a missing lock relation bug in
pg_restore_relation_stats when restoring stats to an index.

This fix follows the steps for proper locking prior to an inplace update
of a relation as specified in aac2c9b4fde8, and mimics the locking
behavior of the analyze command.

Confirmed that test case generates the reported error without the code
fix.

[1] https://www.postgresql.org/message-id/CACJufxGreTY7qsCV8%2BBkuv0p5SXGTScgh%3DD%2BDq6%3D%2B_%3DXTp7FWg%40mail.gmail.com
---
 src/backend/statistics/stat_utils.c        | 40 ++++++++++++++++++++--
 src/test/regress/expected/stats_import.out | 24 +++++++++++++
 src/test/regress/sql/stats_import.sql      | 18 ++++++++++
 3 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/src/backend/statistics/stat_utils.c b/src/backend/statistics/stat_utils.c
index 0d446f55b0..dc510fe253 100644
--- a/src/backend/statistics/stat_utils.c
+++ b/src/backend/statistics/stat_utils.c
@@ -17,14 +17,19 @@
 #include "postgres.h"
 
 #include "access/relation.h"
+#include "catalog/pg_class_d.h"
 #include "catalog/pg_database.h"
+#include "catalog/index.h"
 #include "funcapi.h"
 #include "miscadmin.h"
 #include "statistics/stat_utils.h"
+#include "storage/lockdefs.h"
 #include "utils/acl.h"
 #include "utils/array.h"
 #include "utils/builtins.h"
+#include "utils/lsyscache.h"
 #include "utils/rel.h"
+#include "utils/syscache.h"
 
 /*
  * Ensure that a given argument is not null.
@@ -126,8 +131,39 @@ stats_check_arg_pair(FunctionCallInfo fcinfo,
 void
 stats_lock_check_privileges(Oid reloid)
 {
-	Relation	rel = relation_open(reloid, ShareUpdateExclusiveLock);
-	const char	relkind = rel->rd_rel->relkind;
+	const char	relkind = get_rel_relkind(reloid);
+	Relation	rel;
+	Oid			indrelid = InvalidOid;
+	LOCKMODE	lockmode = ShareUpdateExclusiveLock;
+
+	/*
+	 * For indexes we follow what do_analyze_rel() does so as to avoid any
+	 * deadlocks with analyze/vacuum, with one exception, and that is
+	 * that we will also lock partitioned indexes because if we encounter
+	 * one then we are about to create pg_statistic rows that reference it.
+	 */
+	/* if (relkind == RELKIND_INDEX || relkind == RELKIND_PARTITIONED_INDEX) */
+	if (relkind == RELKIND_INDEX)
+	{
+		Relation	parentrel;
+
+		/*
+		 * Lock the table instead, and get a lesser lock on the index itself.
+		 */
+		indrelid = IndexGetRelation(reloid, false);
+
+		parentrel = relation_open(indrelid, ShareUpdateExclusiveLock);
+
+		relation_close(parentrel, NoLock);
+
+		lockmode = AccessShareLock;
+	}
+
+	rel = relation_open(reloid, lockmode);
+
+	if (OidIsValid(indrelid) &&
+		(indrelid != IndexGetRelation(rel->rd_rel->oid, false)))
+		elog(ERROR, "index parent oid recheck failed for index %u", reloid);
 
 	/* All of the types that can be used with ANALYZE, plus indexes */
 	switch (relkind)
diff --git a/src/test/regress/expected/stats_import.out b/src/test/regress/expected/stats_import.out
index fb50da1cd8..be283b4035 100644
--- a/src/test/regress/expected/stats_import.out
+++ b/src/test/regress/expected/stats_import.out
@@ -182,6 +182,7 @@ CREATE TABLE stats_import.part_child_1
   PARTITION OF stats_import.part_parent
   FOR VALUES FROM (0) TO (10)
   WITH (autovacuum_enabled = false);
+CREATE INDEX part_parent_i ON stats_import.part_parent(i);
 ANALYZE stats_import.part_parent;
 SELECT relpages
 FROM pg_class
@@ -202,6 +203,16 @@ SELECT
  
 (1 row)
 
+-- Check locking of partitioned index
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.part_parent_i'::regclass,
+        relpages => 2::integer);
+ pg_set_relation_stats 
+-----------------------
+ 
+(1 row)
+
 -- nothing stops us from setting it to -1
 SELECT
     pg_catalog.pg_set_relation_stats(
@@ -1414,6 +1425,19 @@ SELECT 3, 'tre', (3, 3.3, 'TRE', '2003-03-03', NULL)::stats_import.complex_type,
 UNION ALL
 SELECT 4, 'four', NULL, int4range(0,100), NULL;
 CREATE INDEX is_odd ON stats_import.test(((comp).a % 2 = 1));
+/* Test for proper locking */
+SELECT * FROM pg_catalog.pg_restore_relation_stats(
+    'relation', 'stats_import.is_odd'::regclass,
+    'version', '180000'::integer,
+    'relpages', '11'::integer,
+    'reltuples', '10000'::real,
+    'relallvisible', '0'::integer
+);
+ pg_restore_relation_stats 
+---------------------------
+ t
+(1 row)
+
 -- Generate statistics on table with data
 ANALYZE stats_import.test;
 CREATE TABLE stats_import.test_clone ( LIKE stats_import.test )
diff --git a/src/test/regress/sql/stats_import.sql b/src/test/regress/sql/stats_import.sql
index d3058bf8f6..298ac90fb3 100644
--- a/src/test/regress/sql/stats_import.sql
+++ b/src/test/regress/sql/stats_import.sql
@@ -127,6 +127,8 @@ CREATE TABLE stats_import.part_child_1
   FOR VALUES FROM (0) TO (10)
   WITH (autovacuum_enabled = false);
 
+CREATE INDEX part_parent_i ON stats_import.part_parent(i);
+
 ANALYZE stats_import.part_parent;
 
 SELECT relpages
@@ -140,6 +142,12 @@ SELECT
         relation => 'stats_import.part_parent'::regclass,
         relpages => 2::integer);
 
+-- Check locking of partitioned index
+SELECT
+    pg_catalog.pg_set_relation_stats(
+        relation => 'stats_import.part_parent_i'::regclass,
+        relpages => 2::integer);
+
 -- nothing stops us from setting it to -1
 SELECT
     pg_catalog.pg_set_relation_stats(
@@ -1062,6 +1070,16 @@ SELECT 4, 'four', NULL, int4range(0,100), NULL;
 
 CREATE INDEX is_odd ON stats_import.test(((comp).a % 2 = 1));
 
+
+/* Test for proper locking */
+SELECT * FROM pg_catalog.pg_restore_relation_stats(
+    'relation', 'stats_import.is_odd'::regclass,
+    'version', '180000'::integer,
+    'relpages', '11'::integer,
+    'reltuples', '10000'::real,
+    'relallvisible', '0'::integer
+);
+
 -- Generate statistics on table with data
 ANALYZE stats_import.test;
 
-- 
2.48.1

