pgbench: Improve error-handling in \sleep command.

MasaoFujii · MasaoFujii · commit 8c6eda2d1c92 · 2021-03-22T12:04:48.000+09:00
This commit improves pgbench \sleep command so that it handles the following three cases more properly. (1) When only one argument was specified in \sleep command and it's not a number, previously pgbench reported a confusing error message like "unrecognized time unit, must be us, ms or s". This commit fixes this so that more proper error message like "invalid sleep time, must be an integer" is reported. (2) When two arguments were specified in \sleep command and the first argument was not a number, previously pgbench treated that argument as the sleep time 0. No error was reported in this case. This commit fixes this so that an error is thrown in this case. (3) When a variable was specified as the first argument in \sleep command and the variable stored non-digit value, previously pgbench treated that argument as the sleep time 0. No error was reported in this case. This commit fixes this so that an error is thrown in this case. Author: Kota Miyake Reviewed-by: Hayato Kuroda, Alvaro Herrera, Fujii Masao Discussion: https://postgr.es/m/23b254daf20cec4332a2d9168505dbc9@oss.nttdata.com
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
@@ -2953,7 +2953,16 @@ evaluateSleep(CState *st, int argc, char **argv, int *usecs)
 			pg_log_error("%s: undefined variable \"%s\"", argv[0], argv[1] + 1);
 			return false;
 		}
+
 		usec = atoi(var);
+
+		/* Raise an error if the value of a variable is not a number */
+		if (usec == 0 && !isdigit((unsigned char) *var))
+		{
+			pg_log_error("%s: invalid sleep time \"%s\" for variable \"%s\"",
+						 argv[0], var, argv[1] + 1);
+			return false;
+		}
 	}
 	else
 		usec = atoi(argv[1]);
@@ -4788,17 +4797,41 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 		 * will be parsed with atoi, which ignores trailing non-digit
 		 * characters.
 		 */
-		if (my_command->argc == 2 && my_command->argv[1][0] != ':')
+		if (my_command->argv[1][0] != ':')
 		{
 			char	   *c = my_command->argv[1];
+			bool		have_digit = false;
+
+			/* Skip sign */
+			if (*c == '+' || *c == '-')
+				c++;
+
+			/* Require at least one digit */
+			if (*c && isdigit((unsigned char) *c))
+				have_digit = true;
 
-			while (isdigit((unsigned char) *c))
+			/* Eat all digits */
+			while (*c && isdigit((unsigned char) *c))
 				c++;
+
 			if (*c)
 			{
-				my_command->argv[2] = c;
-				offsets[2] = offsets[1] + (c - my_command->argv[1]);
-				my_command->argc = 3;
+				if (my_command->argc == 2 && have_digit)
+				{
+					my_command->argv[2] = c;
+					offsets[2] = offsets[1] + (c - my_command->argv[1]);
+					my_command->argc = 3;
+				}
+				else
+				{
+					/*
+					 * Raise an error if argument starts with non-digit
+					 * character (after sign).
+					 */
+					syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
+								 "invalid sleep time, must be an integer",
+								 my_command->argv[1], offsets[1] - start_offset);
+				}
 			}
 		}
 
