Bonjour Michaël,

Good. I was thinking of adding such capability, possibly for handling
connection errors and reconnecting…

round-robin and random make sense.  I am wondering how round-robin
would work with -C, though?  Would you just reuse the same connection
string as the one chosen at the starting point.

Well, not necessarily, but this is debatable.

I was thinking of providing a allowing a list of conninfo strings with
repeated options, eg --conninfo "foo" --conninfo "bla"…

That was my first thought when reading the subject of this thread:
create a list of connection strings and pass one of them to
doConnect() to grab the properties looked for.  That's a bit confusing
though as pgbench does not support directly connection strings,

They are supported because libpq silently assumes that "dbname" can be a full connection string.

and we should be careful to keep fallback_application_name intact.

Hmmm. See attached patch, ISTM that it does the right thing.

Your approach using PGSERVICEFILE also make sense!

I am not sure that's actually needed here, as it is possible to pass
down a service name within a connection string.  I think that you'd
better leave libpq do all the work related to a service file, if
specified.  pgbench does not need to know any of that.

Yes, this is an inconvenient with this approach, part of libpq machinery
is more or less replicated in pgbench, which is quite annoying, and less powerful.

Attached my work-in-progress version, with a few open issues (eg probably not thread safe), but comments about the provided feature are welcome.

I borrowed the "strategy" option, renamed policy, from the initial patch. Pgbench just accepts several connection strings as parameters, eg:

  pgbench ... "service=db1" "service=db2" "service=db3"

The next stage is to map scripts to connections types and connections
to connection types, so that pgbench could run W transactions against a primary and R transactions agains a hot standby, for instance. I have a some design for that, but nothing is implemented.

There is also the combination with the error handling patch to consider: if a connection fails, a connection to a replica could be issued instead.

--
Fabien.
diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 0c60077e1f..7b99344c90 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -29,7 +29,7 @@ PostgreSQL documentation
   <cmdsynopsis>
    <command>pgbench</command>
    <arg rep="repeat"><replaceable>option</replaceable></arg>
-   <arg choice="opt"><replaceable>dbname</replaceable></arg>
+   <arg rep="repeat"><replaceable>dbname or conninfo</replaceable></arg>
   </cmdsynopsis>
  </refsynopsisdiv>
 
@@ -160,6 +160,9 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
         not specified, the environment variable
         <envar>PGDATABASE</envar> is used. If that is not set, the
         user name specified for the connection is used.
+        Alternatively, the <replaceable>dbname</replaceable> can be
+        a standard connection information string.
+        Several connections can be provided.
        </para>
       </listitem>
      </varlistentry>
@@ -840,6 +843,21 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
 
     <variablelist>
 
+     <varlistentry>
+      <term><option>--connection-policy=</option><replaceable>policy</replaceable></term>
+      <listitem>
+       <para>
+        Set the connection policy when multiple connections are available.
+        Default is <literal>round-robin</literal> provided (<literal>ro</literal>).
+        Possible values are:
+        <literal>first</literal> (<literal>f</literal>), 
+        <literal>random</literal> (<literal>ra</literal>), 
+        <literal>round-robin</literal> (<literal>ro</literal>),
+        <literal>working</literal> (<literal>w</literal>).
+       </para>
+      </listitem>
+     </varlistentry>
+
      <varlistentry>
       <term><option>-h</option> <replaceable>hostname</replaceable></term>
       <term><option>--host=</option><replaceable>hostname</replaceable></term>
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index b0e20c46ae..95e58f0573 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -277,13 +277,39 @@ bool		is_connect;			/* establish connection for each transaction */
 bool		report_per_command; /* report per-command latencies */
 int			main_pid;			/* main process id used in log filename */
 
+char	   *logfile_prefix = NULL;
+
+/* main connection definition */
 const char *pghost = NULL;
 const char *pgport = NULL;
 const char *username = NULL;
-const char *dbName = NULL;
-char	   *logfile_prefix = NULL;
 const char *progname;
 
+/* multi connections */
+typedef enum mc_policy_t
+{
+	MC_UNKNOWN = 0,
+	MC_FIRST,
+	MC_RANDOM,
+	MC_ROUND_ROBIN,
+	MC_WORKING
+} mc_policy_t;
+
+/* connection info list */
+typedef struct connection_t
+{
+	const char *connection;		/* conninfo or dbname */
+	int			errors;			/* number of connection errors */
+} connection_t;
+
+static int				n_connections = 0;
+static connection_t	   *connections = NULL;
+static mc_policy_t	mc_policy = MC_ROUND_ROBIN;
+
+/* last used connection */
+// FIXME per thread?
+static int current_connection = 0;
+
 #define WSEP '@'				/* weight separator */
 
 volatile bool timer_exceeded = false;	/* flag from signal handler */
@@ -701,7 +727,7 @@ usage(void)
 {
 	printf("%s is a benchmarking tool for PostgreSQL.\n\n"
 		   "Usage:\n"
-		   "  %s [OPTION]... [DBNAME]\n"
+		   "  %s [OPTION]... [DBNAME or CONNINFO ...]\n"
 		   "\nInitialization options:\n"
 		   "  -i, --initialize         invokes initialization mode\n"
 		   "  -I, --init-steps=[" ALL_INIT_STEPS "]+ (default \"" DEFAULT_INIT_STEPS "\")\n"
@@ -756,6 +782,7 @@ usage(void)
 		   "  -h, --host=HOSTNAME      database server host or socket directory\n"
 		   "  -p, --port=PORT          database server port number\n"
 		   "  -U, --username=USERNAME  connect as specified database user\n"
+		   "  --connection-policy=S    set multiple connection policy (\"first\", \"rand\", \"round-robin\", \"working\")\n"
 		   "  -V, --version            output version information, then exit\n"
 		   "  -?, --help               show this help, then exit\n"
 		   "\n"
@@ -1350,13 +1377,89 @@ tryExecuteStatement(PGconn *con, const char *sql)
 	PQclear(res);
 }
 
+/* store a new connection information string */
+static void
+push_connection(const char *c)
+{
+	connections = pg_realloc(connections, sizeof(connection_t) * (n_connections+1));
+	connections[n_connections].connection = pg_strdup(c);
+	connections[n_connections].errors = 0;
+	n_connections++;
+}
+
+/* switch connection */
+static int
+next_connection(int *pci)
+{
+	int ci;
+
+	ci = ((*pci) + 1) % n_connections;
+	*pci = ci;
+
+	return ci;
+}
+
+/* return the connection index to use for next attempt */
+static int
+choose_connection(int *pci)
+{
+	int ci;
+
+	switch (mc_policy)
+	{
+		case MC_FIRST:
+			ci = 0;
+			break;
+		case MC_RANDOM:
+			// FIXME should use a prng state ; not thread safe ;
+			ci = (int) getrand(&base_random_sequence, 0, n_connections-1);
+			*pci = ci;
+			break;
+		case MC_ROUND_ROBIN:
+			ci = next_connection(pci);
+			break;
+		case MC_WORKING:
+			ci = *pci;
+			break;
+		default:
+			pg_log_fatal("unexpected multi connection policy: %d", mc_policy);
+			exit(1);
+	}
+
+	return ci;
+}
+
+/* return multi-connection policy based on its name or shortest prefix */
+static mc_policy_t
+get_connection_policy(const char *s)
+{
+	if (s == NULL || *s == '\0' || strcmp(s, "first") == 0 || strcmp(s, "f") == 0)
+		return MC_FIRST;
+	else if (strcmp(s, "random") == 0 || strcmp(s, "ra") == 0)
+		return MC_RANDOM;
+	else if (strcmp(s, "round-robin") == 0 || strcmp(s, "ro") == 0)
+		return MC_ROUND_ROBIN;
+	else if (strcmp(s, "working") == 0 || strcmp(s, "w") == 0)
+		return MC_WORKING;
+	else
+		return MC_UNKNOWN;
+}
+
+/* get backend connection info */
+static connection_t *
+getConnection(void)
+{
+	return &connections[choose_connection(&current_connection)];
+}
+
 /* set up a connection to the backend */
 static PGconn *
 doConnect(void)
 {
-	PGconn	   *conn;
-	bool		new_pass;
-	static char *password = NULL;
+	PGconn		   *conn;
+	bool			new_pass;
+	static char    *password = NULL;
+	connection_t   *ci = getConnection();
 
 	/*
 	 * Start the connection.  Loop until we have a password if requested by
@@ -1377,8 +1480,9 @@ doConnect(void)
 		values[2] = username;
 		keywords[3] = "password";
 		values[3] = password;
+		/* dbname can include a full conninfo */
 		keywords[4] = "dbname";
-		values[4] = dbName;
+		values[4] = ci->connection;
 		keywords[5] = "fallback_application_name";
 		values[5] = progname;
 		keywords[6] = NULL;
@@ -1386,11 +1490,12 @@ doConnect(void)
 
 		new_pass = false;
 
+		pg_log_debug("connecting to %s", ci->connection);
 		conn = PQconnectdbParams(keywords, values, true);
 
 		if (!conn)
 		{
-			pg_log_error("connection to database \"%s\" failed", dbName);
+			pg_log_error("connection to database \"%s\" failed", ci->connection);
 			return NULL;
 		}
 
@@ -1409,6 +1514,9 @@ doConnect(void)
 	{
 		pg_log_error("%s", PQerrorMessage(conn));
 		PQfinish(conn);
+		ci->errors += 1;
+		if (mc_policy == MC_WORKING)
+			(void) next_connection(&current_connection);
 		return NULL;
 	}
 
@@ -5787,6 +5895,7 @@ main(int argc, char **argv)
 		{"show-script", required_argument, NULL, 10},
 		{"partitions", required_argument, NULL, 11},
 		{"partition-method", required_argument, NULL, 12},
+		{"connection-policy", required_argument, NULL, 13},
 		{NULL, 0, NULL, 0}
 	};
 
@@ -6140,6 +6249,14 @@ main(int argc, char **argv)
 					exit(1);
 				}
 				break;
+			case 13:
+				mc_policy = get_connection_policy(optarg);
+				if (mc_policy == MC_UNKNOWN)
+				{
+					pg_log_fatal("unexpected connection policy: %s", optarg);
+					exit(1);
+				}
+				break;
 			default:
 				fprintf(stderr, _("Try \"%s --help\" for more information.\n"), progname);
 				exit(1);
@@ -6194,23 +6311,18 @@ main(int argc, char **argv)
 	throttle_delay *= nthreads;
 
 	if (argc > optind)
-		dbName = argv[optind++];
+	{
+		while (optind < argc)
+			push_connection(argv[optind++]);
+	}
 	else
 	{
 		if ((env = getenv("PGDATABASE")) != NULL && *env != '\0')
-			dbName = env;
+			push_connection(env);
 		else if ((env = getenv("PGUSER")) != NULL && *env != '\0')
-			dbName = env;
+			push_connection(env);
 		else
-			dbName = get_user_name_or_exit(progname);
-	}
-
-	if (optind < argc)
-	{
-		pg_log_fatal("too many command-line arguments (first is \"%s\")",
-					 argv[optind]);
-		fprintf(stderr, _("Try \"%s --help\" for more information.\n"), progname);
-		exit(1);
+			push_connection(get_user_name_or_exit(progname));
 	}
 
 	if (is_init_mode)

Reply via email to