From 4bc531d8fae93164431ad88fed0866243c94d4af Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Thu, 21 Jan 2021 08:26:37 +0530
Subject: [PATCH v15] postgres_fdw add keep_connections GUC to not cache
 connections

This patch adds a new GUC postgres_fdw.keep_connections, default
being on, when set to off no remote connections are cached by the
local session.
---
 contrib/postgres_fdw/connection.c             | 12 +++--
 .../postgres_fdw/expected/postgres_fdw.out    | 43 ++++++++++++++++++
 contrib/postgres_fdw/postgres_fdw.c           | 16 +++++++
 contrib/postgres_fdw/postgres_fdw.h           |  3 ++
 contrib/postgres_fdw/sql/postgres_fdw.sql     | 20 +++++++++
 doc/src/sgml/postgres-fdw.sgml                | 44 ++++++++++++++++++-
 6 files changed, 134 insertions(+), 4 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index be0ff43b4d..4e65a4e47a 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -809,6 +809,7 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 	while ((entry = (ConnCacheEntry *) hash_seq_search(&scan)))
 	{
 		PGresult   *res;
+		bool		used_in_current_xact = false;
 
 		/* Ignore cache entry if no open connection right now */
 		if (entry->conn == NULL)
@@ -819,6 +820,8 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 		{
 			bool		abort_cleanup_failure = false;
 
+			used_in_current_xact = true;
+
 			elog(DEBUG3, "closing remote transaction on connection %p",
 				 entry->conn);
 
@@ -953,16 +956,19 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 
 		/*
 		 * If the connection isn't in a good idle state or it is marked as
-		 * invalid, then discard it to recover. Next GetConnection will open a
-		 * new connection.
+		 * invalid or keep_connections GUC is false and this connection is used
+		 * in current xact, then discard it to recover. Next GetConnection will
+		 * open a new connection.
 		 */
 		if (PQstatus(entry->conn) != CONNECTION_OK ||
 			PQtransactionStatus(entry->conn) != PQTRANS_IDLE ||
 			entry->changing_xact_state ||
-			entry->invalidated)
+			entry->invalidated ||
+			(used_in_current_xact && !keep_connections))
 		{
 			elog(DEBUG3, "discarding connection %p", entry->conn);
 			disconnect_pg_server(entry);
+			used_in_current_xact = false;
 		}
 	}
 
diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index 27e6a8f141..ffd0af4931 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -9411,3 +9411,46 @@ SELECT * FROM postgres_fdw_get_connections() ORDER BY 1;
 -------------+-------
 (0 rows)
 
+-- ===================================================================
+-- Test postgres_fdw.keep_connections GUC
+-- ===================================================================
+-- By default, keep_connections GUC is on i.e. local session caches all the
+-- foreign server connections.
+SHOW postgres_fdw.keep_connections;
+ postgres_fdw.keep_connections 
+-------------------------------
+ on
+(1 row)
+
+-- Set it off i.e. the cached connections which are used after this setting are
+-- disconnected at the end of respective xacts.
+SET postgres_fdw.keep_connections TO off;
+-- Make connection using loopback server, connection should not be cached as
+-- the GUC is off.
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Should output nothing.
+SELECT * FROM postgres_fdw_get_connections() ORDER BY 1;
+ server_name | valid 
+-------------+-------
+(0 rows)
+
+-- Set it on i.e. connections are cached.
+SET postgres_fdw.keep_connections TO on;
+SELECT 1 FROM ft1 LIMIT 1;
+ ?column? 
+----------
+        1
+(1 row)
+
+-- Should output loopback.
+SELECT * FROM postgres_fdw_get_connections() ORDER BY 1;
+ server_name | valid 
+-------------+-------
+ loopback    | t
+(1 row)
+
diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c
index 9a31bbb86b..1698563a0e 100644
--- a/contrib/postgres_fdw/postgres_fdw.c
+++ b/contrib/postgres_fdw/postgres_fdw.c
@@ -312,6 +312,8 @@ typedef struct
 	List	   *already_used;	/* expressions already dealt with */
 } ec_member_foreign_arg;
 
+bool keep_connections = true;
+
 /*
  * SQL functions
  */
@@ -527,6 +529,20 @@ static void merge_fdw_options(PgFdwRelationInfo *fpinfo,
 							  const PgFdwRelationInfo *fpinfo_i);
 static int get_batch_size_option(Relation rel);
 
+void
+_PG_init(void)
+{
+	DefineCustomBoolVariable("postgres_fdw.keep_connections",
+							 "Enables postgres_fdw connection caching.",
+							 "When off postgres_fdw will close connections at the end of transaction.",
+							 &keep_connections,
+							 true,
+							 PGC_USERSET,
+							 0,
+							 NULL,
+							 NULL,
+							 NULL);
+}
 
 /*
  * Foreign-data wrapper handler function: return a struct with pointers
diff --git a/contrib/postgres_fdw/postgres_fdw.h b/contrib/postgres_fdw/postgres_fdw.h
index 1f67b4d9fd..ceea46b304 100644
--- a/contrib/postgres_fdw/postgres_fdw.h
+++ b/contrib/postgres_fdw/postgres_fdw.h
@@ -124,9 +124,12 @@ typedef struct PgFdwRelationInfo
 	int			relation_index;
 } PgFdwRelationInfo;
 
+extern bool keep_connections;
+
 /* in postgres_fdw.c */
 extern int	set_transmission_modes(void);
 extern void reset_transmission_modes(int nestlevel);
+extern void _PG_init(void);
 
 /* in connection.c */
 extern PGconn *GetConnection(UserMapping *user, bool will_prep_stmt);
diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql
index 4924cab74f..fc243ce972 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2891,3 +2891,23 @@ SELECT postgres_fdw_disconnect_all();
 -- List all the existing cached connections. No connection exists, so NULL
 -- should be output.
 SELECT * FROM postgres_fdw_get_connections() ORDER BY 1;
+
+-- ===================================================================
+-- Test postgres_fdw.keep_connections GUC
+-- ===================================================================
+-- By default, keep_connections GUC is on i.e. local session caches all the
+-- foreign server connections.
+SHOW postgres_fdw.keep_connections;
+-- Set it off i.e. the cached connections which are used after this setting are
+-- disconnected at the end of respective xacts.
+SET postgres_fdw.keep_connections TO off;
+-- Make connection using loopback server, connection should not be cached as
+-- the GUC is off.
+SELECT 1 FROM ft1 LIMIT 1;
+-- Should output nothing.
+SELECT * FROM postgres_fdw_get_connections() ORDER BY 1;
+-- Set it on i.e. connections are cached.
+SET postgres_fdw.keep_connections TO on;
+SELECT 1 FROM ft1 LIMIT 1;
+-- Should output loopback.
+SELECT * FROM postgres_fdw_get_connections() ORDER BY 1;
diff --git a/doc/src/sgml/postgres-fdw.sgml b/doc/src/sgml/postgres-fdw.sgml
index 738f0cbc85..cb5513344e 100644
--- a/doc/src/sgml/postgres-fdw.sgml
+++ b/doc/src/sgml/postgres-fdw.sgml
@@ -569,6 +569,42 @@ postgres=# SELECT * FROM postgres_fdw_disconnect_all();
 
 </sect2>
 
+<sect2>
+  <title>Configuration Parameters</title>
+
+  <variablelist>
+   <varlistentry>
+
+    <term>
+     <varname>postgres_fdw.keep_connections</varname> (<type>boolean</type>)
+     <indexterm>
+      <primary><varname>postgres_fdw.keep_connections</varname> configuration parameter</primary>
+     </indexterm>
+    </term>
+
+    <listitem>
+
+     <para>
+      Allows <filename>postgres_fdw</filename> to keep or discard the
+      connection made to the foreign server by the local session. Default is
+      <literal>on</literal>. When set to <literal>off</literal> the local
+      session doesn't keep the connections made to the foreign servers. Each
+      connection is discarded at the end of transaction in which it is used.
+     </para>
+
+     <para>
+      Note that setting <varname>postgres_fdw.keep_connections</varname> to
+      <literal>off</literal> does not discard any previously made and still open
+      connections immediately. They will be closed only at the end of future
+      transactions, which operated on them. To close all connections
+      immediately use <function>postgres_fdw_disconnect</function> function.
+     </para>
+
+    </listitem>
+   </varlistentry>
+  </variablelist>
+  </sect2>
+
  <sect2>
   <title>Connection Management</title>
 
@@ -587,12 +623,18 @@ postgres=# SELECT * FROM postgres_fdw_disconnect_all();
    the remote servers are kept idle until they are re-used by the local session.
    This may waste resources if those connections are not frequently used by the
    local session. To address this, the <filename>postgres_fdw</filename>
-   provides following way to remove the connections to the remote servers and
+   provides following ways to remove the connections to the remote servers and
    so the remote sessions:
     
    <function>postgres_fdw_disconnect_all()</function> to discard all the
    connections or <function>postgres_fdw_disconnect(server_name text)</function>
    to discard the connection associated with the given foreign server.
+
+   A configuration parameter, <varname>postgres_fdw.keep_connections</varname>,
+   default being <literal>on</literal>, when set to <literal>off</literal>, the
+   local session doesn't keep remote connections that are made to the foreign
+   servers. Each connection is discarded at the end of transaction in which it
+   is used.
   </para>
  </sect2>
 
-- 
2.25.1

