From 13eccc4c4c4b2ac4548608f37b3518160e8f3283 Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Fri, 1 Jan 2021 14:33:41 +0530
Subject: [PATCH v5 2/3] 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    | 50 +++++++++++++++++++
 contrib/postgres_fdw/postgres_fdw.c           | 16 ++++++
 contrib/postgres_fdw/postgres_fdw.h           |  3 ++
 contrib/postgres_fdw/sql/postgres_fdw.sql     | 32 +++++++++++-
 doc/src/sgml/postgres-fdw.sgml                | 46 ++++++++++++++++-
 6 files changed, 154 insertions(+), 5 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index b9f6050f4b..323242c1d6 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -816,6 +816,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)
@@ -826,6 +827,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);
 
@@ -960,16 +963,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 ba29bc6131..e08d9f38be 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -9448,5 +9448,55 @@ DETAIL:  Such connections are discarded at the end of remote transaction.
 (0 rows)
 
 COMMIT;
+-- ===================================================================
+-- 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 fdw_unnest(postgres_fdw_get_connections()) ORDER BY 1;
+ fdw_unnest 
+------------
+(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 fdw_unnest(postgres_fdw_get_connections()) ORDER BY 1;
+    fdw_unnest    
+------------------
+ (loopback, true)
+(1 row)
+
+-- Discard loopback connection.
+SELECT * FROM postgres_fdw_disconnect();	/* TRUE */
+ postgres_fdw_disconnect 
+-------------------------
+ t
+(1 row)
+
 -- Clean up.
 DROP FUNCTION fdw_unnest;
diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c
index b6c72e1d1e..160f17972b 100644
--- a/contrib/postgres_fdw/postgres_fdw.c
+++ b/contrib/postgres_fdw/postgres_fdw.c
@@ -302,6 +302,8 @@ typedef struct
 	List	   *already_used;	/* expressions already dealt with */
 } ec_member_foreign_arg;
 
+bool keep_connections = true;
+
 /*
  * SQL functions
  */
@@ -506,6 +508,20 @@ static void merge_fdw_options(PgFdwRelationInfo *fpinfo,
 							  const PgFdwRelationInfo *fpinfo_o,
 							  const PgFdwRelationInfo *fpinfo_i);
 
+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 eef410db39..7f1bdb96d6 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 2267aee630..b37aa569d6 100644
--- a/contrib/postgres_fdw/sql/postgres_fdw.sql
+++ b/contrib/postgres_fdw/sql/postgres_fdw.sql
@@ -2875,5 +2875,35 @@ DROP SERVER temploopback2 CASCADE;
 SELECT * FROM fdw_unnest(postgres_fdw_get_connections()) ORDER BY 1; /* WARNING */
 COMMIT;
 
+-- ===================================================================
+-- 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 fdw_unnest(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 fdw_unnest(postgres_fdw_get_connections()) ORDER BY 1;
+
+-- Discard loopback connection.
+SELECT * FROM postgres_fdw_disconnect();	/* TRUE */
+
 -- Clean up.
-DROP FUNCTION fdw_unnest;
+DROP FUNCTION fdw_unnest;
\ No newline at end of file
diff --git a/doc/src/sgml/postgres-fdw.sgml b/doc/src/sgml/postgres-fdw.sgml
index 405923f2aa..15def846f5 100644
--- a/doc/src/sgml/postgres-fdw.sgml
+++ b/doc/src/sgml/postgres-fdw.sgml
@@ -519,6 +519,44 @@ OPTIONS (ADD password_required 'false');
 
 </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>
 
@@ -537,12 +575,18 @@ OPTIONS (ADD password_required 'false');
    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()</function> to discard all the
    connections or <function>postgres_fdw_disconnect(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

