From 78bee027fee0ef1ae5e4873c293d0412197c4ac6 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Thu, 8 Apr 2021 18:23:16 +0530
Subject: [PATCH v2] Simplify backend terminate and wait logic in postgres_fdw
 test

With the recent commit aaf0432572 which introduced a waiting
capability for pg_teriminate_backend function, we can simply
backend terminate and wait logic in postgres_fdw.sql tests.
---
 contrib/postgres_fdw/expected/postgres_fdw.out | 13 ++++++++-----
 contrib/postgres_fdw/sql/postgres_fdw.sql      | 13 ++++++++-----
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 7f69fa0054..97b9431021 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -9198,13 +9198,16 @@ WARNING:  there is no transaction in progress
 -- Terminate the backend having the specified application_name and wait for
 -- the termination to complete.
 CREATE OR REPLACE PROCEDURE terminate_backend_and_wait(appname text) AS $$
+DECLARE
+	pid_v	int;
+	is_terminated bool := false;
 BEGIN
-    PERFORM pg_terminate_backend(pid) FROM pg_stat_activity
-    WHERE application_name = appname;
+	SELECT pid INTO pid_v FROM pg_stat_activity
+		WHERE application_name = appname;
+	SELECT * INTO is_terminated FROM pg_terminate_backend(pid_v);
     LOOP
-        PERFORM * FROM pg_stat_activity WHERE application_name = appname;
-        EXIT WHEN NOT FOUND;
-        PERFORM pg_sleep(1), pg_stat_clear_snapshot();
+		EXIT WHEN is_terminated;
+		SELECT * INTO is_terminated FROM pg_wait_for_backend_termination(pid_v, 1000);
     END LOOP;
 END;
 $$ LANGUAGE plpgsql;
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 7487096eac..accc904c3e 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2789,13 +2789,16 @@ ROLLBACK;
 -- Terminate the backend having the specified application_name and wait for
 -- the termination to complete.
 CREATE OR REPLACE PROCEDURE terminate_backend_and_wait(appname text) AS $$
+DECLARE
+	pid_v	int;
+	is_terminated bool := false;
 BEGIN
-    PERFORM pg_terminate_backend(pid) FROM pg_stat_activity
-    WHERE application_name = appname;
+	SELECT pid INTO pid_v FROM pg_stat_activity
+		WHERE application_name = appname;
+	SELECT * INTO is_terminated FROM pg_terminate_backend(pid_v);
     LOOP
-        PERFORM * FROM pg_stat_activity WHERE application_name = appname;
-        EXIT WHEN NOT FOUND;
-        PERFORM pg_sleep(1), pg_stat_clear_snapshot();
+		EXIT WHEN is_terminated;
+		SELECT * INTO is_terminated FROM pg_wait_for_backend_termination(pid_v, 1000);
     END LOOP;
 END;
 $$ LANGUAGE plpgsql;
-- 
2.25.1

