From 9a00415b91bcae2a0ee1cce4b9593a374b17b83b Mon Sep 17 00:00:00 2001
From: Marina Polyakova <m.polyakova@postgrespro.ru>
Date: Tue, 7 Aug 2018 13:29:38 +0300
Subject: [PATCH v10 3/4] Pgbench errors: use the Variables structure for
 client variables

This is most important when it is used to reset client variables during the
repeating of transactions after serialization/deadlock failures.

Don't allocate Variable structs one by one. Instead, first time allocate 8
variables and then enlarge by duplicating size.
---
 src/bin/pgbench/pgbench.c | 243 ++++++++++++++++++++++++++++++++--------------
 1 file changed, 171 insertions(+), 72 deletions(-)

diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index c45cd44..4f8700b 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -202,6 +202,12 @@ const char *progname;
 volatile bool timer_exceeded = false;	/* flag from signal handler */
 
 /*
+ * The number of varaibles for which the array is allocated for the first time.
+ * If necessary, every next time the array size just doubles.
+ */
+#define DEFAULT_MAX_VARIABLES	8
+
+/*
  * Variable definitions.
  *
  * If a variable only has a string value, "svalue" is that value, and value is
@@ -218,6 +224,24 @@ typedef struct
 	PgBenchValue value;			/* actual variable's value */
 } Variable;
 
+/*
+ * Data structure for client variables.
+ */
+typedef struct
+{
+	Variable   *vars;			/* array of variable definitions */
+	int			nvars;			/* number of variables */
+
+	/*
+	 * The maximum number of variables that we can currently store in 'vars'
+	 * without having to reallocate more space. We must always have max_vars >=
+	 * nvars.
+	 */
+	int			max_vars;
+
+	bool		vars_sorted;	/* are variables sorted by name? */
+} Variables;
+
 #define MAX_SCRIPTS		128		/* max number of SQL scripts allowed */
 #define SHELL_COMMAND_SIZE	256 /* maximum size allowed for shell command */
 
@@ -349,9 +373,7 @@ typedef struct
 	int			command;		/* command number in script */
 
 	/* client variables */
-	Variable   *variables;		/* array of variable definitions */
-	int			nvariables;		/* number of variables */
-	bool		vars_sorted;	/* are variables sorted by name? */
+	Variables   variables;
 
 	/* various times about current transaction */
 	int64		txn_scheduled;	/* scheduled start time of transaction (usec) */
@@ -1248,39 +1270,39 @@ compareVariableNames(const void *v1, const void *v2)
 
 /* Locate a variable by name; returns NULL if unknown */
 static Variable *
-lookupVariable(CState *st, char *name)
+lookupVariable(Variables *variables, char *name)
 {
 	Variable	key;
 
 	/* On some versions of Solaris, bsearch of zero items dumps core */
-	if (st->nvariables <= 0)
+	if (variables->nvars <= 0)
 		return NULL;
 
 	/* Sort if we have to */
-	if (!st->vars_sorted)
+	if (!variables->vars_sorted)
 	{
-		qsort((void *) st->variables, st->nvariables, sizeof(Variable),
+		qsort((void *) variables->vars, variables->nvars, sizeof(Variable),
 			  compareVariableNames);
-		st->vars_sorted = true;
+		variables->vars_sorted = true;
 	}
 
 	/* Now we can search */
 	key.name = name;
 	return (Variable *) bsearch((void *) &key,
-								(void *) st->variables,
-								st->nvariables,
+								(void *) variables->vars,
+								variables->nvars,
 								sizeof(Variable),
 								compareVariableNames);
 }
 
 /* Get the value of a variable, in string form; returns NULL if unknown */
 static char *
-getVariable(CState *st, char *name)
+getVariable(Variables *variables, char *name)
 {
 	Variable   *var;
 	char		stringform[64];
 
-	var = lookupVariable(st, name);
+	var = lookupVariable(variables, name);
 	if (var == NULL)
 		return NULL;			/* not found */
 
@@ -1399,54 +1421,119 @@ valid_variable_name(const char *name)
 }
 
 /*
+ * Make sure there is enough space for 'needed' more variable in the variables
+ * array.
+ * On failure (too many variables are requested): if elevel < ERROR returns
+ * false; exits the program otherwise.
+ */
+static bool
+enlargeVariables(Variables *variables, const char *context, int needed,
+				 ErrorLevel elevel)
+{
+	size_t		new_max_vars;
+	Variable   *new_vars;
+
+	if ((size_t) variables->nvars + needed > INT_MAX)
+	{
+		if (elevel >= log_min_messages)
+		{
+			PQExpBufferData errmsg_buf;
+
+			initPQExpBuffer(&errmsg_buf);
+			if (context)
+				appendPQExpBuffer(&errmsg_buf, "%s: ", context);
+			appendPQExpBuffer(&errmsg_buf,
+							  "too many variables are used (limit is %d)\n",
+							  INT_MAX);
+			pgbench_error(elevel, "%s", errmsg_buf.data);
+			termPQExpBuffer(&errmsg_buf);
+		}
+
+		return false;
+	}
+
+	/* total number of variables required now */
+	needed += variables->nvars;
+
+	/* Because of the above test, we now have needed <= INT_MAX */
+
+	if (needed <= variables->max_vars)
+		return true;			/* got enough space already */
+
+	/*
+	 * We don't want to allocate just a little more space with each addition;
+	 * for efficiency, double the array size each time it overflows.
+	 * Actually, we might need to more than double it if 'needed' is big...
+	 */
+
+	if (variables->max_vars > 0)
+		new_max_vars = 2 * ((size_t) variables->max_vars);
+	else
+		new_max_vars = DEFAULT_MAX_VARIABLES;
+
+	while ((size_t) needed > new_max_vars)
+		new_max_vars = 2 * new_max_vars;
+
+	/*
+	 * Clamp to INT_MAX in case we went past it.  Note we are assuming here
+	 * that INT_MAX <= UINT_MAX/2, else the above loop could overflow.  We
+	 * will still have new_max_vars >= needed_vars.
+	 */
+	if (new_max_vars > (size_t) INT_MAX)
+		new_max_vars = (size_t) INT_MAX;
+
+	new_vars = (Variable *) pg_realloc(variables->vars,
+									   new_max_vars * sizeof(Variable));
+	variables->vars = new_vars;
+	variables->max_vars = new_max_vars;
+	return true;
+}
+
+/*
  * Lookup a variable by name, creating it if need be.
  * Caller is expected to assign a value to the variable.
  * On failure (bad name): if this is a client run returns NULL; exits the
  * program otherwise.
  */
 static Variable *
-lookupCreateVariable(CState *st, const char *context, char *name, bool client)
+lookupCreateVariable(Variables *variables, const char *context, char *name,
+					 bool client)
 {
 	Variable   *var;
 
-	var = lookupVariable(st, name);
+	/*
+	 * About the error level used: if we process client commands, it a normal
+	 * failure; otherwise it is not and we exit the program.
+	 */
+	ErrorLevel elevel = client ? LOG : ERROR;
+
+	var = lookupVariable(variables, name);
 	if (var == NULL)
 	{
-		Variable   *newvars;
-
 		/*
 		 * Check for the name only when declaring a new variable to avoid
 		 * overhead.
 		 */
 		if (!valid_variable_name(name))
 		{
-			/*
-			 * About the error level used: if we process client commands, it a
-			 * normal failure; otherwise it is not and we exit the program.
-			 */
-			pgbench_error(client ? LOG : ERROR,
-						  "%s: invalid variable name: \"%s\"\n", context, name);
+			pgbench_error(elevel, "%s: invalid variable name: \"%s\"\n",
+						  context, name);
 			return NULL;
 		}
 
 		/* Create variable at the end of the array */
-		if (st->variables)
-			newvars = (Variable *) pg_realloc(st->variables,
-											  (st->nvariables + 1) * sizeof(Variable));
-		else
-			newvars = (Variable *) pg_malloc(sizeof(Variable));
-
-		st->variables = newvars;
+		if (!enlargeVariables(variables, context, 1, elevel))
+			return NULL;
 
-		var = &newvars[st->nvariables];
+		var = &(variables->vars[variables->nvars]);
 
 		var->name = pg_strdup(name);
 		var->svalue = NULL;
 		/* caller is expected to initialize remaining fields */
 
-		st->nvariables++;
+		variables->nvars++;
 		/* we don't re-sort the array till we have to */
-		st->vars_sorted = false;
+		variables->vars_sorted = false;
 	}
 
 	return var;
@@ -1455,12 +1542,13 @@ lookupCreateVariable(CState *st, const char *context, char *name, bool client)
 /* Assign a string value to a variable, creating it if need be */
 /* Exits on failure (bad name) */
 static void
-putVariable(CState *st, const char *context, char *name, const char *value)
+putVariable(Variables *variables, const char *context, char *name,
+			const char *value)
 {
 	Variable   *var;
 	char	   *val;
 
-	var = lookupCreateVariable(st, context, name, false);
+	var = lookupCreateVariable(variables, context, name, false);
 
 	/* dup then free, in case value is pointing at this variable */
 	val = pg_strdup(value);
@@ -1477,12 +1565,12 @@ putVariable(CState *st, const char *context, char *name, const char *value)
  * program otherwise.
  */
 static bool
-putVariableValue(CState *st, const char *context, char *name,
+putVariableValue(Variables *variables, const char *context, char *name,
 				 const PgBenchValue *value, bool client)
 {
 	Variable   *var;
 
-	var = lookupCreateVariable(st, context, name, client);
+	var = lookupCreateVariable(variables, context, name, client);
 	if (!var)
 		return false;
 
@@ -1500,13 +1588,13 @@ putVariableValue(CState *st, const char *context, char *name,
  * program otherwise.
  */
 static bool
-putVariableInt(CState *st, const char *context, char *name, int64 value,
-			   bool client)
+putVariableInt(Variables *variables, const char *context, char *name,
+			   int64 value, bool client)
 {
 	PgBenchValue val;
 
 	setIntValue(&val, value);
-	return putVariableValue(st, context, name, &val, client);
+	return putVariableValue(variables, context, name, &val, client);
 }
 
 /*
@@ -1561,7 +1649,7 @@ replaceVariable(char **sql, char *param, int len, char *value)
 }
 
 static char *
-assignVariables(CState *st, char *sql)
+assignVariables(Variables *variables, char *sql)
 {
 	char	   *p,
 			   *name,
@@ -1582,7 +1670,7 @@ assignVariables(CState *st, char *sql)
 			continue;
 		}
 
-		val = getVariable(st, name);
+		val = getVariable(variables, name);
 		free(name);
 		if (val == NULL)
 		{
@@ -1597,12 +1685,13 @@ assignVariables(CState *st, char *sql)
 }
 
 static void
-getQueryParams(CState *st, const Command *command, const char **params)
+getQueryParams(Variables *variables, const Command *command,
+			   const char **params)
 {
 	int			i;
 
 	for (i = 0; i < command->argc - 1; i++)
-		params[i] = getVariable(st, command->argv[i + 1]);
+		params[i] = getVariable(variables, command->argv[i + 1]);
 }
 
 static char *
@@ -2471,7 +2560,7 @@ evaluateExpr(TState *thread, CState *st, PgBenchExpr *expr, PgBenchValue *retval
 			{
 				Variable   *var;
 
-				if ((var = lookupVariable(st, expr->u.variable.varname)) == NULL)
+				if ((var = lookupVariable(&st->variables, expr->u.variable.varname)) == NULL)
 				{
 					pgbench_error(LOG, "undefined variable \"%s\"\n",
 								  expr->u.variable.varname);
@@ -2540,7 +2629,7 @@ getMetaCommand(const char *cmd)
  * Return true if succeeded, or false on error.
  */
 static bool
-runShellCommand(CState *st, char *variable, char **argv, int argc)
+runShellCommand(Variables *variables, char *variable, char **argv, int argc)
 {
 	char		command[SHELL_COMMAND_SIZE];
 	int			i,
@@ -2571,7 +2660,7 @@ runShellCommand(CState *st, char *variable, char **argv, int argc)
 		{
 			arg = argv[i] + 1;	/* a string literal starting with colons */
 		}
-		else if ((arg = getVariable(st, argv[i] + 1)) == NULL)
+		else if ((arg = getVariable(variables, argv[i] + 1)) == NULL)
 		{
 			pgbench_error(LOG, "%s: undefined variable \"%s\"\n",
 						  argv[0], argv[i]);
@@ -2641,7 +2730,7 @@ runShellCommand(CState *st, char *variable, char **argv, int argc)
 					  argv[0], res);
 		return false;
 	}
-	if (!putVariableInt(st, "setshell", variable, retval, true))
+	if (!putVariableInt(variables, "setshell", variable, retval, true))
 		return false;
 
 #ifdef DEBUG
@@ -2695,7 +2784,7 @@ sendCommand(CState *st, Command *command)
 		char	   *sql;
 
 		sql = pg_strdup(command->argv[0]);
-		sql = assignVariables(st, sql);
+		sql = assignVariables(&st->variables, sql);
 
 		pgbench_error(DEBUG, "client %d sending %s\n", st->id, sql);
 		r = PQsendQuery(st->con, sql);
@@ -2706,7 +2795,7 @@ sendCommand(CState *st, Command *command)
 		const char *sql = command->argv[0];
 		const char *params[MAX_ARGS];
 
-		getQueryParams(st, command, params);
+		getQueryParams(&st->variables, command, params);
 
 		pgbench_error(DEBUG, "client %d sending %s\n", st->id, sql);
 		r = PQsendQueryParams(st->con, sql, command->argc - 1,
@@ -2747,7 +2836,7 @@ sendCommand(CState *st, Command *command)
 			st->prepared[st->use_file] = true;
 		}
 
-		getQueryParams(st, command, params);
+		getQueryParams(&st->variables, command, params);
 		preparedStatementName(name, st->use_file, st->command);
 
 		pgbench_error(DEBUG, "client %d sending %s\n", st->id, name);
@@ -2773,14 +2862,14 @@ sendCommand(CState *st, Command *command)
  * of delay, in microseconds.  Returns true on success, false on error.
  */
 static bool
-evaluateSleep(CState *st, int argc, char **argv, int *usecs)
+evaluateSleep(Variables *variables, int argc, char **argv, int *usecs)
 {
 	char	   *var;
 	int			usec;
 
 	if (*argv[1] == ':')
 	{
-		if ((var = getVariable(st, argv[1] + 1)) == NULL)
+		if ((var = getVariable(variables, argv[1] + 1)) == NULL)
 		{
 			pgbench_error(LOG, "%s: undefined variable \"%s\"\n",
 						  argv[0], argv[1]);
@@ -3060,7 +3149,7 @@ doCustom(TState *thread, CState *st, StatsData *agg)
 						 */
 						int			usec;
 
-						if (!evaluateSleep(st, argc, argv, &usec))
+						if (!evaluateSleep(&st->variables, argc, argv, &usec))
 						{
 							commandFailed(st, "sleep", "execution of meta-command failed");
 							st->state = CSTATE_ABORTED;
@@ -3101,8 +3190,8 @@ doCustom(TState *thread, CState *st, StatsData *agg)
 
 						if (command->meta == META_SET)
 						{
-							if (!putVariableValue(st, argv[0], argv[1], &result,
-												  true))
+							if (!putVariableValue(&st->variables, argv[0],
+												  argv[1], &result,  true))
 							{
 								commandFailed(st, "set", "assignment of meta-command failed");
 								st->state = CSTATE_ABORTED;
@@ -3155,7 +3244,9 @@ doCustom(TState *thread, CState *st, StatsData *agg)
 					}
 					else if (command->meta == META_SETSHELL)
 					{
-						bool		ret = runShellCommand(st, argv[1], argv + 2, argc - 2);
+						bool		ret = runShellCommand(&st->variables,
+														  argv[1], argv + 2,
+														  argc - 2);
 
 						if (timer_exceeded) /* timeout */
 						{
@@ -3175,7 +3266,8 @@ doCustom(TState *thread, CState *st, StatsData *agg)
 					}
 					else if (command->meta == META_SHELL)
 					{
-						bool		ret = runShellCommand(st, NULL, argv + 1, argc - 1);
+						bool		ret = runShellCommand(&st->variables, NULL,
+														  argv + 1, argc - 1);
 
 						if (timer_exceeded) /* timeout */
 						{
@@ -5204,7 +5296,7 @@ main(int argc, char **argv)
 					}
 
 					*p++ = '\0';
-					putVariable(&state[0], "option", optarg, p);
+					putVariable(&state[0].variables, "option", optarg, p);
 				}
 				break;
 			case 'F':
@@ -5503,18 +5595,19 @@ main(int argc, char **argv)
 			int			j;
 
 			state[i].id = i;
-			for (j = 0; j < state[0].nvariables; j++)
+			for (j = 0; j < state[0].variables.nvars; j++)
 			{
-				Variable   *var = &state[0].variables[j];
+				Variable   *var = &state[0].variables.vars[j];
 
 				if (var->value.type != PGBT_NO_VALUE)
 				{
-					putVariableValue(&state[i], "startup",
+					putVariableValue(&state[i].variables, "startup",
 									 var->name, &var->value, false);
 				}
 				else
 				{
-					putVariable(&state[i], "startup", var->name, var->svalue);
+					putVariable(&state[i].variables, "startup", var->name,
+								var->svalue);
 				}
 			}
 		}
@@ -5605,24 +5698,30 @@ main(int argc, char **argv)
 	 * :scale variables normally get -s or database scale, but don't override
 	 * an explicit -D switch
 	 */
-	if (lookupVariable(&state[0], "scale") == NULL)
+	if (lookupVariable(&state[0].variables, "scale") == NULL)
 	{
 		for (i = 0; i < nclients; i++)
-			putVariableInt(&state[i], "startup", "scale", scale, false);
+		{
+			putVariableInt(&state[i].variables, "startup", "scale", scale,
+						   false);
+		}
 	}
 
 	/*
 	 * Define a :client_id variable that is unique per connection. But don't
 	 * override an explicit -D switch.
 	 */
-	if (lookupVariable(&state[0], "client_id") == NULL)
+	if (lookupVariable(&state[0].variables, "client_id") == NULL)
 	{
 		for (i = 0; i < nclients; i++)
-			putVariableInt(&state[i], "startup", "client_id", i, false);
+		{
+			putVariableInt(&state[i].variables, "startup", "client_id", i,
+						   false);
+		}
 	}
 
 	/* set default seed for hash functions */
-	if (lookupVariable(&state[0], "default_seed") == NULL)
+	if (lookupVariable(&state[0].variables, "default_seed") == NULL)
 	{
 		uint64		seed = ((uint64) (random() & 0xFFFF) << 48) |
 		((uint64) (random() & 0xFFFF) << 32) |
@@ -5631,18 +5730,18 @@ main(int argc, char **argv)
 
 		for (i = 0; i < nclients; i++)
 		{
-			putVariableInt(&state[i], "startup", "default_seed", (int64) seed,
-						   false);
+			putVariableInt(&state[i].variables, "startup", "default_seed",
+						   (int64) seed, false);
 		}
 	}
 
 	/* set random seed unless overwritten */
-	if (lookupVariable(&state[0], "random_seed") == NULL)
+	if (lookupVariable(&state[0].variables, "random_seed") == NULL)
 	{
 		for (i = 0; i < nclients; i++)
 		{
-			putVariableInt(&state[i], "startup", "random_seed", random_seed,
-						   false);
+			putVariableInt(&state[i].variables, "startup", "random_seed",
+						   random_seed, false);
 		}
 	}
 
-- 
2.7.4

