[PATCH] pgbench: add multiconnect option

-hackers,

This patch adds the concept of "multiconnect" to pgbench (better
terminology welcome).  The basic idea here is to allow connections made
with pgbench to use different auth values or connect to multiple
databases. We implement this using a user-provided PGSERVICEFILE and
choosing a PGSERVICE from this based on a number of strategies.
(Currently the only supported strategies are round robin or random.)

There is definite room for improvement here; at the very least, teaching
`pgbench -i` about all of the distinct DBs referenced in this service
file would ensure that initialization works as expected in all places.
For now, we are punting initialization to the user in this version of
the patch if using more that one database in the given service file.

Best,

David
diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 0c60077e1f..94616c13c2 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -161,6 +161,11 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
         <envar>PGDATABASE</envar> is used. If that is not set, the
         user name specified for the connection is used.
        </para>
+       <para>
+        If <literal>multiconnect</literal> mode is enabled, a defined
+        <literal>dbname</literal> in the chosen service will override this
+        value.
+       </para>
       </listitem>
      </varlistentry>
 
@@ -840,6 +845,39 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
 
     <variablelist>
 
+     <varlistentry>
+      <term><option>-m</option> <replaceable>servicefile</replaceable></term>
+      <term><option>--multiconnect-file=</option><replaceable>servicefile</replaceable></term>
+      <listitem>
+       <para>
+        Turns on <literal>multiconnect</literal> mode and uses the given
+        <literal>pg_service</literal>-style file to derive connection
+        information from.  Any/all connection parameters in this file will
+        overwrite any that were provided in the command-line.
+       </para>
+       <para>
+        Since this behavior will make a connection using
+        the <envar>PGSERVICEFILE</envar> mechanism, it is possible to
+        connect to other databases than the one provided in the original
+        command invocation.  This option assumes that the user has previously
+        run the necessarily initialization steps against all databases that
+        would be accessed via this service file.
+       </para>
+      </listitem>
+     </varlistentry>
+
+     <varlistentry>
+      <term><option>-g roundrobin|random</option></term>
+      <term><option>--multiconnect-strategy=roundrobin|random</option></term>
+      <listitem>
+       <para>
+        Selects the strategy by which <literal>multiconnect</literal> mode
+        uses the connections defined in the indicated service file.  The
+        default value is <literal>roundrobin</literal>.
+       </para>
+      </listitem>
+     </varlistentry>
+
      <varlistentry>
       <term><option>-h</option> <replaceable>hostname</replaceable></term>
       <term><option>--host=</option><replaceable>hostname</replaceable></term>
@@ -847,6 +885,11 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
        <para>
         The database server's host name
        </para>
+       <para>
+        If <literal>multiconnect</literal> mode is enabled, a defined
+        <literal>host</literal> in the chosen service will override this
+        value.
+       </para>
       </listitem>
      </varlistentry>
 
@@ -857,6 +900,11 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
        <para>
         The database server's port number
        </para>
+       <para>
+        If <literal>multiconnect</literal> mode is enabled, a defined
+        <literal>port</literal> in the chosen service will override this
+        value.
+       </para>
       </listitem>
      </varlistentry>
 
@@ -867,6 +915,11 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
        <para>
         The user name to connect as
        </para>
+       <para>
+        If <literal>multiconnect</literal> mode is enabled, a defined
+        <literal>user</literal> in the chosen service will override this
+        value.
+       </para>
       </listitem>
      </varlistentry>
 
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 4aeccd93af..2834c9ef3c 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -69,6 +69,7 @@
 #include "pgbench.h"
 #include "port/pg_bitutils.h"
 #include "portability/instr_time.h"
+#include "lib/stringinfo.h"
 
 #ifndef M_PI
 #define M_PI 3.14159265358979323846
@@ -275,6 +276,8 @@ int            nthreads = 1;        /* number of threads */
 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 */
+int         num_service_names = 0; /* how many service file names are in the indicated service file */
+int         cur_service_index = 0; /* the index of the next service file; used for round-robin */
 
 const char *pghost = NULL;
 const char *pgport = NULL;
@@ -282,6 +285,7 @@ const char *username = NULL;
 const char *dbName = NULL;
 char       *logfile_prefix = NULL;
 const char *progname;
+const char **service_names = NULL;
 
 #define WSEP '@'                /* weight separator */
 
@@ -549,6 +553,14 @@ typedef enum QueryMode
 static QueryMode querymode = QUERY_SIMPLE;
 static const char *QUERYMODE[] = {"simple", "extended", "prepared"};
 
+typedef enum MultiConnectStrategy
+{
+    MC_ROUND_ROBIN,
+    MC_RANDOM
+} MultiConnectStrategy;
+
+static MultiConnectStrategy multiconnect_strategy = MC_ROUND_ROBIN;
+
 /*
  * struct Command represents one command in a script.
  *
@@ -663,7 +675,7 @@ static void clear_socket_set(socket_set *sa);
 static void add_socket_to_set(socket_set *sa, int fd, int idx);
 static int    wait_on_socket_set(socket_set *sa, int64 usecs);
 static bool socket_has_input(socket_set *sa, int fd, int idx);
-
+static const char **availableServiceEntries(const char *serviceFile);
 
 /* callback functions for our flex lexer */
 static const PsqlScanCallbacks pgbench_callbacks = {
@@ -727,6 +739,10 @@ usage(void)
            "  -j, --jobs=NUM           number of threads (default: 1)\n"
            "  -l, --log                write transaction times to log file\n"
            "  -L, --latency-limit=NUM  count transactions lasting more than NUM ms as late\n"
+           "  -m, --multiconnect=FILE  use multiple auth defined in the given service file\n"
+           "  -g, --multiconnect-strategy=roundrobin|random\n"
+           "                           use the given strategy for choosing the service to connect as\n"
+           "                           (default: roundrobin)\n"
            "  -M, --protocol=simple|extended|prepared\n"
            "                           protocol for submitting queries (default: simple)\n"
            "  -n, --no-vacuum          do not run VACUUM before tests\n"
@@ -1351,6 +1367,33 @@ doConnect(void)
     bool        new_pass;
     static char *password = NULL;
 
+    /*
+     * If we are doing a round-robin of service files names, then use/choose the next name
+     */
+    if (num_service_names) {
+        const char *service;
+
+        if (multiconnect_strategy == MC_ROUND_ROBIN)
+        {
+            service = service_names[cur_service_index++];
+
+            if (cur_service_index >= num_service_names)
+                cur_service_index = 0;
+        }
+        else if (multiconnect_strategy == MC_RANDOM)
+        {
+            /*
+             * We need to get random int <= num_service_names; since this is
+             * infrequently-called and just need uniform integer distribution,
+             * we are using system random() instead of one of the more complex
+             * functions available in this file.
+             */
+            service = service_names[ ((unsigned long)random()) % num_service_names ];
+        }
+        pg_log_info("using service: %s", service);
+        setenv("PGSERVICE", service, true);
+    }
+
     /*
      * Start the connection.  Loop until we have a password if requested by
      * backend.
@@ -5724,6 +5767,74 @@ set_random_seed(const char *seed)
     return true;
 }
 
+static const char**
+availableServiceEntries(const char *serviceFile)
+{
+    int            linenr = 0,
+                num_svc = 0,
+                max_svc = 10;
+    FILE       *f;
+    char       *line, **services;
+    StringInfoData linebuf;
+
+    f = fopen(serviceFile, "r");
+    if (f == NULL)
+    {
+        return NULL;
+    }
+
+    initStringInfo(&linebuf);
+
+    services = (char **)pg_malloc0(max_svc * sizeof(char *));
+    
+    while (pg_get_line_buf(f, &linebuf))
+    {
+        linenr++;
+
+        /* ignore whitespace at end of line, especially the newline */
+        while (linebuf.len > 0 &&
+               isspace((unsigned char) linebuf.data[linebuf.len - 1]))
+            linebuf.data[--linebuf.len] = '\0';
+
+        line = linebuf.data;
+
+        /* ignore leading whitespace too */
+        while (*line && isspace((unsigned char) line[0]))
+            line++;
+
+        /* ignore comments and empty lines */
+        if (line[0] == '\0' || line[0] == '#')
+            continue;
+
+        /* Check for groupname section */
+        if (line[0] == '[')
+        {
+            char *endp;
+
+            line++;
+
+            endp = strchr(line, ']');
+            if (endp && (endp - line) > 0) {
+                /* add the literal block to the chunk */
+                services[num_svc] = pnstrdup(line, (endp - line));
+
+                /* possibly expand memory */
+                if (++num_svc >= max_svc) {
+                    max_svc += 10;
+                    services = pg_realloc(services, max_svc * sizeof(char *));
+                }
+
+                /* null out the next possible entry */
+                services[num_svc] = NULL;
+            }
+        }
+    }
+
+    fclose(f);
+    pfree(linebuf.data);
+    return (const char**)services;
+}
+
 int
 main(int argc, char **argv)
 {
@@ -5742,6 +5853,8 @@ main(int argc, char **argv)
         {"jobs", required_argument, NULL, 'j'},
         {"log", no_argument, NULL, 'l'},
         {"latency-limit", required_argument, NULL, 'L'},
+        {"multiconnect", required_argument, NULL, 'm'},
+        {"multiconnect-strategy", required_argument, NULL, 'g'},
         {"no-vacuum", no_argument, NULL, 'n'},
         {"port", required_argument, NULL, 'p'},
         {"progress", required_argument, NULL, 'P'},
@@ -5835,7 +5948,7 @@ main(int argc, char **argv)
         exit(1);
     }
 
-    while ((c = getopt_long(argc, argv, "iI:h:nvp:dqb:SNc:j:Crs:t:T:U:lf:D:F:M:P:R:L:", long_options, &optindex)) !=
-1)
+    while ((c = getopt_long(argc, argv, "iI:h:nvp:dqb:SNc:j:Crs:t:T:U:lf:D:F:M:m:g:P:R:L:", long_options, &optindex))
!=-1)
 
     {
         char       *script;
 
@@ -6008,6 +6121,55 @@ main(int argc, char **argv)
                     exit(1);
                 }
                 break;
+            case 'm':
+                {
+                    char **p;
+
+                    service_names = availableServiceEntries(optarg);
+                    p = (char**)service_names;
+
+                    if (!service_names)
+                    {
+                        pg_log_fatal("Couldn't find any services in file '%s'", optarg);
+                        exit(1);
+                    }
+
+                    while (*(p++))
+                        num_service_names++;
+
+                    /*
+                     * If we found non-zero services in our file then we can set
+                     * the PGSERVICEFILE variable to point to the file we parsed,
+                     * otherwise there is no point.
+                     */
+
+                    if (num_service_names) {
+                        setenv("PGSERVICEFILE", optarg, true);
+
+                        /*
+                         * Warn if number of services exceeds the number of
+                         * clients expected.
+                         */
+
+                        if (num_service_names > nclients)
+                            pg_log_warning("Found %d services defined, but -c is set to %d; did you mean to increase
-c?",
+                                           num_service_names,
+                                           nclients
+                                );
+                    }
+                }
+                break;
+            case 'g':
+                if (strcmp(optarg, "roundrobin") == 0)
+                    multiconnect_strategy = MC_ROUND_ROBIN;
+                else if (strcmp(optarg, "random") == 0)
+                    multiconnect_strategy = MC_RANDOM;
+                else
+                {
+                    pg_log_fatal("Unrecognized multiconnect strategy: %s", optarg);
+                    exit(1);
+                }
+                break;
             case 'M':
                 benchmarking_option_set = true;
                 for (querymode = 0; querymode < NUM_QUERYMODE; querymode++)