From d8cd283a08fe87579e3603eaa8fba4d1eca6a8ca Mon Sep 17 00:00:00 2001
From: Robert Haas <rhaas@postgresql.org>
Date: Wed, 23 Jun 2010 21:54:13 +0000
Subject: [PATCH] Add TCP keepalive support to libpq.

This adds four additional connection parameters to libpq: keepalives,
keepalives_idle, keepalives_count, and keepalives_interval.
keepalives default to on, per discussion, but can be turned off by
specifying keepalives=0.  The remaining parameters, where supported,
can be used to adjust how often keepalives are sent and how many
can be lost before the connection is broken.

The immediate motivation for this patch is to make sure that
walreceiver will eventually notice if the master reboots without
closing the connection cleanly, but it should be helpful in other
cases as well.

Tollef Fog Heen, Fujii Masao, and me.
---
 doc/src/sgml/libpq.sgml           |  53 ++++++++-
 src/interfaces/libpq/fe-connect.c | 182 +++++++++++++++++++++++++++++-
 src/interfaces/libpq/libpq-int.h  |   6 +-
 3 files changed, 238 insertions(+), 3 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index d3a22181825..11161d993f5 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -1,4 +1,4 @@
-<!-- $PostgreSQL: pgsql/doc/src/sgml/libpq.sgml,v 1.308 2010/06/17 16:03:30 tgl Exp $ -->
+<!-- $PostgreSQL: pgsql/doc/src/sgml/libpq.sgml,v 1.309 2010/06/23 21:54:13 rhaas Exp $ -->
 
 <chapter id="libpq">
  <title><application>libpq</application> - C Library</title>
@@ -280,6 +280,57 @@
          </listitem>
         </varlistentry>
 
+        <varlistentry id="libpq-keepalives" xreflabel="keepalives">
+         <term><literal>keepalives</literal></term>
+         <listitem>
+          <para>
+           Controls whether TCP keepalives are used. The default value is 1,
+           meaning on, but you can change this to 0, meaning off, if keepalives
+           are not wanted.  This parameter is ignored for connections made via
+           a Unix-domain socket.
+          </para>
+         </listitem>
+        </varlistentry>
+
+        <varlistentry id="libpq-keepalives-idle" xreflabel="keepalives_idle">
+         <term><literal>keepalives_idle</literal></term>
+         <listitem>
+          <para>
+           On systems that support the <symbol>TCP_KEEPIDLE</symbol> socket
+           option, specifies the number of seconds between sending keepalives
+           on an otherwise idle connection. A value of zero uses the system
+           default.  This parameter is ignored for connections made via a
+           Unix-domain socket, or if keepalives are disabled.
+          </para>
+         </listitem>
+        </varlistentry>
+
+        <varlistentry id="libpq-keepalives-interval" xreflabel="keepalives_interval">
+         <term><literal>keepalives_interval</literal></term>
+         <listitem>
+          <para>
+           On systems that support the <symbol>TCP_KEEPINTVL</symbol> socket
+           option, specifies how long, in seconds, to wait for a response to a
+           keepalive before retransmitting. A value of zero uses the system
+           default.  This parameter is ignored for connections made via a
+           Unix-domain socket, or if keepalives are disabled.
+          </para>
+         </listitem>
+        </varlistentry>
+
+        <varlistentry id="libpq-keepalives-count" xreflabel="keepalives_count">
+         <term><literal>keepalives_count</literal></term>
+         <listitem>
+          <para>
+           On systems that support the <symbol>TCP_KEEPCNT</symbol> socket
+           option, specifies how many keepalives can be lost before the
+           connection is considered dead. A value of zero uses the system
+           default.  This parameter is ignored for connections made via a
+           Unix-domain socket, or if keepalives are disabled.
+          </para>
+         </listitem>
+        </varlistentry>
+
         <varlistentry id="libpq-connect-tty" xreflabel="tty">
          <term><literal>tty</literal></term>
          <listitem>
diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index 3a1465457e0..795b64e55bf 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/interfaces/libpq/fe-connect.c,v 1.393 2010/05/26 21:39:27 tgl Exp $
+ *	  $PostgreSQL: pgsql/src/interfaces/libpq/fe-connect.c,v 1.394 2010/06/23 21:54:13 rhaas Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -183,6 +183,18 @@ static const PQconninfoOption PQconninfoOptions[] = {
 	{"fallback_application_name", NULL, NULL, NULL,
 	"Fallback-Application-Name", "", 64},
 
+	{"keepalives", NULL, NULL, NULL,
+	"TCP-Keepalives", "", 1}, /* should be just '0' or '1' */
+
+	{"keepalives_idle", NULL, NULL, NULL,
+	"TCP-Keepalives-Idle", "", 10}, /* strlen(INT32_MAX) == 10 */
+
+	{"keepalives_interval", NULL, NULL, NULL,
+	"TCP-Keepalives-Interval", "", 10}, /* strlen(INT32_MAX) == 10 */
+
+	{"keepalives_count", NULL, NULL, NULL,
+	"TCP-Keepalives-Count", "", 10}, /* strlen(INT32_MAX) == 10 */
+
 #ifdef USE_SSL
 
 	/*
@@ -552,6 +564,14 @@ fillPGconn(PGconn *conn, PQconninfoOption *connOptions)
 	conn->pgpass = tmp ? strdup(tmp) : NULL;
 	tmp = conninfo_getval(connOptions, "connect_timeout");
 	conn->connect_timeout = tmp ? strdup(tmp) : NULL;
+	tmp = conninfo_getval(connOptions, "keepalives");
+	conn->keepalives = tmp ? strdup(tmp) : NULL;
+	tmp = conninfo_getval(connOptions, "keepalives_idle");
+	conn->keepalives_idle = tmp ? strdup(tmp) : NULL;
+	tmp = conninfo_getval(connOptions, "keepalives_interval");
+	conn->keepalives_interval = tmp ? strdup(tmp) : NULL;
+	tmp = conninfo_getval(connOptions, "keepalives_count");
+	conn->keepalives_count = tmp ? strdup(tmp) : NULL;
 	tmp = conninfo_getval(connOptions, "sslmode");
 	conn->sslmode = tmp ? strdup(tmp) : NULL;
 	tmp = conninfo_getval(connOptions, "sslkey");
@@ -943,6 +963,119 @@ connectFailureMessage(PGconn *conn, int errorno)
 	}
 }
 
+/*
+ * Should we use keepalives?  Returns 1 if yes, 0 if no, and -1 if
+ * conn->keepalives is set to a value which is not parseable as an
+ * integer.
+ */
+static int
+useKeepalives(PGconn *conn)
+{
+	char	   *ep;
+	int			val;
+
+	if (conn->keepalives == NULL)
+		return 1;
+	val = strtol(conn->keepalives, &ep, 10);
+	if (*ep)
+		return -1;
+	return val != 0 ? 1 : 0;
+}
+
+/*
+ * Set the keepalive idle timer.
+ */
+static int
+setKeepalivesIdle(PGconn *conn)
+{
+	int	idle;
+
+	if (conn->keepalives_idle == NULL)
+		return 1;
+
+	idle = atoi(conn->keepalives_idle);
+	if (idle < 0)
+		idle = 0;
+
+#ifdef TCP_KEEPIDLE
+	if (setsockopt(conn->sock, IPPROTO_TCP, TCP_KEEPIDLE,
+				   (char *) &idle, sizeof(idle)) < 0)
+	{
+		char	sebuf[256];
+
+		appendPQExpBuffer(&conn->errorMessage,
+						  libpq_gettext("setsockopt(TCP_KEEPIDLE) failed: %s\n"),
+						  SOCK_STRERROR(SOCK_ERRNO, sebuf, sizeof(sebuf)));
+		return 0;
+	}
+#endif
+
+	return 1;
+}
+
+/*
+ * Set the keepalive interval.
+ */
+static int
+setKeepalivesInterval(PGconn *conn)
+{
+	int	interval;
+
+	if (conn->keepalives_interval == NULL)
+		return 1;
+
+	interval = atoi(conn->keepalives_interval);
+	if (interval < 0)
+		interval = 0;
+
+#ifdef TCP_KEEPINTVL
+	if (setsockopt(conn->sock, IPPROTO_TCP, TCP_KEEPINTVL,
+				   (char *) &interval, sizeof(interval)) < 0)
+	{
+		char	sebuf[256];
+
+		appendPQExpBuffer(&conn->errorMessage,
+						  libpq_gettext("setsockopt(TCP_KEEPINTVL) failed: %s\n"),
+						  SOCK_STRERROR(SOCK_ERRNO, sebuf, sizeof(sebuf)));
+		return 0;
+	}
+#endif
+
+	return 1;
+}
+
+/*
+ * Set the count of lost keepalive packets that will trigger a connection
+ * break.
+ */
+static int
+setKeepalivesCount(PGconn *conn)
+{
+	int	count;
+
+	if (conn->keepalives_count == NULL)
+		return 1;
+
+	count = atoi(conn->keepalives_count);
+	if (count < 0)
+		count = 0;
+
+#ifdef TCP_KEEPCNT
+	if (setsockopt(conn->sock, IPPROTO_TCP, TCP_KEEPCNT,
+				   (char *) &count, sizeof(count)) < 0)
+	{
+		char	sebuf[256];
+
+		appendPQExpBuffer(&conn->errorMessage,
+						  libpq_gettext("setsockopt(TCP_KEEPCNT) failed: %s\n"),
+						  SOCK_STRERROR(SOCK_ERRNO, sebuf, sizeof(sebuf)));
+		return 0;
+	}
+#endif
+
+	return 1;
+}
+
 
 /* ----------
  * connectDBStart -
@@ -1329,6 +1462,45 @@ keep_going:						/* We will come back to here until there is
 					}
 #endif   /* F_SETFD */
 
+					if (!IS_AF_UNIX(addr_cur->ai_family))
+					{
+						int		on = 1;
+						int		usekeepalives = useKeepalives(conn);
+						int		err = 0;
+
+						if (usekeepalives < 0)
+						{
+							appendPQExpBuffer(&conn->errorMessage,
+											  libpq_gettext("keepalives parameter must be an integer\n"));
+							err = 1;
+						}
+						else if (usekeepalives == 0)
+						{
+							/* Do nothing */
+						}
+						else if (setsockopt(conn->sock,
+											SOL_SOCKET, SO_KEEPALIVE,
+											(char *) &on, sizeof(on)) < 0)
+						{
+							appendPQExpBuffer(&conn->errorMessage,
+											  libpq_gettext("setsockopt(SO_KEEPALIVE) failed: %s\n"),
+							  SOCK_STRERROR(SOCK_ERRNO, sebuf, sizeof(sebuf)));
+							err = 1;
+						}
+						else if (!setKeepalivesIdle(conn)
+								 || !setKeepalivesInterval(conn)
+								 || !setKeepalivesCount(conn))
+							err = 1;
+
+						if (err)
+						{
+							closesocket(conn->sock);
+							conn->sock = -1;
+							conn->addr_cur = addr_cur->ai_next;
+							continue;
+						}
+					}
+
 					/*----------
 					 * We have three methods of blocking SIGPIPE during
 					 * send() calls to this socket:
@@ -2290,6 +2462,14 @@ freePGconn(PGconn *conn)
 		free(conn->pguser);
 	if (conn->pgpass)
 		free(conn->pgpass);
+	if (conn->keepalives)
+		free(conn->keepalives);
+	if (conn->keepalives_idle)
+		free(conn->keepalives_idle);
+	if (conn->keepalives_interval)
+		free(conn->keepalives_interval);
+	if (conn->keepalives_count)
+		free(conn->keepalives_count);
 	if (conn->sslmode)
 		free(conn->sslmode);
 	if (conn->sslcert)
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 6fe96ab1684..23b914abbc5 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -12,7 +12,7 @@
  * Portions Copyright (c) 1996-2010, PostgreSQL Global Development Group
  * Portions Copyright (c) 1994, Regents of the University of California
  *
- * $PostgreSQL: pgsql/src/interfaces/libpq/libpq-int.h,v 1.150 2010/03/13 14:55:57 momjian Exp $
+ * $PostgreSQL: pgsql/src/interfaces/libpq/libpq-int.h,v 1.151 2010/06/23 21:54:13 rhaas Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -300,6 +300,10 @@ struct pg_conn
 	char	   *replication;	/* connect as the replication standby? */
 	char	   *pguser;			/* Postgres username and password, if any */
 	char	   *pgpass;
+	char	   *keepalives;		/* use TCP keepalives? */
+	char	   *keepalives_idle;	/* time between TCP keepalives */
+	char	   *keepalives_interval;	/* time between TCP keepalive retransmits */
+	char	   *keepalives_count;	/* maximum number of TCP keepalive retransmits */
 	char	   *sslmode;		/* SSL mode (require,prefer,allow,disable) */
 	char	   *sslkey;			/* client key filename */
 	char	   *sslcert;		/* client certificate filename */
-- 
GitLab