From 4cad81d643df66e11c44e7b3e10edfc0748cd086 Mon Sep 17 00:00:00 2001
From: Onur Tirtir <onurcantirtir@gmail.com>
Date: Fri, 31 Jan 2025 14:59:23 +0300
Subject: [PATCH] Fix flaky citus upgrade test

---
 .../upgrade_pg_dist_cleanup_after_0.out         |  9 +++++++++
 .../upgrade_pg_dist_cleanup_before_0.out        | 17 +++++++++++++++++
 .../sql/upgrade_pg_dist_cleanup_after.sql       |  5 +++++
 .../sql/upgrade_pg_dist_cleanup_before.sql      | 10 ++++++++++
 4 files changed, 41 insertions(+)

diff --git a/src/test/regress/expected/upgrade_pg_dist_cleanup_after_0.out b/src/test/regress/expected/upgrade_pg_dist_cleanup_after_0.out
index d71fad887..168c64cca 100644
--- a/src/test/regress/expected/upgrade_pg_dist_cleanup_after_0.out
+++ b/src/test/regress/expected/upgrade_pg_dist_cleanup_after_0.out
@@ -28,3 +28,12 @@ SELECT * FROM pg_dist_cleanup;
 CALL citus_cleanup_orphaned_resources();
 NOTICE:  cleaned up 1 orphaned resources
 DROP TABLE table_with_orphaned_shards;
+-- Re-enable automatic shard cleanup by maintenance daemon as
+-- we have disabled it in upgrade_pg_dist_cleanup_before.sql
+ALTER SYSTEM RESET citus.defer_shard_delete_interval;
+SELECT pg_reload_conf();
+ pg_reload_conf
+---------------------------------------------------------------------
+ t
+(1 row)
+
diff --git a/src/test/regress/expected/upgrade_pg_dist_cleanup_before_0.out b/src/test/regress/expected/upgrade_pg_dist_cleanup_before_0.out
index a0cf9ceb1..dd6c8868e 100644
--- a/src/test/regress/expected/upgrade_pg_dist_cleanup_before_0.out
+++ b/src/test/regress/expected/upgrade_pg_dist_cleanup_before_0.out
@@ -30,6 +30,23 @@ SELECT COUNT(*) FROM pg_dist_placement WHERE shardstate = 1 AND shardid IN (SELE
 (1 row)
 
 -- create an orphaned placement based on an existing one
+--
+-- But before doing that, first disable automatic shard cleanup
+-- by maintenance daemon so that we can reliably test the cleanup
+-- in upgrade_pg_dist_cleanup_after.sql.
+ALTER SYSTEM SET citus.defer_shard_delete_interval TO -1;
+SELECT pg_reload_conf();
+ pg_reload_conf
+---------------------------------------------------------------------
+ t
+(1 row)
+
+SELECT pg_sleep(0.1);
+ pg_sleep
+---------------------------------------------------------------------
+
+(1 row)
+
 INSERT INTO pg_dist_placement(placementid, shardid, shardstate, shardlength, groupid)
     SELECT nextval('pg_dist_placement_placementid_seq'::regclass), shardid, 4, shardlength, 3-groupid
     FROM pg_dist_placement
diff --git a/src/test/regress/sql/upgrade_pg_dist_cleanup_after.sql b/src/test/regress/sql/upgrade_pg_dist_cleanup_after.sql
index e84c35b60..333ac60ca 100644
--- a/src/test/regress/sql/upgrade_pg_dist_cleanup_after.sql
+++ b/src/test/regress/sql/upgrade_pg_dist_cleanup_after.sql
@@ -13,3 +13,8 @@ SELECT COUNT(*) FROM pg_dist_placement WHERE shardid IN (SELECT shardid FROM pg_
 SELECT * FROM pg_dist_cleanup;
 CALL citus_cleanup_orphaned_resources();
 DROP TABLE table_with_orphaned_shards;
+
+-- Re-enable automatic shard cleanup by maintenance daemon as
+-- we have disabled it in upgrade_pg_dist_cleanup_before.sql
+ALTER SYSTEM RESET citus.defer_shard_delete_interval;
+SELECT pg_reload_conf();
diff --git a/src/test/regress/sql/upgrade_pg_dist_cleanup_before.sql b/src/test/regress/sql/upgrade_pg_dist_cleanup_before.sql
index 62ec8a1fb..ec0eef353 100644
--- a/src/test/regress/sql/upgrade_pg_dist_cleanup_before.sql
+++ b/src/test/regress/sql/upgrade_pg_dist_cleanup_before.sql
@@ -16,6 +16,16 @@ SELECT create_distributed_table('table_with_orphaned_shards', 'a');
 -- show all 32 placements are active
 SELECT COUNT(*) FROM pg_dist_placement WHERE shardstate = 1 AND shardid IN (SELECT shardid FROM pg_dist_shard WHERE logicalrelid='table_with_orphaned_shards'::regclass);
 -- create an orphaned placement based on an existing one
+--
+-- But before doing that, first disable automatic shard cleanup
+-- by maintenance daemon so that we can reliably test the cleanup
+-- in upgrade_pg_dist_cleanup_after.sql.
+
+ALTER SYSTEM SET citus.defer_shard_delete_interval TO -1;
+SELECT pg_reload_conf();
+
+SELECT pg_sleep(0.1);
+
 INSERT INTO pg_dist_placement(placementid, shardid, shardstate, shardlength, groupid)
     SELECT nextval('pg_dist_placement_placementid_seq'::regclass), shardid, 4, shardlength, 3-groupid
     FROM pg_dist_placement