From 144e1dedae12a8460c07cafcd5e19adf2660f70f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fabr=C3=ADzio=20de=20Royes=20Mello?=
 <fabriziomello@gmail.com>
Date: Fri, 24 Nov 2023 12:50:59 -0300
Subject: [PATCH] Fix flaky pg_dump test
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This is a very known flaky test where we need to create another database
as a template of the TEST_DBNAME but some background workers didn't had
enough time to properly shutdown and disconnect from the database.

Fixed it by forcing other processes to terminate just after waiting for
background workers to discinnect from the database.

Closes #4766

Signed-off-by: Fabrízio de Royes Mello <fabriziomello@gmail.com>
---
 test/expected/pg_dump.out | 16 ++++++++++++++--
 test/sql/pg_dump.sql      | 12 ++++++++++--
 test/src/test_utils.c     | 22 ++++++++++++++--------
 3 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/test/expected/pg_dump.out b/test/expected/pg_dump.out
index 01a20c4f30f..27b78a999fe 100644
--- a/test/expected/pg_dump.out
+++ b/test/expected/pg_dump.out
@@ -36,7 +36,7 @@ INSERT INTO "two_Partitions"("timeCustom", device_id, series_0, series_1) VALUES
 \set QUIET on
 \o
 \c :TEST_DBNAME :ROLE_SUPERUSER
-CREATE OR REPLACE FUNCTION bgw_wait(database TEXT, timeout INT)
+CREATE OR REPLACE FUNCTION bgw_wait(database TEXT, timeout INT, raise_error BOOLEAN DEFAULT TRUE)
 RETURNS VOID
 AS :MODULE_PATHNAME, 'ts_bgw_wait'
 LANGUAGE C VOLATILE;
@@ -609,12 +609,24 @@ SELECT timescaledb_pre_restore();
  t
 (1 row)
 
-SELECT bgw_wait(:'TEST_DBNAME', 60);
+SELECT bgw_wait(:'TEST_DBNAME', 60, FALSE);
  bgw_wait 
 ----------
  
 (1 row)
 
+-- Force other sessions connected to the TEST_DBNAME to be finished
+SET client_min_messages TO ERROR;
+SELECT COUNT(pg_catalog.pg_terminate_backend(pid))>=0
+FROM pg_stat_activity
+WHERE pid <> pg_backend_pid()
+AND datname = ':TEST_DBNAME';
+ ?column? 
+----------
+ t
+(1 row)
+
+RESET client_min_messages;
 CREATE DATABASE :TEST_DBNAME_EXTRA WITH TEMPLATE :TEST_DBNAME;
 -- Connect to the database and do some basic stuff to check that the
 -- extension works.
diff --git a/test/sql/pg_dump.sql b/test/sql/pg_dump.sql
index d288980215b..e958298325a 100644
--- a/test/sql/pg_dump.sql
+++ b/test/sql/pg_dump.sql
@@ -9,7 +9,7 @@
 \o
 \c :TEST_DBNAME :ROLE_SUPERUSER
 
-CREATE OR REPLACE FUNCTION bgw_wait(database TEXT, timeout INT)
+CREATE OR REPLACE FUNCTION bgw_wait(database TEXT, timeout INT, raise_error BOOLEAN DEFAULT TRUE)
 RETURNS VOID
 AS :MODULE_PATHNAME, 'ts_bgw_wait'
 LANGUAGE C VOLATILE;
@@ -197,7 +197,15 @@ drop function get_sqlstate(TEXT);
 -- cannot have any backends connected to the database when cloning it.
 \c :TEST_DBNAME :ROLE_SUPERUSER
 SELECT timescaledb_pre_restore();
-SELECT bgw_wait(:'TEST_DBNAME', 60);
+SELECT bgw_wait(:'TEST_DBNAME', 60, FALSE);
+
+-- Force other sessions connected to the TEST_DBNAME to be finished
+SET client_min_messages TO ERROR;
+SELECT COUNT(pg_catalog.pg_terminate_backend(pid))>=0
+FROM pg_stat_activity
+WHERE pid <> pg_backend_pid()
+AND datname = ':TEST_DBNAME';
+RESET client_min_messages;
 
 CREATE DATABASE :TEST_DBNAME_EXTRA WITH TEMPLATE :TEST_DBNAME;
 
diff --git a/test/src/test_utils.c b/test/src/test_utils.c
index 3447e643271..36cee439cf4 100644
--- a/test/src/test_utils.c
+++ b/test/src/test_utils.c
@@ -251,6 +251,7 @@ ts_bgw_wait(PG_FUNCTION_ARGS)
 	/* The timeout is given in seconds, so we compute the number of iterations
 	 * necessary to get a coverage of that time */
 	uint32 iterations = PG_ARGISNULL(1) ? 5 : (PG_GETARG_UINT32(1) + 4) / 5;
+	bool raise_error = PG_ARGISNULL(2) ? true : PG_GETARG_BOOL(2);
 	Oid dboid = get_database_oid(text_to_cstring(datname), false);
 
 	/* This function contains a timeout of 5 seconds, so we iterate a few
@@ -269,14 +270,19 @@ ts_bgw_wait(PG_FUNCTION_ARGS)
 						   notherbackends,
 						   npreparedxacts)));
 	}
-	ereport(ERROR,
-			(errcode(ERRCODE_OBJECT_IN_USE),
-			 errmsg("source database \"%s\" is being accessed by other users",
-					text_to_cstring(datname)),
-			 errdetail("There are %d other session(s) and %d prepared transaction(s) using the "
-					   "database.",
-					   notherbackends,
-					   npreparedxacts)));
+
+	if (raise_error)
+	{
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_IN_USE),
+				 errmsg("source database \"%s\" is being accessed by other users",
+						text_to_cstring(datname)),
+				 errdetail("There are %d other session(s) and %d prepared transaction(s) using the "
+						   "database.",
+						   notherbackends,
+						   npreparedxacts)));
+	}
+
 	pg_unreachable();
 }