Acquire AccessShareLock before updating table statistics (#5155)

src/backend/distributed/metadata/metadata_utility.c

@@ -893,11 +893,22 @@ GenerateAllShardStatisticsQueryForNode(WorkerNode *workerNode, List *citusTableI
 	Oid relationId = InvalidOid;
 	foreach_oid(relationId, citusTableIds)
 	{
-		List *shardIntervalsOnNode = ShardIntervalsOnWorkerGroup(workerNode, relationId);
+		/*
-		char *shardStatisticsQuery =
+		 * Ensure the table still exists by trying to acquire a lock on it
-			GenerateShardStatisticsQueryForShardList(shardIntervalsOnNode,
+		 * If function returns NULL, it means the table doesn't exist
-													 useShardMinMaxQuery);
+		 * hence we should skip
-		appendStringInfoString(allShardStatisticsQuery, shardStatisticsQuery);
+		 */
+		Relation relation = try_relation_open(relationId, AccessShareLock);
+		if (relation != NULL)
+		{
+			List *shardIntervalsOnNode = ShardIntervalsOnWorkerGroup(workerNode,
+																	 relationId);
+			char *shardStatisticsQuery =
+				GenerateShardStatisticsQueryForShardList(shardIntervalsOnNode,
+														 useShardMinMaxQuery);
+			appendStringInfoString(allShardStatisticsQuery, shardStatisticsQuery);
+			relation_close(relation, AccessShareLock);
+		}
 	}
 
 	/* Add a dummy entry so that UNION ALL doesn't complain */

src/backend/distributed/operations/stage_protocol.c

@@ -380,7 +380,32 @@ citus_update_table_statistics(PG_FUNCTION_ARGS)
 
 	Oid distributedTableId = PG_GETARG_OID(0);
 
-	UpdateTableStatistics(distributedTableId);
+	/*
+	 * Ensure the table still exists by trying to acquire a lock on it
+	 * If function returns NULL, it means the table doesn't exist
+	 * hence we should skip
+	 */
+	Relation relation = try_relation_open(distributedTableId, AccessShareLock);
+
+	if (relation != NULL)
+	{
+		UpdateTableStatistics(distributedTableId);
+
+		/*
+		 * We release the lock here since citus_update_table_statistics
+		 * is usually used in the following command:
+		 * SELECT citus_update_table_statistics(logicalrelid) FROM pg_dist_partition;
+		 * In this way we avoid holding the locks on distributed tables for a long time:
+		 * If we close the relation with NoLock, the locks on the distributed tables will
+		 * be held until the above command is finished (all distributed tables are updated).
+		 */
+		relation_close(relation, AccessShareLock);
+	}
+	else
+	{
+		ereport(NOTICE, (errmsg("relation with OID %u does not exist, skipping",
+								distributedTableId)));
+	}
 
 	PG_RETURN_VOID();
 }

src/test/regress/expected/isolation_citus_update_table_statistics.out

@@ -0,0 +1,90 @@
+Parsed test spec with 2 sessions
+
+starting permutation: s1-begin s1-drop-table s2-citus-update-table-statistics s1-commit
+create_distributed_table
+
+
+step s1-begin:
+    BEGIN;
+
+step s1-drop-table:
+    DROP TABLE dist_table;
+
+step s2-citus-update-table-statistics:
+    SET client_min_messages TO NOTICE;
+    SELECT citus_update_table_statistics(logicalrelid) FROM pg_dist_partition;
+ <waiting ...>
+step s1-commit: 
+    COMMIT;
+
+step s2-citus-update-table-statistics: <... completed>
+s2: NOTICE:  relation with OID XXXX does not exist, skipping
+citus_update_table_statistics
+
+
+
+
+starting permutation: s1-begin s1-drop-table s2-citus-shards s1-commit
+create_distributed_table
+
+
+step s1-begin:
+    BEGIN;
+
+step s1-drop-table:
+    DROP TABLE dist_table;
+
+step s2-citus-shards:
+    SELECT 1 AS result FROM citus_shards GROUP BY result;
+ <waiting ...>
+step s1-commit: 
+    COMMIT;
+
+step s2-citus-shards: <... completed>
+result
+
+1
+
+starting permutation: s2-begin s2-citus-shards s1-drop-table s2-commit
+create_distributed_table
+
+
+step s2-begin:
+    BEGIN;
+
+step s2-citus-shards:
+    SELECT 1 AS result FROM citus_shards GROUP BY result;
+
+result
+
+1
+step s1-drop-table:
+    DROP TABLE dist_table;
+
+step s2-commit:
+    COMMIT;
+
+
+starting permutation: s2-begin s2-citus-update-table-statistics s1-drop-table s2-commit
+create_distributed_table
+
+
+step s2-begin:
+    BEGIN;
+
+step s2-citus-update-table-statistics:
+    SET client_min_messages TO NOTICE;
+    SELECT citus_update_table_statistics(logicalrelid) FROM pg_dist_partition;
+
+citus_update_table_statistics
+
+
+
+step s1-drop-table:
+    DROP TABLE dist_table;
+ <waiting ...>
+step s2-commit: 
+    COMMIT;
+
+step s1-drop-table: <... completed>
+error in steps s2-commit s1-drop-table: ERROR:  tuple concurrently deleted

src/test/regress/isolation_schedule

@@ -65,6 +65,7 @@ test: isolation_multiuser_locking
 test: shared_connection_waits
 test: isolation_cancellation
 test: isolation_undistribute_table
+test: isolation_citus_update_table_statistics
 
 # Rebalancer
 test: isolation_blocking_move_single_shard_commands

src/test/regress/spec/isolation_citus_update_table_statistics.spec

@@ -0,0 +1,59 @@
+setup
+{
+    CREATE TABLE dist_table(a INT, b INT);
+    SELECT create_distributed_table('dist_table', 'a');
+}
+
+teardown
+{
+    DROP TABLE IF EXISTS dist_table;
+}
+
+session "s1"
+
+step "s1-begin"
+{
+    BEGIN;
+}
+
+step "s1-drop-table"
+{
+    DROP TABLE dist_table;
+}
+
+step "s1-commit"
+{
+    COMMIT;
+}
+
+session "s2"
+
+step "s2-begin"
+{
+    BEGIN;
+}
+
+step "s2-citus-update-table-statistics"
+{
+    SET client_min_messages TO NOTICE;
+    SELECT citus_update_table_statistics(logicalrelid) FROM pg_dist_partition;
+}
+
+step "s2-citus-shards"
+{
+    SELECT 1 AS result FROM citus_shards GROUP BY result;
+}
+
+step "s2-commit"
+{
+    COMMIT;
+}
+
+permutation "s1-begin" "s1-drop-table" "s2-citus-update-table-statistics" "s1-commit"
+permutation "s1-begin" "s1-drop-table" "s2-citus-shards" "s1-commit"
+permutation "s2-begin" "s2-citus-shards" "s1-drop-table" "s2-commit"
+
+// ERROR:  tuple concurrently deleted -- is expected in the following permutation
+// Check the explanation at PR #5155 in the following comment
+// https://github.com/citusdata/citus/pull/5155#issuecomment-897028194
+permutation "s2-begin" "s2-citus-update-table-statistics" "s1-drop-table" "s2-commit"